#!/usr/bin/php
<?php

/* 
 * Create an AMI in the specified region from a local FS image.
 *
 * To do this we need ssh access to an account in disk group in an
 * EC2 instance. We call this ec2 build host. 
 * 
 * We call local machine the local build host.
 * 
 * Pre-requisites:
 *  - Must have an AWS account and the pair of access keys
 * 	AWS console/My Account/Security Credentials/Access Keys
 *  - Must have an EC2 instance running with ssh root access (or sudoer)
 *  - install AWS PHP on local build host.
 *       http://aws.amazon.com/sdkforphp/
 *	 (tested with version 2)
 *
 * Complete region list as of 2/23/2014:
	ap-northeast-1
	ap-southeast-1
	ap-southeast-2
	eu-west-1
	sa-east-1
	us-east-1
	us-west-1
	us-west-2
 */
require 'vendor/autoload.php';

use Aws\Ec2\Ec2Client;

/* ============== CONFIGS =============== */

/* configs for AMI image */
define('AMI_IMAGE_FILE', '/opt/ec2/images/vyatta-64bit.img');
//define('AMI_IMAGE_FILE', '/opt/ec2/images/simple.img');
define('AMI_EBS_SIZE', 2);	// in GB units
define('AMI_NAME', 'Vyatta v6.6r1-js140325 (PV EBS x86_64)');
define('AMI_DESCRIPTION', AMI_NAME . ' (built from community dasiy branch)');
//define('AMI_REGIONS', 'eu-west-1 sa-east-1 ap-northeast-1 ap-southeast-1 ap-southeast-2 us-west-1 us-west-2 us-east-1'); // list of regions AMI shold be created
define('AMI_REGIONS', 'us-east-1'); // list of regions AMI shold be created
define('AMI_ARCH', 'x86_64'); // i386 or x86_64

/* configs for ec2 build host */
define('EC2_BUILD_ZONE', 'us-west-1c');
define('EC2_BUILD_REGION', 'us-west-1');
define('EC2_BUILD_INSTANCE', 'i-9a0c7bdc');
define('EC2_BUILD_NAME', 'dev.jinggling.com');	// IP addr is OK
define('EC2_BUILD_PORT', '2222');
define('EC2_BUILD_USER', 'jsun');

/* configs of credential */
define('EC2_ACCESS_KEY','XXXXXXXXXXXX');
define('EC2_ACCESS_SECRET', 'XXXXXXXXXXXXXXX');

/* ============== PVGRUB AKI =============== */

/* we use v 1.04, around 2/23/2014.
 * One can update this list by (change region as necessary)

ec2-describe-images --owner amazon --region us-west-1 | grep "amazon\/pv-grub-hd0" | awk '{ print $1, $2, $3, $5, $7 }'

 Note: we use hd0 version which is partitionless. hd00 is for image with 
 partition table.

 NOTE: this can be computed dynamically by script. We just go easy for now.
  TODO.
 */

$aki=array(
	'ap-northeast-1' => array(
		'i386' => 'aki-136bf512',
		'x86_64' => 'aki-176bf516',
	),
	'ap-southeast-1' => array(
		'i386' => 'aki-ae3973fc',
		'x86_64' => 'aki-503e7402',
	),
	'ap-southeast-2' => array(
		'i386' => 'aki-cd62fff7',
		'x86_64' => 'aki-c362fff9',
	),
	'eu-west-1' => array(
		'i386' => 'aki-68a3451f',
		'x86_64' => 'aki-52a34525',
	),
	'sa-east-1' => array(
		'i386' => 'aki-5b53f446',
		'x86_64' => 'aki-5553f448',
	),
	'us-east-1' => array(
		'i386' => 'aki-8f9dcae6',
		'x86_64' => 'aki-919dcaf8',
	),
	'us-west-1' => array(
		'i386' => 'aki-8e0531cb',
		'x86_64' => 'aki-880531cd',
	),
	'us-west-2' => array(
		'i386' => 'aki-f08f11c0',
		'x86_64' => 'aki-fc8f11cc',
	),
);

