#!/usr/bin/env python3
# https://raw.githubusercontent.com/elpy1/ssm-tool/master/ssm-tool
import argparse
import os
import shlex
import textwrap

import boto3
from prettytable import PrettyTable


def get_profiles():
    try:
        return boto3.session.Session().available_profiles
    except Exception as e:
        raise SystemExit(f"get_profiles failed with {e.__class__.__name__}: {e}")


def get_ec2_name(iid, session):
    try:
        tags = session.resource("ec2").Instance(iid).tags
        name = [tag["Value"] for tag in tags if tag["Key"] == "Name"]
        return str(*name) if name else ""
    except Exception:
        return ""


def ssm_update_agent(ssmi, session, max_instances=50):
    try:
        for c in range(0, len(ssmi), max_instances):
            instances = [i.get("InstanceId") for i in ssmi[c : c + max_instances]]
            session.client("ssm").send_command(
                InstanceIds=instances,
                DocumentName="AWS-UpdateSSMAgent",
                DocumentVersion="$LATEST",
            )
    except Exception as e:
        raise SystemExit(f"ssm_update_agent failed with {e.__class__.__name__}: {e}")
    else:
        print("success")


def ssm_list_instances(platform, session):
    ssm = session.client("ssm")
    filters = {"Key": "PlatformTypes", "Values": platform}

    try:
        data = ssm.describe_instance_information(MaxResults=50, Filters=[filters])
        instances = data.get("InstanceInformationList")

        while data.get("NextToken"):
            data = ssm.describe_instance_information(
                MaxResults=50, Filters=[filters], NextToken=data["NextToken"]
            )
            instances.extend(data.get("InstanceInformationList"))

        for i in instances:
            i.update({"Name": get_ec2_name(i["InstanceId"], session)})

        return instances
    except Exception as e:
        raise SystemExit(f"ssm_list_instances failed with {e.__class__.__name__}: {e}")


def build_table(instances):
    table = PrettyTable(
        ["tag[name]", "instance", "ip address", "ssm-agent*", "platform"]
    )
    table.align = "l"
    data = list(
        zip(
            [i.get("Name")[:40] for i in instances],
            [i.get("InstanceId")[:40] for i in instances],
            [i.get("IPAddress")[:40] for i in instances],
            [i.get("IsLatestVersion") for i in instances],
            [i.get("PlatformName")[:40] for i in instances],
        )
    )

    for r in data:
        table.add_row(r)

    return table


def build_ssh_conf(instances, session, profile="env"):
    ssh_users = dict(centos="centos", amazon="ec2-user", ubuntu="ubuntu")
    conf_file = "{}/.ssh/ssmtool-{}".format(os.path.expanduser("~"), profile)

    try:
        os.umask(0)

        with open(
            os.open(conf_file, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o600), "w"
        ) as f:
            if profile == "env":
                f.write(
                    f"# config created without AWS profile set "
                    f"(must have appropriate env vars set when connecting)"
                )

            for i in instances:
                name = i.get("Name").replace(" ", "_")
                hostname = ".".join(i.get("ComputerName").split(".")[:2])
                platform = i.get("PlatformName").split()[0].lower()

                f.write(
                    f"\nHost {name} {i.get('IPAddress')} {hostname}\n"
                    f"  Hostname {i.get('InstanceId')}\n"
                )

                if any(u in platform for u in ssh_users):
                    f.write(f"  User {ssh_users[platform]}\n")
    except Exception as e:
        raise SystemExit(f"build_ssh_conf failed with {e.__class__.__name__}: {e}")
    else:
        print(f"\nssh config fragment generated and saved to -> {conf_file}")
    finally:
        f.close()


def exec_sh(command, session):
    envs = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]

    try:
        cmd = shlex.split(str(command))
        creds = [(c or "") for c in session.get_credentials().get_frozen_credentials()]

        for k, v in list(zip(envs, creds)):
            os.environ[k] = str(v)

        os.execvp(cmd[0], cmd)
    except OSError as e:
        raise SystemExit(f"could not execute {cmd[0]}: {os.strerror(e.errno)}")
    except Exception as e:
        raise SystemExit(f"exec_sh failed with {e.__class__.__name__}: {e}")


def parse_args():
    parser = argparse.ArgumentParser(
        description="ssm toolkit",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent(
            """
            A wrapper toolkit that interacts with AWS SSM and SSH to simplify
            instance visibility and access. Designed to work in conjunction with
            ssh-over-ssm[0] for access to linux-based instances.

            If run without a command the script will list all available instances.

            AWS credentials will be read from appropriate environment variables if
            provided. However, if using in confjunction with ssh-over-ssm it is
            recommended to use the --profile argument or AWS_PROFILE environment
            variable, which take precedence over the others. Examples below.

            For usage examples and more information see the README[1].

            [0] https://github.com/elpy1/ssh-over-ssm
            [1] https://github.com/elpy1/ssm-tool
            """
        ),
    )

    group = parser.add_mutually_exclusive_group()

    parser.add_argument(
        "--profile",
        dest="profile",
        action="store",
        default=os.getenv("AWS_PROFILE") or None,
        metavar="profile".upper(),
        choices=get_profiles(),
        help="AWS CLI profile to use (default: 'default')",
    )

    parser.add_argument(
        "-x",
        "--linux",
        dest="platforms",
        action="append_const",
        const="Linux",
        help="return only linux-based instances (default: all)",
    )

    parser.add_argument(
        "-w",
        "--windows",
        dest="platforms",
        action="append_const",
        const="Windows",
        help="return only windows-based instances (default: all)",
    )

    group.add_argument(
        "--text",
        dest="text",
        action="store_true",
        help="return data in plain text columns (default: false)",
    )

    group.add_argument("--iid", action="store_true", help="returns only instance ids")

    group.add_argument(
        "--update",
        dest="update",
        action="store_true",
        help="updates ssm-agent on instances",
    )

    group.add_argument(
        "--session",
        action="store",
        metavar="instance".upper(),
        help="an instance to connect to using ssm start-session",
    )

    group.add_argument(
        "--ssh",
        action="store",
        metavar="command".upper(),
        help="a ssh command to execute",
    )

    group.add_argument(
        "--ssh-conf",
        dest="sshconf",
        action="store_true",
        help="generates and saves a ssh config fragment",
    )

    parser.add_argument(
        "search",
        metavar="string".upper(),
        nargs="?",
        help="an optional search term to filter results",
    )

    return parser.parse_args()


def main():
    args = parse_args()
    session = boto3.session.Session(profile_name=args.profile)

    if args.session:
        exec_sh((f"aws ssm start-session --target {args.session}"), session)

    if args.ssh:
        exec_sh((f"ssh {args.ssh}"), session)

    instances = ssm_list_instances((args.platforms or ["Linux", "Windows"]), session)

    if args.search:
        search = args.search
        instances = [
            i
            for i in instances
            if any([search.lower() in str(v).lower() for v in i.values()])
        ]

    if args.sshconf:
        build_ssh_conf(instances, session, args.profile)
    elif args.update:
        ssm_update_agent(instances, session)
    elif args.iid:
        print(*[i.get("InstanceId") for i in instances], sep="\n")
    else:
        table = build_table(instances)

        if args.text:
            print(table.get_string(border=False, header=False))
        else:
            print(
                f"{table}\n * ssm-agent column refers to whether the agent is up-to-date"
            )


main()
