#!/usr/bin/python
#
#    ubuntu-ec2-run: ec2-run-instances that support human readable aliases for AMI's
#
#    Copyright (C) 2011 Dustin Kirkland <kirkland@ubuntu.com>
#
#    Authors: Dustin Kirkland <kirkland@ubuntu.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, version 3 of the License.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

USAGE = """
Usage: ubuntu-ec2-run [ options ] arguments

  Run an ec2 instance of Ubuntu.

  options:
    --dry-run: only report what would be done
    
   All non-understood options are passed through to $EC2_PRE-run-instances

   ubuntu-ec2-run passes the following arguments to cloud-image-query 
   in order to select an AMI to run.  Defaults are marked with a '*':

     releases: lucid* maverick natty oneiric precise
     stream: release* daily 
     arch: amd64*, x86_64, i386
     store: ebs*, instance-store, instance
     pvtype: pv*, hvm, paravirtual

   Note, that --instance-type/-t will modify arch appropriately

  Example:
   * ubuntu-ec2-run oneiric daily --dry-run
     # us-east-1/ebs/ubuntu-oneiric-daily-amd64-server-20110902
     ec2-run-instances --instance-type=t1.micro ami-0ba16262
   * EC2_PRE=euca- ubuntu-ec2-run lucid released --dry-run
     # us-east-1/ebs/ubuntu-oneiric-daily-amd64-server-20110902
     euca-run-instances released --instance-type=t1.micro ami-0ba16262
   * ubuntu-ec2-run oneiric hvm --dry-run
     # us-east-1/hvm/ubuntu-oneiric-11.10-beta1-amd64-server-20110831
     ec2-run-instances ./bin/ubuntu-ec2-run --instance-type=cc1.4xlarge \\
         --block-device-mapping /dev/sdb=ephemeral0 \\
         --block-device-mapping /dev/sdc=ephemeral1 ami-b79754de
   * ubuntu-ec2-run --region us-west-1 --instance-type \\
         m1.small oneiric instance --dry-run
     # us-west-1/instance-store/ubuntu-oneiric-11.10-beta1-i386-server-20110831
     ec2-run-instances --region us-west-1 --instance-type m1.small ami-39bfe27c
"""

import os, string, sys, urllib2
import subprocess

# This could/should use `distro-info --supported`
aliases = [
  "amd64", "x86_64", "i386",
  "server", "desktop",
  "release", "daily",
  "ebs", "instance-store", "instance",
  "hvm", "paravirtual", "pv",
]

def get_argopt(args, optnames):
   ret = None
   i = 0
   while i < len(args):
      cur = args[i];
      for opt in optnames:
         if opt.startswith("--"):
            if cur == opt:
               ret = args[i+1]
               i = i + 1
               break
            elif cur.startswith("%s=" % opt):
               ret = args[i].split("=")[1]
               break
         else:
            if args[i] == opt:
               ret = args[i+1]
               i = i + 1
               break
      i = i + 1
   return ret

def get_block_device_mappings(itype):
   # cleaned from http://aws.amazon.com/ec2/instance-types/
   # t1.micro     NONE    # m2.2xlarge    850   # c1.xlarge    1690
   # m1.small      160    # m1.large      850   # m1.xlarge    1690
   # c1.medium     350    # cc1.4xlarge  1690   # cc1.4xlarge  1690
   # m2.xlarge     420    # m2.4xlarge   1690   # cg1.4xlarge  1690
   #                                            # cc2.8xlarge  3370
   bdmaps = [ ]
   if itype in ("t1.micro", "m1.small", "c1.medium"):
      pass # the first one is always attached. ephemeral0=sda2
   elif itype in ("m2.xlarge"):
      pass # one 420 for m2.xlarge
   elif ( itype in ("m1.large", "m2.2xlarge") or 
          itype.startswith("cg1.") or itype.startswith("cc1.")):
         bdmaps=[ "/dev/sdb=ephemeral0", "/dev/sdc=ephemeral1", ]
   elif ( itype in ("m1.xlarge", "m2.4xlarge", "c1.xlarge") or
          itype.startswith("cc2.8xlarge")):
         bdmaps=[ "sdb=ephemeral0", "sdc=ephemeral1",
                  "sdd=ephemeral2", "sde=ephemeral3", ]
   args = [ ]
   for m in bdmaps:
      args.extend(("--block-device-mapping", m,))
   return(args)

if "--help" in sys.argv or "-h" in sys.argv:
   sys.stdout.write(USAGE)
   sys.exit(0)

if len(sys.argv) == 1:
   sys.stderr.write(USAGE)
   sys.exit(1)

pre = "ec2-"
for name in ("EC2_PRE", "EC2PRE"):
   if name in os.environ:
      pre=os.environ[name]

# if the prefix is something like "myec2 "
# then assume that 'myec2' is a command itself
if pre.strip() == pre:
   ri_cmd = [ "%srun-instances" % pre ]
else:
   ri_cmd = [ pre.strip(), "run-instances" ]

query_cmd = [ "ubuntu-cloudimg-query", "--format=%{ami}\n%{itype}\n%{summary}\n%{store}\n" ]


# Get the list of releases.  If they have 'ubuntu-distro-info', then use that
# otherwise, fall back to our builtin list of releases
try:
   out = subprocess.check_output(["ubuntu-distro-info", "--all"])
   all_rels = out.strip().split("\n")
   releases = [ ]
   seen_lucid = False
   for r in all_rels:
      if seen_lucid or r == "lucid":
         seen_lucid = True
         releases.append(r)
except OSError as e:
   releases = [ "lucid", "maverick", "natty", "oneiric", "precise", ]


# each arg_group is a list of arguments and a default if not found
# ec2-run-instances default instance-type is m1.small
arg_groups = (
   ("--region",),
   ("--instance-type", "-t"),
   ("--block-device-mapping", "-b"),
)

flags = { }
for opts in arg_groups:
   arg_value = get_argopt(sys.argv, opts)
   if arg_value != None:
      query_cmd.append(arg_value)
   flags[opts[0]]=arg_value

dry_run = False

for arg in sys.argv[1:]:
   if arg in aliases or arg in releases:
      query_cmd.append(arg)
   elif arg == "--dry-run":
      dry_run = True
   else:
      ri_cmd.append(arg)

cmd=""
for i in query_cmd:
   cmd += " '%s'" % i.replace("\n","\\n")
cmd=cmd[1:]

try:
   (ami, itype, summary, store, endl) = subprocess.check_output(query_cmd).split("\n")
   if endl.strip():
      sys.stderr.write("Unexpected output of command:\n  %s" % cmd)
except subprocess.CalledProcessError as e:
   sys.stderr.write("Failed. The following command returned failure:\n")
   sys.stderr.write("  %s\n" % cmd)
   sys.exit(1)
except OSError as e:
   sys.stderr.write("You do not have '%s' in your path\n" % query_cmd[0])
   sys.exit(1)

if flags.get("--instance-type",None) == None:
   ri_cmd.append("--instance-type=%s" % itype)

if store == "ebs" and flags.get("--block-device-mapping",None) == None:
   ri_cmd.extend(get_block_device_mappings(itype))

ri_cmd.append(ami)

sys.stderr.write("# %s\n" % summary)
if dry_run:
   print ' '.join(ri_cmd)
else:
   os.execvp(ri_cmd[0], ri_cmd)
################################################################################