/* ============== check CONFIGS =============== */
function myerror($msg) {
	echo $msg . "\n";
	exit(-1);
}

function check_configs() {
	global $aki;
	if ( PHP_INT_SIZE != 8) 
		myerror("This script only supports 64bit machine!". PHP_INT_SIZE);

	// check file
	if (!file_exists(AMI_IMAGE_FILE)) 
		myerror("Cannot find image file " . AMI_IMAGE_FILE);

	// check size
	if (filesize(AMI_IMAGE_FILE) > AMI_EBS_SIZE*1024*1024*1024)
		myerror("Image file size is bigger than EBS volume size: " . filesize(AMI_IMAGE_FILE) . " > " . AMI_EBS_SIZE*1024*1024*1024);

	// check build region
	if (!array_key_exists(EC2_BUILD_REGION, $aki))
		myerror("Build host region is not found, " . EC2_BUILD_REGION);

	// check region and arch
	$regions = preg_split('/\s+/', AMI_REGIONS);
	foreach ($regions as $r) 
		if (!array_key_exists($r, $aki)) 
			myerror("One of the target region is not found in aki list: $r");

	if (!array_key_exists(AMI_ARCH, $aki[EC2_BUILD_REGION]))
		myerror("AMI_ARCH is not found in aki list: " . AMI_ARCH);
}

/* ============== check build host =============== */
function ssh_cmd($cmd) {
	$cmd="ssh -p " . EC2_BUILD_PORT . " -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -o LogLevel=quiet " . EC2_BUILD_USER . "@" . EC2_BUILD_NAME . " \"$cmd\"";
	$result=shell_exec($cmd);
	if ($result == NULL) 
		myerror("Failed to execute SSH command: $cmd\n");
	return $result;
}

function check_build_host() {
	// check user group
	$result=ssh_cmd("groups");
	if (strpos($result,  "disk") == NULL) 
		myerror("User " . EC2_BUILD_USER . " does not belong to disk group. Quitting.");

	// check existing partitions
	$result=ssh_cmd("cat /proc/partitions");
	if (strpos($result,  "xvdf") != NULL) 
		myerror("partition /dev/xvdf aleady exist on the ec2 buiil host. Please detach it first before proceeding.");
}

/* ============== snapshot/volume state =============== */
function image_state($p) {
	$c=$p[0];
	$ami_id=$p[1];
	$result=$c->describeImages(array('ImageIds'=>array($ami_id)));
	return $result['Images'][0]['State'];
}

function snapshot_state($ss_id) {
	global $c;
	try {
		$result=$c->describeSnapshots(array('SnapshotIds'=>array($ss_id)));
	} catch (Exception $e) {
		if($e->getExceptionCode() == "InvalidSnapshot.NotFound")
			return "deleted";
		else
			return "unknown(" . $e->getExceptionCode() . ")";
	}
	return $result['Snapshots'][0]['State'];
}
function volume_state($vol_id) {
	global $c;
	try {
		$result=$c->describeVolumes(array('VolumeIds'=>array($vol_id)));
	} catch (Exception $e) {
		if($e->getExceptionCode() == "InvalidVolume.NotFound")
			return "deleted";
		else
			return "unknown(" . $e->getExceptionCode() . ")";
	}
	return $result['Volumes'][0]['State'];
}

function wait($func, $arg, $res) {
	// echo "Waiting for $func to be $res ...\n";
	for($i=0;$i<10;$i++) {
		if ($func($arg) == $res) return;
		sleep(5);
	}
	echo "Failed to reach expected state. Bail out.\n";
}
/* ============== delete ami/snapshots =============== */
function find_image($r) {
	global $ca;
	$result=$ca[$r]->describeImages(array('Owners'=>array('self')));
	foreach ($result['Images'] as $i) {
		if ($i['Name'] == AMI_NAME) 
			return $i['ImageId'];
	}
	return NULL;
}

