lite

#!/usr/bin/env python

from __future__ import print_function
import sys
import os
import shutil
import subprocess
import platform
import string
import signal
import re


DEFAULT_OPS_FILES = [
    "local-dns.yml",
    "virtualbox/cpi.yml",
    "virtualbox/outbound-network.yml",
    "bosh-lite.yml",
    "bosh-lite-runc.yml",
    "jumpbox-user.yml",
]
USAGE = """\
usage: lite create env-name
       lite delete env-name
       lite set-routes env-name
       lite unset-routes env-name
       lite tunnel env-name deployment instance_group tunnel
       lite ssh-command env-name deployment instance_group cmd
       lite proc-status env-name deployment instance_group
       lite go-routines env-name deployment instance_group

set environment variable 'BOSH_CMD' to set bosh-cli path. Defaults to
'/usr/local/bin/bosh'.
"""
BAD_ENV_NAME_CHARS = """ !@#$%^&*()_+-={}|[]\:";'<>?,/~`"""


def bosh_cmd():
    return os.environ.get("BOSH_CMD", "/usr/local/bin/bosh")


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def create_env(env_name, args):
    if len(args) != 0:
        eprint(USAGE)
        sys.exit(1)

    eprint("creating env:", env_name)
    eprint("on network: 192.168.%d.0/24" % class_c(env_name))
    create_deployment_dir(env_name)
    write_ops(env_name, DEFAULT_OPS_FILES)
    code = subprocess.call([
        bosh_cmd(),
        "create-env",
        "~/workspace/bosh-deployment/bosh.yml",
        "--state", dir_name(env_name) + "/state.json",
        "--vars-store", dir_name(env_name) + "/creds.yml",
        "-v", "director_name='Bosh Lite Director %s'" % env_name,
        "-v", "internal_cidr=192.168.%d.0/24" % class_c(env_name),
        "-v", "internal_gw=192.168.%d.1" % class_c(env_name),
        "-v", "internal_ip=192.168.%d.6" % class_c(env_name),
        "-v", "outbound_network_name=NatNetwork",
    ] + ops_args(DEFAULT_OPS_FILES), stdout=sys.stdout, stderr=sys.stderr)
    if code != 0:
        eprint("encountered error with creating env")
        delete_deployment_dir(env_name)
        sys.exit(code)

    ca_cert = subprocess.check_output([
        bosh_cmd(),
        "int",
        dir_name(env_name) + "/creds.yml",
        "--path", "/director_ssl/ca",
    ])
    password = subprocess.check_output([
        bosh_cmd(),
        "int",
        dir_name(env_name) + "/creds.yml",
        "--path", "/admin_password",
    ])
    subprocess.call([
        bosh_cmd(),
        "-e", "192.168.%d.6" % class_c(env_name),
        "--ca-cert", ca_cert,
        "--client", "admin",
        "--client-secret", password,
        "alias-env",
        env_name,
    ], stdout=sys.stdout, stderr=sys.stderr)
    insert_secrets(env_name, "admin", password)
    subprocess.call([
        bosh_cmd(),
        "-n",
        "-e", env_name,
        "upload-stemcell",
        "https://bosh.io/d/stemcells/bosh-warden-boshlite-ubuntu-trusty-go_agent",
    ], stdout=sys.stdout, stderr=sys.stderr)
    subprocess.call([
        bosh_cmd(),
        "-n",
        "-e", env_name,
        "update-cloud-config",
        os.path.expanduser("~/workspace/bosh-deployment/warden/cloud-config.yml"),
    ], stdout=sys.stdout, stderr=sys.stderr)
    set_routes(env_name, [])


def create_deployment_dir(env_name):
    if env_exists(env_name):
        eprint("env directory already exists:", dir_name(env_name))
        sys.exit(1)
    os.makedirs(dir_name(env_name))


def env_exists(env_name):
    return os.path.isdir(dir_name(env_name))


def delete_deployment_dir(env_name):
    shutil.rmtree(dir_name(env_name))


