cli.py

#!/usr/bin/python3
from __future__ import annotations

import argparse
import os
import re
import sys

import boto3
import git
from botocore.credentials import InstanceMetadataProvider
from botocore.exceptions import ClientError
from botocore.utils import InstanceMetadataFetcher

from lib.aws.attacks import Attacks
from lib.aws.ingestor import IngestionManager
from lib.aws.ingestor import Ingestor
from lib.aws.profile import Profile
from lib.aws.resources import RESOURCES
from lib.graph.db import Neo4j
from lib.util.console import console

SERVICES = list(Ingestor.__subclasses__())


def handle_update(args):
    """
    Function to handle the following subcommand :
    awspx update
    This function uses gitpython to update the repository and the npm modules.
    """

    repo = git.Repo("/opt/awspx")
    head = repo.head.commit

    repo.remotes.origin.set_url("https://github.com/FSecureLABS/awspx.git")
    repo.git.reset("--hard")
    repo.remotes.origin.pull()

    if head == repo.head.commit:
        console.notice("Already up to date")
        return

    console.task(
        f"Updating to {repo.head.commit}",
        os.system,
        args=["cd /opt/awspx/www && npm install >/dev/null 2>&1"],
        done=f"Updated to {repo.head.commit}",
    )


def handle_profile(args, console=console):
    """
    Function to handle the following subcommand :
    awspx profile
    This function will create, list or delete a profile.
    """

    profile = Profile(console=console)

    if args.create_profile:
        profile.create(args.create_profile)
        if args.create_profile in Profile().credentials.sections():
            console.notice(f"Saved profile '{args.create_profile}'")
        else:
            sys.exit()

    elif args.list_profiles:
        profile.list()

    elif args.delete_profile:
        profile.delete(args.delete_profile)


def handle_ingest(args):
    """
    Function to handle the following subcommand :
    awspx ingest
    This function will use the specified credentials to gather info about the
    aws environmnent.
    """

    session = None

    # Get credentials from environment variables
    if args.env:
        session = boto3.session.Session(region_name=args.region)

    # Use existing profile
    elif args.profile in Profile().credentials.sections():
        session = boto3.session.Session(
            profile_name=args.profile, region_name=args.region
        )
    # Use instance profile
    elif args.profile == "default":
        try:
            provider = InstanceMetadataProvider(
                iam_role_fetcher=InstanceMetadataFetcher()
            )
            creds = provider.load()

            session = boto3.session.Session(
                region_name=args.region,
                aws_access_key_id=creds.access_key,
                aws_secret_access_key=creds.secret_key,
                aws_session_token=creds.token,
            )
        except:
            pass

    # Specified profile doesn't exist, offer to create it
    if not session:
        profile = console.item("Create profile") if args.pretty else console

        profile.notice(
            f"The profile '{args.profile}' doesn't exist. "
            "Please enter your AWS credentials.\n"
            "(this information will be saved automatically)"
        )

        args.create_profile = args.profile
        handle_profile(args, console=profile)

        session = boto3.session.Session(
            profile_name=args.profile, region_name=args.region
        )
    # Ancillary operations
    try:
        if args.mfa_device:
            session_token = session.client("sts").get_session_token(
                SerialNumber=args.mfa_device,
                TokenCode=args.mfa_token,
                DurationSeconds=args.mfa_duration,
            )["Credentials"]

            session = boto3.session.Session(
                aws_access_key_id=session_token["AccessKeyId"],
                aws_secret_access_key=session_token["SecretAccessKey"],
                aws_session_token=session_token["SessionToken"],
                region_name=args.region,
            )

        if args.role_to_assume:
            assume_role_args = {
                "RoleArn": args.role_to_assume,
                "RoleSessionName": "awspx",
                "DurationSeconds": args.role_to_assume_duration,
                **dict(
                    {"ExternalId": args.role_to_assume_external_id}
                    if args.role_to_assume_external_id
                    else {}
                ),
            }

            assumed_role = session.client("sts").assume_role(
                **assume_role_args
            )["Credentials"]

            session = boto3.session.Session(
                aws_access_key_id=assumed_role["AccessKeyId"],
                aws_secret_access_key=assumed_role["SecretAccessKey"],
                aws_session_token=assumed_role["SessionToken"],
                region_name=args.region,
            )

    except ClientError as e:
        console.critical(e)

    # Ingest profile data
    ingestor = IngestionManager(
        session=session,
        console=console,
        services=args.services,
        db=args.database,
        quick=args.quick,
        skip_actions=args.skip_actions_all,
        only_types=args.only_types,
        skip_types=args.skip_types,
        only_arns=args.only_arns,
        skip_arns=args.skip_arns,
    )

    assert ingestor.zip is not None, "Ingestion failed"

    args.load_zips = [ingestor.zip]

    # Create the db
    handle_db(args, console=console.item("Creating Database"))

    if not (args.skip_attacks_all or args.skip_actions_all):
        handle_attacks(args, console=console.item("Updating Attack paths"))


