dbt-platform-helper 11.0.1__py3-none-any.whl → 11.1.0__py3-none-any.whl

dbt_platform_helper/domain/database_copy.py

@@ -0,0 +1,220 @@
+import re
+from collections.abc import Callable
+from pathlib import Path
+
+import boto3
+import click
+from boto3 import Session
+
+from dbt_platform_helper.constants import PLATFORM_CONFIG_FILE
+from dbt_platform_helper.domain.maintenance_page import MaintenancePageProvider
+from dbt_platform_helper.exceptions import AWSException
+from dbt_platform_helper.utils.application import Application
+from dbt_platform_helper.utils.application import ApplicationNotFoundError
+from dbt_platform_helper.utils.application import load_application
+from dbt_platform_helper.utils.aws import Vpc
+from dbt_platform_helper.utils.aws import get_connection_string
+from dbt_platform_helper.utils.aws import get_vpc_info_by_name
+from dbt_platform_helper.utils.messages import abort_with_error
+from dbt_platform_helper.utils.validation import load_and_validate_platform_config
+
+
+class DatabaseCopy:
+    def __init__(
+        self,
+        app: str,
+        database: str,
+        auto_approve: bool = False,
+        load_application_fn: Callable[[str], Application] = load_application,
+        vpc_config_fn: Callable[[Session, str, str, str], Vpc] = get_vpc_info_by_name,
+        db_connection_string_fn: Callable[
+            [Session, str, str, str, Callable], str
+        ] = get_connection_string,
+        maintenance_page_provider: Callable[
+            [str, str, list[str], str, str], None
+        ] = MaintenancePageProvider(),
+        input_fn: Callable[[str], str] = click.prompt,
+        echo_fn: Callable[[str], str] = click.secho,
+        abort_fn: Callable[[str], None] = abort_with_error,
+    ):
+        self.app = app
+        self.database = database
+        self.auto_approve = auto_approve
+        self.vpc_config_fn = vpc_config_fn
+        self.db_connection_string_fn = db_connection_string_fn
+        self.maintenance_page_provider = maintenance_page_provider
+        self.input_fn = input_fn
+        self.echo_fn = echo_fn
+        self.abort_fn = abort_fn
+
+        if not self.app:
+            if not Path(PLATFORM_CONFIG_FILE).exists():
+                self.abort_fn("You must either be in a deploy repo, or provide the --app option.")
+
+            config = load_and_validate_platform_config(disable_aws_validation=True)
+            self.app = config["application"]
+
+        try:
+            self.application = load_application_fn(self.app)
+        except ApplicationNotFoundError:
+            abort_fn(f"No such application '{app}'.")
+
+    def _execute_operation(self, is_dump: bool, env: str, vpc_name: str):
+        vpc_name = self.enrich_vpc_name(env, vpc_name)
+
+        environments = self.application.environments
+        environment = environments.get(env)
+        if not environment:
+            self.abort_fn(
+                f"No such environment '{env}'. Available environments are: {', '.join(environments.keys())}"
+            )
+
+        env_session = environment.session
+
+        try:
+            vpc_config = self.vpc_config_fn(env_session, self.app, env, vpc_name)
+        except AWSException as ex:
+            self.abort_fn(str(ex))
+
+        database_identifier = f"{self.app}-{env}-{self.database}"
+
+        try:
+            db_connection_string = self.db_connection_string_fn(
+                env_session, self.app, env, database_identifier
+            )
+        except Exception as exc:
+            self.abort_fn(f"{exc} (Database: {database_identifier})")
+
+        try:
+            task_arn = self.run_database_copy_task(
+                env_session, env, vpc_config, is_dump, db_connection_string
+            )
+        except Exception as exc:
+            self.abort_fn(f"{exc} (Account id: {self.account_id(env)})")
+
+        if is_dump:
+            message = f"Dumping {self.database} from the {env} environment into S3"
+        else:
+            message = f"Loading data into {self.database} in the {env} environment from S3"
+
+        self.echo_fn(message, fg="white", bold=True)
+        self.echo_fn(
+            f"Task {task_arn} started. Waiting for it to complete (this may take some time)...",
+            fg="white",
+        )
+        self.tail_logs(is_dump, env)
+
+    def enrich_vpc_name(self, env, vpc_name):
+        if not vpc_name:
+            if not Path(PLATFORM_CONFIG_FILE).exists():
+                self.abort_fn(
+                    "You must either be in a deploy repo, or provide the vpc name option."
+                )
+            config = load_and_validate_platform_config(disable_aws_validation=True)
+            vpc_name = config.get("environments", {}).get(env, {}).get("vpc")
+        return vpc_name
+
+    def run_database_copy_task(
+        self,
+        session: boto3.session.Session,
+        env: str,
+        vpc_config: Vpc,
+        is_dump: bool,
+        db_connection_string: str,
+    ) -> str:
+        client = session.client("ecs")
+        action = "dump" if is_dump else "load"
+        env_vars = [
+            {"name": "DATA_COPY_OPERATION", "value": action.upper()},
+            {"name": "DB_CONNECTION_STRING", "value": db_connection_string},
+        ]
+        if not is_dump:
+            env_vars.append({"name": "ECS_CLUSTER", "value": f"{self.app}-{env}"})
+
+        response = client.run_task(
+            taskDefinition=f"arn:aws:ecs:eu-west-2:{self.account_id(env)}:task-definition/{self.app}-{env}-{self.database}-{action}",
+            cluster=f"{self.app}-{env}",
+            capacityProviderStrategy=[
+                {"capacityProvider": "FARGATE", "weight": 1, "base": 0},
+            ],
+            networkConfiguration={
+                "awsvpcConfiguration": {
+                    "subnets": vpc_config.subnets,
+                    "securityGroups": vpc_config.security_groups,
+                    "assignPublicIp": "DISABLED",
+                }
+            },
+            overrides={
+                "containerOverrides": [
+                    {
+                        "name": f"{self.app}-{env}-{self.database}-{action}",
+                        "environment": env_vars,
+                    }
+                ]
+            },
+        )
+
+        return response.get("tasks", [{}])[0].get("taskArn")
+
+    def dump(self, env: str, vpc_name: str):
+        self._execute_operation(True, env, vpc_name)
+
+    def load(self, env: str, vpc_name: str):
+        if self.is_confirmed_ready_to_load(env):
+            self._execute_operation(False, env, vpc_name)
+
+    def copy(
+        self,
+        from_env: str,
+        to_env: str,
+        from_vpc: str,
+        to_vpc: str,
+        services: tuple[str],
+        template: str,
+        no_maintenance_page: bool = False,
+    ):
+        to_vpc = self.enrich_vpc_name(to_env, to_vpc)
+        if not no_maintenance_page:
+            self.maintenance_page_provider.activate(self.app, to_env, services, template, to_vpc)
+        self.dump(from_env, from_vpc)
+        self.load(to_env, to_vpc)
+        if not no_maintenance_page:
+            self.maintenance_page_provider.deactivate(self.app, to_env)
+
+    def is_confirmed_ready_to_load(self, env: str) -> bool:
+        if self.auto_approve:
+            return True
+
+        user_input = self.input_fn(
+            f"\nWARNING: the load operation is destructive and will delete the {self.database} database in the {env} environment. Continue? (y/n)"
+        )
+        return user_input.lower().strip() in ["y", "yes"]
+
+    def tail_logs(self, is_dump: bool, env: str):
+        action = "dump" if is_dump else "load"
+        log_group_name = f"/ecs/{self.app}-{env}-{self.database}-{action}"
+        log_group_arn = f"arn:aws:logs:eu-west-2:{self.account_id(env)}:log-group:{log_group_name}"
+        self.echo_fn(f"Tailing {log_group_name} logs", fg="yellow")
+        session = self.application.environments[env].session
+        response = session.client("logs").start_live_tail(logGroupIdentifiers=[log_group_arn])
+
+        stopped = False
+        for data in response["responseStream"]:
+            if stopped:
+                break
+            results = data.get("sessionUpdate", {}).get("sessionResults", [])
+            for result in results:
+                message = result.get("message")
+
+                if message:
+                    match = re.match(r"(Stopping|Aborting) data (load|dump).*", message)
+                    if match:
+                        if match.group(1) == "Aborting":
+                            self.abort_fn("Task aborted abnormally. See logs above for details.")
+                        stopped = True
+                    self.echo_fn(message)
+
+    def account_id(self, env):
+        envs = self.application.environments
+        if env in envs:
+            return envs.get(env).account_id