def dir_name(env_name):
    return os.path.expanduser("~/deployments/"+env_name)


def write_ops(env_name, ops):
    with open(dir_name(env_name) + "/ops", "w") as f:
        for op in ops:
            f.write(op + "\n")


def ops_args(ops):
    args = []
    for op in ops:
        args.append("-o")
        args.append(os.path.expanduser("~/workspace/bosh-deployment/" + op))
    return args


def class_c(env_name):
    return hash(env_name) % 204 + 50


def insert_secrets(env_name, username, password):
    with open(os.path.expanduser("~/.bosh/config")) as f:
        data = f.read()
    search_line = "  alias: %s\n" % env_name
    pos = string.find(data, search_line)
    if pos == -1:
        eprint("unable to add secrets to alias-env.")
        return
    secrets = """
  username: %s
  password: %s
""" % (username, password)
    with open(os.path.expanduser("~/.bosh/config"), "w") as f:
        f.write(data[:pos] + secrets + data[pos:])


def delete_env(env_name, args):
    if len(args) != 0:
        eprint(USAGE)
        sys.exit(1)

    eprint("deleting env:", env_name)
    ensure_env_exists(env_name)
    ops = read_ops(env_name)
    code = subprocess.call([
        bosh_cmd(),
        "delete-env",
        "~/workspace/bosh-deployment/bosh.yml",
        "--state", dir_name(env_name) + "/state.json",
        "--vars-store", dir_name(env_name) + "/creds.yml",
        "-v", "director_name='Bosh Lite Director %s'" % env_name,
        "-v", "internal_cidr=192.168.%d.0/24" % class_c(env_name),
        "-v", "internal_gw=192.168.%d.1" % class_c(env_name),
        "-v", "internal_ip=192.168.%d.6" % class_c(env_name),
        "-v", "outbound_network_name=NatNetwork",
    ] + ops_args(ops), stdout=sys.stdout, stderr=sys.stderr)
    if code != 0:
        eprint("encountered error with deleting env")
        sys.exit(code)
    eprint("deleting routes")
    unset_routes(env_name)
    eprint("deleting deployment dir")
    delete_deployment_dir(env_name)
    log_out(env_name)


def ensure_env_exists(env_name):
    if not env_exists(env_name):
        eprint("env does not exist")
        sys.exit(1)


def read_ops(env_name):
    with open(dir_name(env_name) + "/ops") as f:
        return [l.strip() for l in f.readlines()]

def log_out(env_name):
    eprint("logging out")
    code = subprocess.call([
        "bosh",
        "--environment", env_name,
        "log-out"],
        stdout=sys.stdout, stderr=sys.stderr)
    if code != 0:
        eprint("encountered error with logging out")
        sys.exit(code)


def unset_routes(env_name):
    if platform.system() == "Linux":
        subprocess.call([
            "sudo",
            "route",
            "del",
            "-net",
            "10.244.0.0/16",
            "gw",
            "192.168.%d.6" % class_c(env_name),
        ])
    elif platform.system() == "Windows":
        subprocess.call([
            "route",
            "del",
            "10.244.0.0/16",
            "192.168.%d.6" % class_c(env_name),
        ])
    elif platform.system() == "Darwin":
        subprocess.call([
            "sudo",
            "route",
            "-n",
            "delete",
            "10.244.0.0/16",
            "192.168.%d.6" % class_c(env_name),
        ])
    else:
        eprint("unable to detect platform")
        sys.exit(1)


def set_routes(env_name, args):
    if len(args) != 0:
        eprint(USAGE)
        sys.exit(1)

    if platform.system() == "Linux":
        subprocess.call([
            "sudo",
            "route",
            "add",
            "-net",
            "10.244.0.0/16",
            "gw",
            "192.168.%d.6" % class_c(env_name),
        ])
    elif platform.system() == "Windows":
        subprocess.call([
            "route",
            "add",
            "10.244.0.0/16",
            "192.168.%d.6" % class_c(env_name),
        ])
    elif platform.system() == "Darwin":
        subprocess.call([
            "sudo",
            "route",
            "-n",
            "add",
            "10.244.0.0/16",
            "192.168.%d.6" % class_c(env_name),
        ])
    else:
        eprint("unable to detect platform")
        sys.exit(1)