def handle_attacks(args, console=console):
    """
    Function to handle the following subcommand :
    awspx attacks
    This function will compute the different attack paths supported.
    """

    attacks = Attacks(
        skip_conditional_actions=args.include_conditional_attacks is False,
        skip_attacks=args.skip_attacks,
        only_attacks=args.only_attacks,
        max_search_depth=str(
            args.max_attack_depth if args.max_attack_depth is not None else ""
        ),
        console=console,
    )

    attacks.compute(max_iterations=args.max_attack_iterations)


def handle_db(args, console=console):
    """
    Function to handle the following subcommand :
    awspx db
    This function will load a db from a zip file, list the existing dbs or use
    an existing db.
    """

    db = Neo4j(console=console)

    if args.load_zips:
        db.load_zips(
            archives=args.load_zips,
            db=args.database if "database" in args else "default",
        )

    elif args.list_dbs:
        db.list()

    elif args.use_db:
        db.use(args.use_db)


def main():
    """Main function defining the arguments and calling the handler functions."""

    def profile(p):
        """Verify that the provided profile has the right format"""
        if p in list(Profile().credentials.sections()):
            raise argparse.ArgumentTypeError(f"profile '{p}' already exists")
        return p

    def database(p):
        """Verify that the provided database has the right format"""
        if re.compile("[A-Za-z0-9-_]+").match(p) is None:
            suggestion = re.sub(r"[^A-Za-z0-9-_]+", "", p)
            raise argparse.ArgumentTypeError(
                f"'{p}' is invalid, perhaps you meant '{suggestion}' instead?"
            )
        return p

    def service(service):
        """Verify that the provided service has the right format"""
        match = next(
            (
                s
                for s in SERVICES
                if s.__name__.upper() == service.replace(",", "").upper()
            ),
            None,
        )
        if match is None:
            raise argparse.ArgumentTypeError(
                f"'{service}' is not a supported service "
                f"(eg: {', '.join([str(v.__name__) for v in SERVICES])})"
            )
        return match

    def resource(resource):
        """Verify that the provided resource has the right format"""
        match = next(
            (r for r in RESOURCES if r.upper() == resource.upper()), None
        )
        if match is None:
            raise argparse.ArgumentTypeError(
                f"'{resource}' is not a supported resource type "
                "(see lib/aws/resources.py for details)"
            )
        return match

    def ARN(arn):
        """Verify that the provided arn has the right format"""
        if (
            re.compile(
                r"arn:aws:([a-zA-Z0-9]+):([a-z0-9-]*):(\d{12}|aws)?:(.*)"
            ).match(arn)
            is None
        ):
            raise argparse.ArgumentTypeError(f"'{arn}' is not a valid ARN")

        return arn

    def attack(name):
        """Verify that the provided attack has the right format"""
        match = next(
            (a for a in Attacks.definitions if a.upper() == name.upper()), None
        )
        if match is None:
            raise argparse.ArgumentTypeError(
                f"'{name}' is not a supported attack "
                "(see lib/aws/attacks.py for details)"
            )
        return match

    parser = argparse.ArgumentParser(
        prog="awspx",
        description=(
            "awspx is a graph-based tool for visualizing effective "
            "access and resource relationships within AWS. "
        ),
    )

    subparsers = parser.add_subparsers(title="commands")

    #
    # awspx update
    #
    update_parser = subparsers.add_parser(
        "update", help="Update awspx to the latest version."
    )
    update_parser.set_defaults(func=handle_update)

    #
    # awspx profile
    #
    profile_parser = subparsers.add_parser(
        "profile", help="Manage AWS credential profiles used for ingestion."
    )
    profile_parser.set_defaults(func=handle_profile)

    profile_group = profile_parser.add_mutually_exclusive_group(required=True)

    profile_group.add_argument(
        "--create",
        dest="create_profile",
        default=None,
        type=profile,
        help="Create a new profile using `aws configure`.",
    )
    profile_group.add_argument(
        "--list",
        dest="list_profiles",
        action="store_true",
        help="List saved profiles.",
    )
    profile_group.add_argument(
        "--delete",
        dest="delete_profile",
        choices=Profile().credentials.sections(),
        help="Delete a saved profile.",
    )
    #
    # awspx ingest
    #
    ingest_parser = subparsers.add_parser(
        "ingest", help="Ingest data from an AWS account."
    )
    ingest_parser.set_defaults(func=handle_ingest)

    # Profile & region args
    pnr = ingest_parser.add_argument_group("Profile and region")
    pnr.add_argument(
        "--env",
        action="store_true",
        help="Use AWS credential environment variables.",
    )
    pnr.add_argument(
        "--profile",
        dest="profile",
        default="default",
        help="""Profile to use for ingestion
        (corresponds to a `[section]` in `~/.aws/credentials).""",
    )
    pnr.add_argument(
        "--mfa-device",
        dest="mfa_device",
        help="ARN of the MFA device to authenticate with.",
    )
    pnr.add_argument(
        "--mfa-token", dest="mfa_token", help="Current MFA token."
    )
    pnr.add_argument(
        "--mfa-duration",
        dest="mfa_duration",
        type=int,
        default=3600,
        help="Maximum session duration in seconds (for MFA session).",
    )
    pnr.add_argument(
        "--assume-role",
        dest="role_to_assume",
        help="""ARN of a role to assume for ingestion
        (useful for cross-account ingestion).""",
    )
    pnr.add_argument(
        "--assume-role-duration",
        dest="role_to_assume_duration",
        type=int,
        default=3600,
        help="Maximum session duration in seconds (for --assume-role).",
    )
    pnr.add_argument(
        "--assume-role-external-id",
        dest="role_to_assume_external_id",
        help="External ID for the role to assume.",
    )
    pnr.add_argument(
        "--region",
        dest="region",
        default="eu-west-1",
        choices=Profile.regions,
        help="""Region to ingest
        (defaults to profile region, or `eu-west-1` if not set).""",
    )
    pnr.add_argument(
        "--database",
        dest="database",
        default=None,
        type=database,
        help="Database to store results (defaults to <profile>).",
    )

    # Services & resources args
    snr = ingest_parser.add_argument_group("Services and resources")
    snr.add_argument(
        "--services",
        dest="services",
        default=SERVICES,
        nargs="+",
        type=service,
        help=(
            f"One or more services to ingest (eg: {' '.join([s.__name__ for s in SERVICES])})."  # noqa: E501
        ),
    )
    snr.add_argument(
        "--quick",
        dest="quick",
        action="store_true",
        default=False,
        help=(
            "Skips supplementary ingestion functions "
            "(i.e. speed at the cost of information)."
        ),
    )

    type_args = snr.add_mutually_exclusive_group()
    type_args.add_argument(
        "--only-types",
        dest="only_types",
        default=[],
        nargs="+",
        type=resource,
        help="""Resource to include by type,
        all other resource types will be excluded.""",
    )
    type_args.add_argument(
        "--skip-types",
        dest="skip_types",
        nargs="+",
        default=[],
        type=resource,
        help="Resources to exclude by type.",
    )

    # ARN args
    arn_args = snr.add_mutually_exclusive_group()
    arn_args.add_argument(
        "--only-arns",
        dest="only_arns",
        default=[],
        nargs="+",
        type=ARN,
        help="""Resources to include by ARN,
        all other resources will be excluded.""",
    )
    arn_args.add_argument(
        "--skip-arns",
        dest="skip_arns",
        default=[],
        nargs="+",
        type=ARN,
        help="Resources to exclude by ARN.",
    )

    actions = ingest_parser.add_argument_group("Actions")
    actions.add_argument(
        "--skip-actions-all",
        dest="skip_actions_all",
        action="store_true",
        default=False,
        help="Skip policy resolution (actions will not be processed).",
    )

    #
    # awspx attacks
    #
    attacks_parser = subparsers.add_parser(
        "attacks", help="Compute attacks using the active database."
    )
    attacks_parser.set_defaults(func=handle_attacks)

    # Add args to ingest, attacks
    for p in [ingest_parser, attacks_parser]:
        ag = p.add_argument_group("Attack computation")
        g = ag.add_mutually_exclusive_group()

        g.add_argument(
            "--skip-attacks",
            dest="skip_attacks",
            default=[],
            nargs="+",
            type=attack,
            help="Attacks to exclude by name.",
        )
        g.add_argument(
            "--only-attacks",
            dest="only_attacks",
            default=[],
            nargs="+",
            type=attack,
            help="""Attacks to include by name,
            all other attacks will be excluded.""",
        )
        ag.add_argument(
            "--max-attack-iterations",
            dest="max_attack_iterations",
            default=5,
            type=int,
            help="""Maximum number of iterations to run each attack
            (default: 5).""",
        )
        ag.add_argument(
            "--max-attack-depth",
            dest="max_attack_depth",
            default=None,
            type=int,
            help="Maximum search depth for attacks (default: None).",
        )
        ag.add_argument(
            "--include-conditional-attacks",
            dest="include_conditional_attacks",
            action="store_true",
            default=False,
            help="""Include conditional actions when computing attacks
            (default: False).""",
        )

        if p is ingest_parser:
            ag.add_argument(
                "--skip-attacks-all",
                dest="skip_attacks_all",
                action="store_true",
                default=False,
                help="""Skip attack path computation
                (it can be run later with `awspx attacks`).""",
            )

    #
    # awspx db
    #
    db_parser = subparsers.add_parser(
        "db",
        help="""Manage databases used for visualization, ingestion,
        and attack computation.""",
    )

    db_parser.set_defaults(func=handle_db)

    db_group = db_parser.add_mutually_exclusive_group(required=True)

    db_group.add_argument(
        "--use",
        dest="use_db",
        choices=Neo4j.databases,
        help="Switch to the specified database.",
    )
    db_group.add_argument(
        "--list",
        dest="list_dbs",
        action="store_true",
        help="List available databases.",
    )
    db_group.add_argument(
        "--load-zip",
        dest="load_zips",
        choices=sorted(Neo4j.zips),
        action="append",
        help="Create/overwrite database using ZIP file content.",
    )

    # Add --pretty to ingest, attacks, db
    for p in [ingest_parser, attacks_parser, db_parser]:
        p.add_argument(
            "--pretty",
            dest="pretty",
            action="store_true",
            default=False,
            help="Enable pretty output (slower).",
        )

    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()

    # Unless a database has been defined for ingest, default to <profile>
    if "database" in args and args.database is None:
        args.database = f"{args.profile}"

    if "pretty" in args and not args.pretty:
        console.verbose()
    else:
        console.start()

    try:
        # Call the handler function defined in arguments
        args.func(args)

    except (KeyboardInterrupt, SystemExit):
        console.stop()
        os._exit(1)

    except BaseException as e:
        console.critical(e)
        os._exit(1)

    console.stop()


main()