#!/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/>.

KNOWN_RELEASES = ["lucid", "maverick", "natty", "oneiric", "precise",
                  "quantal", "raring"]

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: %(rels)s
     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
""" % {'rels': ' '.join(KNOWN_RELEASES)}

import os
import string
import subprocess
import sys
import urllib2

# 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        0   # m1.large      850   # cg1.4xlarge  1690
    # m1.small      160   # m2.2xlarge    850   # m1.xlarge    1690
    # c1.medium     350   # c1.xlarge    1690   # m2.4xlarge   1690
    # m1.medium     410   # cc1.4xlarge  1690   # hi1.4xlarge  2048
    # m2.xlarge     420   # cc1.4xlarge  1690   # cc2.8xlarge  3370
    # m3.xlarge       0
    # m3.2xlarge      0
    bdmaps = []
    if (itype in ("t1.micro", "m1.small", "c1.medium") or
        itype.startswith("m3.")):
        pass  # the first one is always attached. ephemeral0=sda2
    elif itype in ("m2.xlarge", "m1.medium"):
        bdmaps = ["/dev/sdb=ephemeral0"]
    elif (itype in ("m1.large", "m2.2xlarge", "hi1.4xlarge") 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 = KNOWN_RELEASES


# each arg_group is a list of arguments and a boolean that indicates
# if the value of that argument should be passed to query_cmd
# ec2-run-instances default instance-type is m1.small
arg_groups = (
    (("--region",), True),
    (("--instance-type", "-t"), True),
    (("--block-device-mapping", "-b"), False),
)

flags = {}
for opts, passthrough in arg_groups:
    arg_value = get_argopt(sys.argv, opts)
    if arg_value is not None and passthrough:
        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) is None:
    ri_cmd.append("--instance-type=%s" % itype)

if store == "ebs" and flags.get("--block-device-mapping", None) is 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)
###############################################################################

# vi: ts=4 expandtab