def tunnel(env_name, args):
    if len(args) != 3:
        eprint(USAGE)
        sys.exit(1)

    deployment = args[0]
    instance_group = args[1]
    tunnel = args[2]

    f = open(os.devnull, 'w')

    eprint("*******************************")
    eprint("****** Use ctrl-d to stop *****")
    eprint("*******************************")

    p = subprocess.Popen([
        bosh_cmd(),
        "--environment",
        env_name,
        "--deployment",
        deployment,
        "ssh",
        instance_group,
        "--opts=-L {}".format(tunnel),
    ],
    stdout=f,
    ).wait()


def run_ssh_command(env_name, deployment, instance_group, cmd):
    p = subprocess.Popen([
        bosh_cmd(),
        "--environment",
        env_name,
        "--deployment",
        deployment,
        "ssh",
        instance_group,
        "--command",
        cmd,
    ],
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE,
    )
    p.wait()

    vms_out = dict()
    vms_err = dict()
    for line in p.stdout:
        m = re.search('(.+)\s+(stdout|stderr)\s+\|\s+(.+)', line)
        if m is None or len(m.groups()) < 3:
            continue

        name = m.groups()[0]
        outType = m.groups()[1]
        output = m.groups()[2]

        vms = vms_out if outType == "stdout" else vms_err

        if name not in vms:
            vms[name]=[]

        vms[name].append(output.strip())
    return (vms_out, vms_err)


def ssh_command(env_name, args):
    if len(args) != 3:
        eprint(USAGE)
        sys.exit(1)

    deployment = args[0]
    instance_group = args[1]
    cmd = args[2]

    if instance_group == "*":
        instance_group = ""

    vms_out, vms_err = run_ssh_command(env_name, deployment, instance_group, cmd)

    def print_vms(vms, title_color, printer):
      for name, outputs in vms.items():
        printer("{}{}{}\n".format(title_color, name, bcolors.ENDC))
        for output in outputs:
          printer("{}\n".format(output))

    print_vms(vms_out, bcolors.OKGREEN, print)
    print_vms(vms_err, bcolors.FAIL, eprint)


def proc_status(env_name, deployment, instance_group):
    class Process:
        name = ""
        cpu = ""
        mem = ""
        pid = ""

    def parsestatus(status):
        procs = []
        for procblock in re.split('\nProcess\s+', status):
            procname = re.findall("'(.+)'", procblock)
            if len(procname) == 0:
                continue

            proc = Process()
            procs.append(proc)
            proc.name = procname[0]

            for line in procblock.split('\n'):
                # TODO: This is a bit kludgy
                cpu = re.findall("^\s*cpu percent\s+([0-9]+\.[0-9]+)", line)
                mem = re.findall("^\s*memory percent\s+([0-9]+\.[0-9]+)", line)
                pid = re.findall("^\s*pid\s+([0-9]+)", line)
                if len(cpu) != 0:
                    proc.cpu = cpu[0]
                    continue

                if len(mem) != 0:
                    proc.mem = mem[0]
                    continue

                if len(pid) != 0:
                    proc.pid = pid[0]
                    continue
        return procs

    cmd = "sudo /var/vcap/bosh/bin/monit status"
    statuses, _ = run_ssh_command(env_name, deployment, instance_group, cmd)

    def join(l):
        result = ""
        for i in l:
            result = result+"\n"+i
        return result

    result = {}
    for vm, status in statuses.iteritems():
        result[vm] = parsestatus(join(status))

    return result