function delete_images() {
	global $regions;
	global $ca;
	foreach ($regions as $r) {
		$result= find_image($r);
		if ($result) {
			echo "deleting AMI $result in $r ...\n";
			$ca[$r]->deregisterImage(array(
				'ImageId' => $result,
			));
		}
	}
}

function find_snapshots($r) {
	global $ca;
	$a=array();
	$result=$ca[$r]->describeSnapshots(array('OwnerIds'=>array('self')));
	foreach ($result['Snapshots'] as $s) {
		if (strpos($s['Description'], AMI_NAME) != NULL)
			array_push($a, $s['SnapshotId']);
	}
	return $a;
}

function delete_snapshots() {
	global $ca;
	global $regions;
	foreach ($regions as $r) {
		$result= find_snapshots($r);
		foreach ($result as $s) {
			echo "deleting snapshot $s in $r ...\n";
			$ca[$r]->deleteSnapshot(array(
				'SnapshotId' => $s,
			));
		}
	}		
}

/* ============== main =============== */
function usage($name) {
	echo "\n";
	echo "Usage:\n";
	echo "\n";
	echo "$name [--run|--delete-image|--delete-snapshot]\n";
	echo "\n";
	exit(1);
}

if ($argc != 2) usage($argv[0]);

check_configs();

check_build_host();

// create a client for build region
$c = Ec2Client::factory(array(
	'key' => EC2_ACCESS_KEY,
	'secret' => EC2_ACCESS_SECRET,
	'region' => EC2_BUILD_REGION,
));

// create a client for each target region
$regions = preg_split('/\s+/', AMI_REGIONS);
foreach ($regions as $r) {
	if ($r == EC2_BUILD_REGION) 
		$ca[$r] = $c;
	else
		$ca[$r]= Ec2Client::factory(array(
			'key' => EC2_ACCESS_KEY,
			'secret' => EC2_ACCESS_SECRET,
			'region' => $r,
		));
}

if ($argv[1] == '--delete-image') {
	delete_images();
	exit();
} elseif ($argv[1] == '--delete-snapshot') {
	delete_images();
	delete_snapshots();
	exit();
} elseif ($argv[1] != '--run') {
	usage($argv[0]);
}

// check to see if target ami already exists
foreach ($regions as $r) {
	if (find_image($r) != NULL) 
		myerror("Target AMI already exists in region $r. Quitting.");
}

// create volume
echo "create ebs volume ... ";
$result=$c->createVolume(array(
	'Size' => AMI_EBS_SIZE,
	'AvailabilityZone' => EC2_BUILD_ZONE,
));
$vol_id=$result['VolumeId'];
//$result=$c->waitUntilVolumeAvailable(array('VolumeIds'=>array($vol_id),'waiter.interval'=>10, 'waiter.max_attempts'=>3));
//$w=$c->getWaiter('VolumeInUse')
//	->setConfig(array('VolumeIds'=>array($vol_id)))
//	->setInterval(10)
//	->setMaxAttempts(3);
//$w->wait();
wait('volume_state', $vol_id, "available");
echo  $vol_id . "\n";

// attach volume to instance
echo "attach volume to the ec2 instance ... ";
$result=$c->attachVolume(array(
	'VolumeId' => $vol_id,
	'InstanceId' => EC2_BUILD_INSTANCE,
	'Device' => '/dev/sdf',
));
//$c->waitUntilVolumeInUse(array('VolumeIds'=>array($vol_id)));
wait('volume_state', $vol_id, "in-use");
echo "done \n";