dbt_platform_helper/domain/maintenance_page.py

@@ -0,0 +1,459 @@
+import itertools
+import random
+import re
+import string
+from pathlib import Path
+from typing import List
+from typing import Union
+
+import boto3
+import click
+
+from dbt_platform_helper.providers.load_balancers import ListenerNotFoundError
+from dbt_platform_helper.providers.load_balancers import ListenerRuleNotFoundError
+from dbt_platform_helper.providers.load_balancers import LoadBalancerNotFoundError
+from dbt_platform_helper.providers.load_balancers import find_https_listener
+from dbt_platform_helper.utils.application import Environment
+from dbt_platform_helper.utils.application import Service
+from dbt_platform_helper.utils.application import load_application
+
+
+class MaintenancePageProvider:
+    def activate(self, app, env, svc, template, vpc):
+        application = load_application(app)
+        application_environment = get_app_environment(app, env)
+
+        if "*" in svc:
+            services = [
+                s for s in application.services.values() if s.kind == "Load Balanced Web Service"
+            ]
+        else:
+            all_services = [get_app_service(app, s) for s in list(svc)]
+            services = [s for s in all_services if s.kind == "Load Balanced Web Service"]
+
+        if not services:
+            click.secho(f"No services deployed yet to {app} environment {env}", fg="red")
+            raise click.Abort
+
+        try:
+            https_listener = find_https_listener(application_environment.session, app, env)
+            current_maintenance_page = get_maintenance_page(
+                application_environment.session, https_listener
+            )
+            remove_current_maintenance_page = False
+            if current_maintenance_page:
+                remove_current_maintenance_page = click.confirm(
+                    f"There is currently a '{current_maintenance_page}' maintenance page for the {env} "
+                    f"environment in {app}.\nWould you like to replace it with a '{template}' "
+                    f"maintenance page?"
+                )
+                if not remove_current_maintenance_page:
+                    raise click.Abort
+
+            if remove_current_maintenance_page or click.confirm(
+                f"You are about to enable the '{template}' maintenance page for the {env} "
+                f"environment in {app}.\nWould you like to continue?"
+            ):
+                if current_maintenance_page and remove_current_maintenance_page:
+                    remove_maintenance_page(application_environment.session, https_listener)
+
+                allowed_ips = get_env_ips(vpc, application_environment)
+
+                add_maintenance_page(
+                    application_environment.session,
+                    https_listener,
+                    app,
+                    env,
+                    services,
+                    allowed_ips,
+                    template,
+                )
+                click.secho(
+                    f"Maintenance page '{template}' added for environment {env} in application {app}",
+                    fg="green",
+                )
+            else:
+                raise click.Abort
+
+        except LoadBalancerNotFoundError:
+            click.secho(
+                f"No load balancer found for environment {env} in the application {app}.", fg="red"
+            )
+            raise click.Abort
+
+        except ListenerNotFoundError:
+            click.secho(
+                f"No HTTPS listener found for environment {env} in the application {app}.", fg="red"
+            )
+            raise click.Abort
+
+    def deactivate(self, app, env):
+        application_environment = get_app_environment(app, env)
+
+        try:
+            https_listener = find_https_listener(application_environment.session, app, env)
+            current_maintenance_page = get_maintenance_page(
+                application_environment.session, https_listener
+            )
+            if not current_maintenance_page:
+                click.secho("There is no current maintenance page to remove", fg="red")
+                raise click.Abort
+
+            if not click.confirm(
+                f"There is currently a '{current_maintenance_page}' maintenance page, "
+                f"would you like to remove it?"
+            ):
+                raise click.Abort
+
+            remove_maintenance_page(application_environment.session, https_listener)
+            click.secho(
+                f"Maintenance page removed from environment {env} in application {app}", fg="green"
+            )
+
+        except LoadBalancerNotFoundError:
+            click.secho(
+                f"No load balancer found for environment {env} in the application {app}.", fg="red"
+            )
+            raise click.Abort
+
+        except ListenerNotFoundError:
+            click.secho(
+                f"No HTTPS listener found for environment {env} in the application {app}.", fg="red"
+            )
+            raise click.Abort
+
+
+def get_app_service(app_name: str, svc_name: str) -> Service:
+    application = load_application(app_name)
+    application_service = application.services.get(svc_name)
+
+    if not application_service:
+        click.secho(
+            f"The service {svc_name} was not found in the application {app_name}. "
+            f"It either does not exist, or has not been deployed.",
+            fg="red",
+        )
+        raise click.Abort
+
+    return application_service
+
+
+def get_app_environment(app_name: str, env_name: str) -> Environment:
+    application = load_application(app_name)
+    application_environment = application.environments.get(env_name)
+
+    if not application_environment:
+        click.secho(
+            f"The environment {env_name} was not found in the application {app_name}. "
+            f"It either does not exist, or has not been deployed.",
+            fg="red",
+        )
+        raise click.Abort
+
+    return application_environment
+
+
+def get_maintenance_page(session: boto3.Session, listener_arn: str) -> Union[str, None]:
+    lb_client = session.client("elbv2")
+
+    rules = lb_client.describe_rules(ListenerArn=listener_arn)["Rules"]
+    tag_descriptions = get_rules_tag_descriptions(rules, lb_client)
+
+    maintenance_page_type = None
+    for description in tag_descriptions:
+        tags = {t["Key"]: t["Value"] for t in description["Tags"]}
+        if tags.get("name") == "MaintenancePage":
+            maintenance_page_type = tags.get("type")
+
+    return maintenance_page_type
+
+
+def remove_maintenance_page(session: boto3.Session, listener_arn: str):
+    lb_client = session.client("elbv2")
+
+    rules = lb_client.describe_rules(ListenerArn=listener_arn)["Rules"]
+    # TODO: The next line doesn't appear to do anything.
+    tag_descriptions = get_rules_tag_descriptions(rules, lb_client)
+    # TODO: In fact the following line seems to do the same but better.
+    tag_descriptions = lb_client.describe_tags(ResourceArns=[r["RuleArn"] for r in rules])[
+        "TagDescriptions"
+    ]
+
+    for name in ["MaintenancePage", "AllowedIps", "BypassIpFilter", "AllowedSourceIps"]:
+        deleted = delete_listener_rule(tag_descriptions, name, lb_client)
+
+        if name == "MaintenancePage" and not deleted:
+            raise ListenerRuleNotFoundError()
+
+
+def get_rules_tag_descriptions(rules: list, lb_client):
+    tag_descriptions = []
+    chunk_size = 20
+
+    for i in range(0, len(rules), chunk_size):
+        chunk = rules[i : i + chunk_size]
+        resource_arns = [r["RuleArn"] for r in chunk]
+        response = lb_client.describe_tags(ResourceArns=resource_arns)
+        tag_descriptions.extend(response["TagDescriptions"])
+
+    return tag_descriptions
+
+
+def delete_listener_rule(tag_descriptions: list, tag_name: str, lb_client: boto3.client):
+    current_rule_arn = None
+
+    for description in tag_descriptions:
+        tags = {t["Key"]: t["Value"] for t in description["Tags"]}
+        if tags.get("name") == tag_name:
+            current_rule_arn = description["ResourceArn"]
+            if current_rule_arn:
+                lb_client.delete_rule(RuleArn=current_rule_arn)
+
+    return current_rule_arn
+
+
+def add_maintenance_page(
+    session: boto3.Session,
+    listener_arn: str,
+    app: str,
+    env: str,
+    services: List[Service],
+    allowed_ips: tuple,
+    template: str = "default",
+):
+    lb_client = session.client("elbv2")
+    maintenance_page_content = get_maintenance_page_template(template)
+    bypass_value = "".join(random.choices(string.ascii_lowercase + string.digits, k=12))
+
+    rule_priority = itertools.count(start=1)
+
+    for svc in services:
+        target_group_arn = find_target_group(app, env, svc.name, session)
+
+        # not all of an application's services are guaranteed to have been deployed to an environment
+        if not target_group_arn:
+            continue
+
+        for ip in allowed_ips:
+            create_header_rule(
+                lb_client,
+                listener_arn,
+                target_group_arn,
+                "X-Forwarded-For",
+                [ip],
+                "AllowedIps",
+                next(rule_priority),
+            )
+            create_source_ip_rule(
+                lb_client,
+                listener_arn,
+                target_group_arn,
+                [ip],
+                "AllowedSourceIps",
+                next(rule_priority),
+            )
+
+        create_header_rule(
+            lb_client,
+            listener_arn,
+            target_group_arn,
+            "Bypass-Key",
+            [bypass_value],
+            "BypassIpFilter",
+            next(rule_priority),
+        )
+
+        click.secho(
+            f"\nUse a browser plugin to add `Bypass-Key` header with value {bypass_value} to your requests. For more detail, visit https://platform.readme.trade.gov.uk/activities/holding-and-maintenance-pages/",
+            fg="green",
+        )
+
+    lb_client.create_rule(
+        ListenerArn=listener_arn,
+        Priority=next(rule_priority),
+        Conditions=[
+            {
+                "Field": "path-pattern",
+                "PathPatternConfig": {"Values": ["/*"]},
+            }
+        ],
+        Actions=[
+            {
+                "Type": "fixed-response",
+                "FixedResponseConfig": {
+                    "StatusCode": "503",
+                    "ContentType": "text/html",
+                    "MessageBody": maintenance_page_content,
+                },
+            }
+        ],
+        Tags=[
+            {"Key": "name", "Value": "MaintenancePage"},
+            {"Key": "type", "Value": template},
+        ],
+    )
+
+
+def get_maintenance_page_template(template) -> str:
+    template_contents = (
+        Path(__file__)
+        .parent.parent.joinpath(
+            f"templates/svc/maintenance_pages/{template}.html",
+        )
+        .read_text()
+        .replace("\n", "")
+    )
+
+    # [^\S]\s+ - Remove any space that is not preceded by a non-space character.
+    return re.sub(r"[^\S]\s+", "", template_contents)
+
+
+def find_target_group(app: str, env: str, svc: str, session: boto3.Session) -> str:
+    rg_tagging_client = session.client("resourcegroupstaggingapi")
+    response = rg_tagging_client.get_resources(
+        TagFilters=[
+            {
+                "Key": "copilot-application",
+                "Values": [
+                    app,
+                ],
+                "Key": "copilot-environment",
+                "Values": [
+                    env,
+                ],
+                "Key": "copilot-service",
+                "Values": [
+                    svc,
+                ],
+            },
+        ],
+        ResourceTypeFilters=[
+            "elasticloadbalancing:targetgroup",
+        ],
+    )
+    for resource in response["ResourceTagMappingList"]:
+        tags = {tag["Key"]: tag["Value"] for tag in resource["Tags"]}
+
+        if (
+            "copilot-service" in tags
+            and tags["copilot-service"] == svc
+            and "copilot-environment" in tags
+            and tags["copilot-environment"] == env
+            and "copilot-application" in tags
+            and tags["copilot-application"] == app
+        ):
+            return resource["ResourceARN"]
+
+    click.secho(
+        f"No target group found for application: {app}, environment: {env}, service: {svc}",
+        fg="red",
+    )
+
+    return None
+
+
+def create_header_rule(
+    lb_client: boto3.client,
+    listener_arn: str,
+    target_group_arn: str,
+    header_name: str,
+    values: list,
+    rule_name: str,
+    priority: int,
+):
+    conditions = get_host_conditions(lb_client, listener_arn, target_group_arn)
+
+    # add new condition to existing conditions
+    combined_conditions = [
+        {
+            "Field": "http-header",
+            "HttpHeaderConfig": {"HttpHeaderName": header_name, "Values": values},
+        }
+    ] + conditions
+
+    lb_client.create_rule(
+        ListenerArn=listener_arn,
+        Priority=priority,
+        Conditions=combined_conditions,
+        Actions=[{"Type": "forward", "TargetGroupArn": target_group_arn}],
+        Tags=[
+            {"Key": "name", "Value": rule_name},
+        ],
+    )
+
+    click.secho(
+        f"Creating listener rule {rule_name} for HTTPS Listener with arn {listener_arn}.\n\nIf request header {header_name} contains one of the values {values}, the request will be forwarded to target group with arn {target_group_arn}.",
+        fg="green",
+    )
+
+
+def create_source_ip_rule(
+    lb_client: boto3.client,
+    listener_arn: str,
+    target_group_arn: str,
+    values: list,
+    rule_name: str,
+    priority: int,
+):
+    conditions = get_host_conditions(lb_client, listener_arn, target_group_arn)
+
+    # add new condition to existing conditions
+    combined_conditions = [
+        {
+            "Field": "source-ip",
+            "SourceIpConfig": {"Values": [value + "/32" for value in values]},
+        }
+    ] + conditions
+
+    lb_client.create_rule(
+        ListenerArn=listener_arn,
+        Priority=priority,
+        Conditions=combined_conditions,
+        Actions=[{"Type": "forward", "TargetGroupArn": target_group_arn}],
+        Tags=[
+            {"Key": "name", "Value": rule_name},
+        ],
+    )
+
+    click.secho(
+        f"Creating listener rule {rule_name} for HTTPS Listener with arn {listener_arn}.\n\nIf request source ip matches one of the values {values}, the request will be forwarded to target group with arn {target_group_arn}.",
+        fg="green",
+    )
+
+
+def get_host_conditions(lb_client: boto3.client, listener_arn: str, target_group_arn: str):
+    rules = lb_client.describe_rules(ListenerArn=listener_arn)["Rules"]
+
+    # Get current set of forwarding conditions for the target group
+    for rule in rules:
+        for action in rule["Actions"]:
+            if action["Type"] == "forward" and action["TargetGroupArn"] == target_group_arn:
+                conditions = rule["Conditions"]
+
+    # filter to host-header conditions
+    conditions = [
+        {i: condition[i] for i in condition if i != "Values"}
+        for condition in conditions
+        if condition["Field"] == "host-header"
+    ]
+
+    # remove internal hosts
+    conditions[0]["HostHeaderConfig"]["Values"] = [
+        v for v in conditions[0]["HostHeaderConfig"]["Values"]
+    ]
+
+    return conditions
+
+
+def get_env_ips(vpc: str, application_environment: Environment) -> List[str]:
+    account_name = f"{application_environment.session.profile_name}-vpc"
+    vpc_name = vpc if vpc else account_name
+    ssm_client = application_environment.session.client("ssm")
+
+    try:
+        param_value = ssm_client.get_parameter(Name=f"/{vpc_name}/EGRESS_IPS")["Parameter"]["Value"]
+    except ssm_client.exceptions.ParameterNotFound:
+        click.secho(f"No parameter found with name: /{vpc_name}/EGRESS_IPS")
+        raise click.Abort
+
+    return [ip.strip() for ip in param_value.split(",")]