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

import os
import string
import subprocess
import sys

KNOWN_RELEASES = ["lucid", "maverick", "natty", "oneiric", "precise",
                  "quantal", "raring", "trusty", "utopic", "vivid",
                  "wily", "xenial", "yakkety"]

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)}

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

SSD = "ssd"
SPIN = "spin"

# cleaned from http://aws.amazon.com/ec2/instance-types/
# (vcpu, compute-units, mem, disknum, disksize, diskback)
SIZE_DATA = {
    'm3.medium': (1, 3, 3.75, 1, 4, SSD),
    'm3.large': (2, 6.5, 7.5, 1, 32, SSD),
    'm3.xlarge': (4, 13, 15, 2, 40, SSD),
    'm3.2xlarge': (8, 26, 30, 2, 80, SSD),
    'm1.small': (1, 1, 1.7, 1, 160, SPIN),
    'm1.medium': (1, 2, 3.75, 1, 410, SPIN),
    'm1.large': (2, 4, 7.5, 2, 420, SPIN),
    'm1.xlarge': (4, 8, 15, 4, 420, SPIN),
    'c3.large': (2, 7, 3.75, 2, 16, SSD),
    'c3.xlarge': (4, 14, 7.5, 2, 40, SSD),
    'c3.2xlarge': (8, 28, 15, 2, 80, SSD),
    'c3.4xlarge': (16, 55, 30, 2, 160, SSD),
    'c3.8xlarge': (32, 108, 60, 2, 320, SSD),
    'c1.medium': (2, 5, 1.7, 1, 350, SPIN),
    'c1.xlarge': (8, 20, 7, 4, 420, SPIN),
    'cc2.8xlarge': (32, 88, 60.5, 4, 840, SPIN),
    'g2.2xlarge': (8, 26, 15, 1, 60, SSD),
    'cg1.4xlarge': (16, 33.5, 22.5, 2, 840, SPIN),
    'm2.xlarge': (2, 6.5, 17.1, 1, 420, SPIN),
    'm2.2xlarge': (4, 13, 34.2, 1, 850, SPIN),
    'm2.4xlarge': (8, 26, 68.4, 2, 840, SPIN),
    'cr1.8xlarge': (32, 88, 244, 2, 120, SSD),
    'i2.xlarge': (4, 14, 30.5, 1, 800, SSD),
    'i2.2xlarge': (8, 27, 61, 2, 800, SSD),
    'i2.4xlarge': (16, 53, 122, 4, 800, SSD),
    'i2.8xlarge': (32, 104, 244, 8, 800, SSD),
    'hs1.8xlarge': (16, 35, 117, 2, 2048, SPIN),
    'hi1.4xlarge': (16, 35, 60.5, 2, 1024, SSD),
    't1.micro': (1, .1, 0.615, 0, 0, None),
}


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):
    bdmaps = []
    allmaps = ["/dev/sdb=ephemeral0", "/dev/sdc=ephemeral1",
               "/dev/sdd=ephemeral2", "/dev/sde=ephemeral3"]
    if itype in SIZE_DATA:
        (vcpu, ec2, mem, disknum, disksize, diskback) = SIZE_DATA[itype]
        bdmaps = allmaps[0:disknum]

    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.decode().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).decode().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