// scp image to the drive
echo "scp " . AMI_IMAGE_FILE . " over to the EBS volume (may take long time) ... ";
sleep(10);		// wait for linux to detect the volume
$cmd="gzip -c " . AMI_IMAGE_FILE . " | ssh -p " . EC2_BUILD_PORT . " -oStrictHostKeyChecking=no -oUserKnownHostsFile=/dev/null -o LogLevel=quiet " . EC2_BUILD_USER . "@" . EC2_BUILD_NAME . " \"gunzip > /dev/xvdf\"";
system($cmd, $ret);
if ($ret != 0) {
	echo "ssh failed: $ret\n";
	exit(1);
}
echo "Done!\n";

// just to be paranoid; flush before detach; supposedly this is not necessary
// last "ls" is to prevent the function from barfing errors.
ssh_cmd("sync;sync;sync;sleep 5;sync;sync;sync;ls /");

// detach volume from instance
echo "detach volume from the ec2 instance ... ";
$result=$c->detachVolume(array(
	'VolumeId' => $vol_id,
	'InstanceId' => EC2_BUILD_INSTANCE,
	'Device' => '/dev/sdf',
));
//$c->waitUntilVolumeAvailable(array('VolumeIds'=>array($vol_id)));
wait('volume_state', $vol_id, "available");
echo  "Done!\n";

// create snapshot
echo "create snapshot (may take long time) ...";
$snapshot_desc='Root image for AMI "'  . AMI_NAME . '", ' . date("Y/m/d H:i:s");
$result=$c->createSnapshot(array(
	'VolumeId' => $vol_id,
	'Description' => $snapshot_desc,
));
$snapshot_id=$result['SnapshotId'];
$c->waitUntilSnapshotCompleted(array('SnapshotIds'=>array($snapshot_id)));
echo $snapshot_id . " in " . EC2_BUILD_REGION . "\n";

// delete volume
echo "deleting volume $vol_id ... ";
$result=$c->deleteVolume(array('VolumeId' => $vol_id));
//$c->waitUntilVolumeDeleted(array('VolumeIds'=>array($vol_id)));
// wait('volume_state', $vol_id, "deleted");  // jsun: no need to wait
echo "Done!\n";

// copy snapshot to target regions
echo "Initiate copying snapshot to target regions ... ";
foreach ($regions as $r) {
	if ($r == EC2_BUILD_REGION) {
		$target_ss_id[$r] = $snapshot_id;
		continue;
	}

	$result=$ca[$r]->copySnapshot(array(
		'SourceRegion' => EC2_BUILD_REGION,
		'SourceSnapshotId' => $snapshot_id,
		'Description' => $snapshot_desc . " (copy from $snapshot_id in " . EC2_BUILD_REGION .")",
	));
	$target_ss_id[$r]=$result['SnapshotId'];
}
echo "Done!\n";

// wait for snapshot copying to be done
echo "Creating AMI in each target region ... \n";
foreach ($regions as $r) {
	echo "\t$r (" . $target_ss_id[$r] . ") ... ";
	$ca[$r]->waitUntilSnapshotCompleted(array('SnapshotIds'=>array($target_ss_id[$r])));

	$result=$ca[$r]->registerImage(array(
		'Name' => AMI_NAME,
		'Description' => AMI_DESCRIPTION,
		'Architecture' => AMI_ARCH,
		'KernelId' => $aki[$r][AMI_ARCH],
		'RootDeviceName' => '/dev/sda1',
		'BlockDeviceMappings' => array(
			array(
				'DeviceName' => '/dev/sda1',
				'Ebs' => array(
					'SnapshotId' => $target_ss_id[$r],
					'VolumeSize' => AMI_EBS_SIZE,
					'DeleteOnTermination' => true,
				),
			),
		),
	));
	$target_ami[$r]=$result['ImageId'];
	echo " " . $target_ami[$r] . "\n";
}
echo "Done!\n";

// wait for ami to be done
echo "Wait for ami to be done ...";
foreach ($regions as $r) {
	wait('image_state', array($ca[$r], $target_ami[$r]), 'available');
}
echo "All done!\n";

?>