def print_proc_status(env_name, args):
    if len(args) != 2:
        eprint(USAGE)
        sys.exit(1)

    deployment = args[0]
    instance_group = args[1]

    if instance_group == "*":
        instance_group = ""

    def label(name, value):
        return "{}{}{} {}%".format(bcolors.FAIL, name, bcolors.ENDC, value)

    for vm, status in proc_status(env_name, deployment, instance_group).iteritems():
      print("{}{}{}\n".format(bcolors.OKGREEN, vm, bcolors.ENDC))
      for p in status:
          name = "{}{}{}".format(bcolors.HEADER, p.name, bcolors.ENDC)
          cpu = label("CPU", p.cpu)
          mem = label("MEM", p.mem)
          print("{} {} {}\n".format(name, cpu, mem))


def listening_ports(env_name, deployment, instance_group, pid):
    cmd = "sudo lsof -p {} | grep LISTEN".format(pid)
    out, _ = run_ssh_command(env_name, deployment, instance_group, cmd)

    results = []
    for vm, output in out.iteritems():
        for line in output:
            ports = re.findall('\:([0-9]+)\s*\(LISTEN\)', line)
            if len(ports) != 0:
                results.append(ports[0])
                continue

    return results


def parse_go_routines(env_name, deployment, instance_group, port):
    cmd = "curl localhost:{}/debug/pprof/".format(port)
    out, _ = run_ssh_command(env_name, deployment, instance_group, cmd)

    for vm, lines in out.iteritems():
        for line in lines:
            goroutines = re.findall('\<tr\>\<td align\=right\>([0-9]+)\<td\>\<a href\=\"goroutine\?debug\=1\"\>goroutine\<\/a\>', line)
            if len(goroutines) != 0:
                return int(goroutines[0])

    return None


def go_routines(env_name, args):
    if len(args) != 2:
        eprint(USAGE)
        sys.exit(1)

    deployment = args[0]
    instance_group = args[1]

    if instance_group == "*":
        instance_group = ""

    for vm, ps in proc_status(env_name, deployment, instance_group).iteritems():
        print("{}{}{}\n".format(bcolors.OKGREEN, vm, bcolors.ENDC))
        for p in ps:
            ports = listening_ports(env_name, deployment, instance_group, p.pid)
            for port in ports:
                goroutines = parse_go_routines(env_name, deployment, instance_group, port)
                if goroutines is not None:
                    name = "{}{}{}".format(bcolors.HEADER, p.name, bcolors.ENDC)
                    print("{} {}\n".format(name, goroutines))
                    break


def main():
    # Validate bosh command is present and executable
    if not os.path.isfile(bosh_cmd()) or not os.access(bosh_cmd(), os.X_OK):
        eprint("'{}' is not a valid bosh-cli".format(bosh_cmd()))
        eprint("please set BOSH_CMD to a valid bosh-cli path")
        sys.exit(1)

    action, env_name, args = parse_action()
    if action == "create":
        create_env(env_name, args)
    elif action == "delete":
        delete_env(env_name, args)
    elif action == "set-routes":
        set_routes(env_name, args)
    elif action == "unset-routes":
        unset_routes(env_name)
    elif action == "tunnel":
        tunnel(env_name, args)
    elif action == "ssh-command":
        ssh_command(env_name, args)
    elif action == "proc-status":
        print_proc_status(env_name, args)
    elif action == "go-routines":
        go_routines(env_name, args)
    else:
        eprint(USAGE)
        sys.exit(1)


def parse_action():
    if len(sys.argv) < 3:
        eprint(USAGE)
        sys.exit(1)
    action, env_name = sys.argv[1], sys.argv[2]
    if action not in ["create", "delete", "set-routes", "unset-routes", "tunnel", "ssh-command", "proc-status", "go-routines"]:
        eprint(USAGE)
        sys.exit(1)
    for c in BAD_ENV_NAME_CHARS:
        if c in env_name:
            eprint("invalid character name in env name:", c)
            sys.exit(1)

    return action, env_name, sys.argv[3:]


if __name__ == "__main__":
    main()