#!/usr/bin/env python3
#
"""
Sets a single target group ARN in an ALB listener default action.

This disgusting hack was necessary to facilitate the migration of 
Web from EC2 to Fargate due to idiosynchrasies in Terraform.

See: https://jira.benchling.team/browse/BNCH-24083
"""
import argparse
import logging
import subprocess
import re
import json

import boto3


logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def parse_args():
    """
    Define an argument parser and return the parsed arguments
    """
    parser = argparse.ArgumentParser(
        prog="crumbled", description="works magic with alb listener default actions"
    )
    parser.add_argument("--name", help="the name of the deploy to patch", required=True)
    parser.add_argument(
        "--dry-run", help="do a dry-run", default=False, action="store_true"
    )
    parser.add_argument(
        "--verbose",
        help="increase logging verbosity",
        default=False,
        action="store_true",
    )
    return parser.parse_args()


def parse_arn(stream, target):
    """
    Parse out an arn from running `terraform state show ...`.

    Since the output of `terraform state show` embeds ASCII escape
    characters, we have to strip them out.
    """
    regex = re.compile(
        r"^\s+arn\s+= \"(arn:aws:elasticloadbalancing:.+:.+:{}.+)\"".format(target)
    )
    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")

    matches = []
    for line in stream.split("\n"):
        line = ansi_escape.sub("", line)
        match = re.findall(regex, line)
        if len(match) == 1:
            matches.append(match[0])

    if len(matches) > 1:
        raise RuntimeError(
            "Too many matches in input data: {}. Expected 1".format(len(matches))
        )

    return matches[0]


def main():
    args = parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    logger.warn(
        "If this script seems to hang, it may be because `aws-okta` is asking you for an auth token."
        " This script has no logic to handle that case interactively!"
    )

    logger.info("You selected {}".format(args.name))

    # Get the aws_lb_listener
    cmd = [
        "terraform",
        "state",
        "show",
        "module.{}.aws_lb_listener.web_https".format(args.name),
    ]

    logger.info("Running command -> `{}`".format(" ".join(cmd)))
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, check=True)
    logger.debug("\n{}".format(proc.stdout))

    arn_lb_listener = parse_arn(proc.stdout.decode(), "listener")
    logger.info("lb listener arn is {}".format(arn_lb_listener))

    # Get the lb_target_group
    cmd = [
        "terraform",
        "state",
        "show",
        "module.{}.aws_lb_target_group.web".format(args.name),
    ]

    logger.info("Running command -> `{}`".format(" ".join(cmd)))
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, check=True)
    logger.debug("\n{}".format(proc.stdout.decode()))

    arn_target_group = parse_arn(proc.stdout.decode(), "targetgroup")
    logger.info("lb target group arn is {}".format(arn_target_group))

    if args.dry_run:
        return

    # Modify the listener
    client = boto3.client("elbv2")

    response = client.modify_listener(
        ListenerArn=arn_lb_listener,
        DefaultActions=[{"Type": "forward", "TargetGroupArn": arn_target_group}],
    )
    print(json.dumps(response))


if __name__ == "__main__":
    main()

# EOF