tflocal

#!/usr/bin/env python

"""
Thin wrapper around the "terraform" command line interface (CLI) for use
with LocalStack.

The "tflocal" CLI allows you to easily interact with your local services
without having to specify the local endpoints in the "provider" section of
your TF config.
"""

import os
import sys
import glob
import subprocess
import json
import textwrap

from packaging import version
from urllib.parse import urlparse
from typing import Iterable, Optional

PARENT_FOLDER = os.path.realpath(os.path.join(os.path.dirname(__file__), ".."))
if os.path.isdir(os.path.join(PARENT_FOLDER, ".venv")):
    sys.path.insert(0, PARENT_FOLDER)

from localstack_client import config  # noqa: E402
import hcl2  # noqa: E402

DRY_RUN = str(os.environ.get("DRY_RUN")).strip().lower() in ["1", "true"]
DEFAULT_REGION = "us-east-1"
DEFAULT_ACCESS_KEY = "test"
AWS_ENDPOINT_URL = os.environ.get("AWS_ENDPOINT_URL")
CUSTOMIZE_ACCESS_KEY = str(os.environ.get("CUSTOMIZE_ACCESS_KEY")).strip().lower() in [
    "1",
    "true",
]
LOCALHOST_HOSTNAME = "localhost.localstack.cloud"
S3_HOSTNAME = os.environ.get("S3_HOSTNAME") or f"s3.{LOCALHOST_HOSTNAME}"
USE_EXEC = str(os.environ.get("USE_EXEC")).strip().lower() in ["1", "true"]
TF_CMD = os.environ.get("TF_CMD") or "terraform"
ADDITIONAL_TF_OVERRIDE_LOCATIONS = os.environ.get(
    "ADDITIONAL_TF_OVERRIDE_LOCATIONS", default=""
)
TF_UNPROXIED_CMDS = (
    os.environ.get("TF_UNPROXIED_CMDS").split(sep=",")
    if os.environ.get("TF_UNPROXIED_CMDS")
    else ("fmt", "validate", "version")
)
LS_PROVIDERS_FILE = (
    os.environ.get("LS_PROVIDERS_FILE") or "localstack_providers_override.tf"
)
LOCALSTACK_HOSTNAME = (
    urlparse(AWS_ENDPOINT_URL).hostname
    or os.environ.get("LOCALSTACK_HOSTNAME")
    or "localhost"
)
EDGE_PORT = int(urlparse(AWS_ENDPOINT_URL).port or os.environ.get("EDGE_PORT") or 4566)
TF_VERSION: Optional[version.Version] = None
TF_PROVIDER_CONFIG = """
provider "aws" {
  access_key                  = "<access_key>"
  secret_key                  = "test"
  skip_credentials_validation = true
  skip_metadata_api_check     = true
  <configs>
  endpoints {
<endpoints>
 }
}
"""
TF_S3_BACKEND_CONFIG = """
terraform {
  backend "s3" {<configs>
  }
}
"""
PROCESS = None
# some services have aliases which are mutually exclusive to each other
# see https://registry.terraform.io/providers/hashicorp/aws/latest/docs/guides/custom-service-endpoints#available-endpoint-customizations
SERVICE_ALIASES = [
    ("amp", "prometheus", "prometheusservice"),
    ("appautoscaling", "applicationautoscaling"),
    ("appintegrations", "appintegrationsservice"),
    ("ce", "costexplorer"),
    ("cloudcontrol", "cloudcontrolapi"),
    ("cloudhsmv2", "cloudhsm"),
    ("cognitoidp", "cognitoidentityprovider"),
    ("configservice", "config"),
    ("cur", "costandusagereportservice"),
    ("deploy", "codedeploy"),
    ("dms", "databasemigration", "databasemigrationservice"),
    ("ds", "directoryservice"),
    ("elasticbeanstalk", "beanstalk"),
    ("elasticsearch", "es", "elasticsearchservice"),
    ("elb", "elasticloadbalancing"),
    ("elbv2", "elasticloadbalancingv2"),
    ("events", "eventbridge", "cloudwatchevents"),
    ("evidently", "cloudwatchevidently"),
    ("grafana", "managedgrafana", "amg"),
    ("inspector2", "inspectorv2"),
    ("kafka", "msk"),
    ("lexmodels", "lexmodelbuilding", "lexmodelbuildingservice", "lex"),
    ("lexv2models", "lexmodelsv2"),
    ("location", "locationservice"),
    ("logs", "cloudwatchlog", "cloudwatchlogs"),
    ("oam", "cloudwatchobservabilityaccessmanager"),
    ("opensearch", "opensearchservice"),
    ("osis", "opensearchingestion"),
    ("rbin", "recyclebin"),
    ("redshiftdata", "redshiftdataapiservice"),
    ("resourcegroupstaggingapi", "resourcegroupstagging"),
    ("rum", "cloudwatchrum"),
    ("s3", "s3api"),
    ("serverlessrepo", "serverlessapprepo", "serverlessapplicationrepository"),
    ("servicecatalogappregistry", "appregistry"),
    ("sfn", "stepfunctions"),
    ("simpledb", "sdb"),
    ("transcribe", "transcribeservice"),
]
# service names to be excluded (not yet available in TF)
SERVICE_EXCLUSIONS = ["meteringmarketplace"]
# maps services to be replaced with alternative names
# skip services which do not have equivalent endpoint overrides
# see https://registry.terraform.io/providers/hashicorp/aws/latest/docs/guides/custom-service-endpoints
SERVICE_REPLACEMENTS = {
    "apigatewaymanagementapi": "",
    "appconfigdata": "",
    "ce": "costexplorer",
    "dynamodbstreams": "",
    "edge": "",
    "emrserverless": "",
    "iotdata": "",
    "ioteventsdata": "",
    "iotjobsdata": "",
    "iotwireless": "",
    "logs": "cloudwatchlogs",
    "mediastoredata": "",
    "qldbsession": "",
    "rdsdata": "",
    "sagemakerruntime": "",
    "support": "",
    "timestream": "",
    "timestreamquery": "",
}


# ---
# CONFIG GENERATION UTILS
# ---


def create_provider_config_file(provider_file_path: str, provider_aliases=None) -> None:
    provider_aliases = provider_aliases or []

    # Force service alias replacements
    SERVICE_REPLACEMENTS.update(
        {
            alias: alias_pairs[0]
            for alias_pairs in SERVICE_ALIASES
            for alias in alias_pairs
            if alias != alias_pairs[0]
        }
    )

    # create list of service names
    services = list(config.get_service_ports())
    services = [srvc for srvc in services if srvc not in SERVICE_EXCLUSIONS]
    services = [s.replace("-", "") for s in services]
    for old, new in SERVICE_REPLACEMENTS.items():
        try:
            services.remove(old)
            if new and new not in services:
                services.append(new)
        except ValueError:
            pass
    services = sorted(services)

    # add default (non-aliased) provider, if not defined yet
    default_provider = [p for p in provider_aliases if not p.get("alias")]
    if not default_provider:
        provider_aliases.append({"region": get_region()})

    # create provider configs
    provider_configs = []
    for provider in provider_aliases:
        provider_config = TF_PROVIDER_CONFIG.replace(
            "<access_key>",
            get_access_key(provider) if CUSTOMIZE_ACCESS_KEY else DEFAULT_ACCESS_KEY,
        )
        endpoints = "\n".join(
            [f'    {s} = "{get_service_endpoint(s)}"' for s in services]
        )
        provider_config = provider_config.replace("<endpoints>", endpoints)
        additional_configs = []
        if use_s3_path_style():
            additional_configs += [" s3_use_path_style = true"]
        alias = provider.get("alias")
        if alias:
            if isinstance(alias, list):
                alias = alias[0]
            additional_configs += [f' alias = "{alias}"']
        region = provider.get("region") or get_region()
        if isinstance(region, list):
            region = region[0]
        additional_configs += [f'region = "{region}"']
        provider_config = provider_config.replace(
            "<configs>", "\n".join(additional_configs)
        )
        provider_configs.append(provider_config)

    # construct final config file content
    tf_config = "\n".join(provider_configs)

    # create s3 backend config
    tf_config += generate_s3_backend_config()

    # write temporary config file
    write_provider_config_file(provider_file_path, tf_config)


def write_provider_config_file(providers_file, tf_config):
    """Write provider config into file"""
    with open(providers_file, mode="w") as fp:
        fp.write(tf_config)


def get_default_provider_folder_path() -> str:
    """Determine the folder under which the providers override file should be stored"""
    chdir = [arg for arg in sys.argv if arg.startswith("-chdir=")]
    base_dir = "."
    if chdir:
        base_dir = chdir[0].removeprefix("-chdir=")

    return os.path.abspath(base_dir)


def get_providers_file_path(base_dir) -> str:
    """Retrieve the path under which the providers override file should be stored"""
    return os.path.join(base_dir, LS_PROVIDERS_FILE)


def determine_provider_aliases() -> list:
    """Return a list of providers (and aliases) configured in the *.tf files (if any)"""
    skipped = str(os.environ.get("SKIP_ALIASES") or "").strip().split(",")
    result = []
    tf_files = parse_tf_files()
    for _file, obj in tf_files.items():
        try:
            providers = ensure_list(obj.get("provider", []))
            aws_providers = [
                prov["aws"]
                for prov in providers
                if prov.get("aws") and prov.get("aws").get("alias") not in skipped
            ]
            result.extend(aws_providers)
        except Exception as e:
            print(f"Warning: Unable to extract providers from {_file}:", e)
    return result


def generate_s3_backend_config() -> str:
    """Generate an S3 `backend {..}` block with local endpoints, if configured"""
    is_tf_legacy = TF_VERSION < version.Version("1.6")
    backend_config = None
    tf_files = parse_tf_files()
    for filename, obj in tf_files.items():
        if LS_PROVIDERS_FILE == filename:
            continue
        tf_configs = ensure_list(obj.get("terraform", []))
        for tf_config in tf_configs:
            backend_config = ensure_list(tf_config.get("backend"))
            if backend_config:
                backend_config = backend_config[0]
                break
    backend_config = backend_config and backend_config.get("s3")
    if not backend_config:
        return ""

    legacy_endpoint_mappings = {
        "endpoint": "s3",
        "iam_endpoint": "iam",
        "sts_endpoint": "sts",
        "dynamodb_endpoint": "dynamodb",
    }

    configs = {
        # note: default values, updated by `backend_config` further below...
        "bucket": "tf-test-state",
        "key": "terraform.tfstate",
        "dynamodb_table": "tf-test-state",
        "region": get_region(),
        "skip_credentials_validation": True,
        "skip_metadata_api_check": True,
        "secret_key": "test",
        "endpoints": {
            "s3": get_service_endpoint("s3"),
            "iam": get_service_endpoint("iam"),
            "sso": get_service_endpoint("sso"),
            "sts": get_service_endpoint("sts"),
            "dynamodb": get_service_endpoint("dynamodb"),
        },
    }
    # Merge in legacy endpoint configs if not existing already
    if is_tf_legacy and backend_config.get("endpoints"):
        print(
            "Warning: Unsupported backend option(s) detected (`endpoints`). Please make sure you always use the corresponding options to your Terraform version."
        )
        exit(1)
    for legacy_endpoint, endpoint in legacy_endpoint_mappings.items():
        if (
            legacy_endpoint in backend_config
            and backend_config.get("endpoints")
            and endpoint in backend_config["endpoints"]
        ):
            del backend_config[legacy_endpoint]
            continue
        if legacy_endpoint in backend_config and (
            not backend_config.get("endpoints")
            or endpoint not in backend_config["endpoints"]
        ):
            if not backend_config.get("endpoints"):
                backend_config["endpoints"] = {}
            backend_config["endpoints"].update(
                {endpoint: backend_config[legacy_endpoint]}
            )
            del backend_config[legacy_endpoint]
    # Add any missing default endpoints
    if backend_config.get("endpoints"):
        backend_config["endpoints"] = {
            k: backend_config["endpoints"].get(k) or v
            for k, v in configs["endpoints"].items()
        }
    backend_config["access_key"] = (
        get_access_key(backend_config) if CUSTOMIZE_ACCESS_KEY else DEFAULT_ACCESS_KEY
    )
    configs.update(backend_config)
    if not DRY_RUN:
        get_or_create_bucket(configs["bucket"])
        get_or_create_ddb_table(configs["dynamodb_table"], region=configs["region"])
    result = TF_S3_BACKEND_CONFIG
    config_options = ""
    for key, value in sorted(configs.items()):
        if isinstance(value, bool):
            value = str(value).lower()
        elif isinstance(value, dict):
            if key == "endpoints" and is_tf_legacy:
                for legacy_endpoint, endpoint in legacy_endpoint_mappings.items():
                    config_options += (
                        f'\n    {legacy_endpoint} = "{configs[key][endpoint]}"'
                    )
                continue
            else:
                value = textwrap.indent(
                    text=f"{key} = {{\n"
                    + "\n".join([f'  {k} = "{v}"' for k, v in value.items()])
                    + "\n}",
                    prefix=" " * 4,
                )
                config_options += f"\n{value}"
                continue
        elif isinstance(value, list):
            # TODO this will break if it's a list of dicts or other complex object
            # this serialization logic should probably be moved to a separate recursive function
            as_string = [f'"{item}"' for item in value]
            value = f"[{', '.join(as_string)}]"
        else:
            value = f'"{str(value)}"'
        config_options += f"\n    {key} = {value}"
    result = result.replace("<configs>", config_options)
    return result


def check_override_file(providers_file: str) -> None:
    """Checks override file existance"""
    if os.path.exists(providers_file):
        msg = f"Providers override file {providers_file} already exists"
        err_msg = msg + " - please delete it first, exiting..."
        if DRY_RUN:
            msg += ". File will be overwritten."
            print(msg)
            print("\tOnly 'yes' will be accepted to approve.")
            if input("\tEnter a value: ") == "yes":
                return
        print(err_msg)
        exit(1)


# ---
# AWS CLIENT UTILS
# ---


def use_s3_path_style() -> bool:
    """
    Whether to use S3 path addressing (depending on the configured S3 endpoint)
    If the endpoint starts with the `s3.` prefix, LocalStack will recognize virtual host addressing. If the endpoint
    does not start with it, use path style. This also allows overriding the endpoint to always use path style in case of
    inter container communications in Docker.
    """
    try:
        host = urlparse(get_service_endpoint("s3")).hostname
    except ValueError:
        host = ""

    return not host.startswith("s3.")


def get_region() -> str:
    region = str(os.environ.get("AWS_DEFAULT_REGION") or "").strip()
    if region:
        return region
    try:
        # If boto3 is installed, try to get the region from local credentials.
        # Note that boto3 is currently not included in the dependencies, to
        # keep the library lightweight.
        import boto3

        region = boto3.session.Session().region_name
    except Exception:
        pass
    # fall back to default region
    return region or DEFAULT_REGION


def get_access_key(provider: dict) -> str:
    access_key = str(
        os.environ.get("AWS_ACCESS_KEY_ID") or provider.get("access_key", "")
    ).strip()
    if access_key and access_key != DEFAULT_ACCESS_KEY:
        # Change live access key to mocked one
        return deactivate_access_key(access_key)
    try:
        # If boto3 is installed, try to get the access_key from local credentials.
        # Note that boto3 is currently not included in the dependencies, to
        # keep the library lightweight.
        import boto3

        access_key = boto3.session.Session().get_credentials().access_key
    except Exception:
        pass
    # fall back to default region
    return deactivate_access_key(access_key or DEFAULT_ACCESS_KEY)


def deactivate_access_key(access_key: str) -> str:
    """Safe guarding user from accidental live credential usage by deactivating access key IDs.
    See more: https://docs.localstack.cloud/references/credentials/"""
    return "L" + access_key[1:] if access_key[0] == "A" else access_key


def get_service_endpoint(service: str) -> str:
    """Get the service endpoint URL for the given service name"""
    # allow configuring a custom endpoint via the environment
    env_name = f"{service.replace('-', '_').upper().strip()}_ENDPOINT"
    env_endpoint = os.environ.get(env_name, "").strip()
    if env_endpoint:
        if "://" not in env_endpoint:
            env_endpoint = f"http://{env_endpoint}"
        return env_endpoint

    # some services need specific hostnames
    hostname = LOCALSTACK_HOSTNAME
    if service == "s3":
        hostname = S3_HOSTNAME
    elif service == "mwaa":
        hostname = f"mwaa.{LOCALHOST_HOSTNAME}"

    return f"http://{hostname}:{EDGE_PORT}"


def connect_to_service(service: str, region: str = None):
    import boto3

    region = region or get_region()
    return boto3.client(
        service,
        endpoint_url=get_service_endpoint(service),
        region_name=region,
        aws_access_key_id="test",
        aws_secret_access_key="test",
    )


def get_or_create_bucket(bucket_name: str):
    """Get or create a bucket in the current region."""
    s3_client = connect_to_service("s3")
    try:
        return s3_client.head_bucket(Bucket=bucket_name)
    except Exception:
        region = s3_client.meta.region_name
        kwargs = {}
        if region != "us-east-1":
            kwargs = {"CreateBucketConfiguration": {"LocationConstraint": region}}
        return s3_client.create_bucket(Bucket=bucket_name, **kwargs)


def get_or_create_ddb_table(table_name: str, region: str = None):
    """Get or create a DynamoDB table with the given name."""
    ddb_client = connect_to_service("dynamodb", region=region)
    try:
        return ddb_client.describe_table(TableName=table_name)
    except Exception:
        return ddb_client.create_table(
            TableName=table_name,
            BillingMode="PAY_PER_REQUEST",
            KeySchema=[{"AttributeName": "LockID", "KeyType": "HASH"}],
            AttributeDefinitions=[{"AttributeName": "LockID", "AttributeType": "S"}],
        )


# ---
# TF UTILS
# ---
def is_override_needed(args) -> bool:
    if any(map(lambda x: x in args, TF_UNPROXIED_CMDS)):
        return False
    return True


def parse_tf_files() -> dict:
    """Parse the local *.tf files and return a dict of <filename> -> <resource_dict>"""
    result = {}
    for _file in glob.glob("*.tf"):
        try:
            with open(_file, "r") as fp:
                result[_file] = hcl2.load(fp)
        except Exception as e:
            print(f'Unable to parse "{_file}" as HCL file: {e}')
    return result


def get_tf_version(env):
    global TF_VERSION
    output = subprocess.run(
        [f"{TF_CMD}", "version", "-json"], env=env, check=True, capture_output=True
    ).stdout.decode("utf-8")
    TF_VERSION = version.parse(json.loads(output)["terraform_version"])


def run_tf_exec(cmd, env):
    """Run terraform using os.exec - can be useful as it does not require any I/O
    handling for stdin/out/err. Does *not* allow us to perform any cleanup logic."""
    os.execvpe(cmd[0], cmd, env=env)


def run_tf_subprocess(cmd, env):
    """Run terraform in a subprocess - useful to perform cleanup logic at the end."""
    global PROCESS

    # register signal handlers
    import signal

    signal.signal(signal.SIGINT, signal_handler)

    PROCESS = subprocess.Popen(
        cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stdout
    )
    PROCESS.communicate()
    sys.exit(PROCESS.returncode)


def cleanup_override_files(override_files: Iterable[str]):
    for file_path in override_files:
        try:
            os.remove(file_path)
        except Exception:
            print(
                f"Count not clean up '{file_path}'. This is not normally a problem but you can delete this file manually."
            )


def get_folder_paths_that_require_an_override_file() -> Iterable[str]:
    if not is_override_needed(sys.argv[1:]):
        return

    yield get_default_provider_folder_path()
    for path in ADDITIONAL_TF_OVERRIDE_LOCATIONS.split(sep=","):
        if path.strip():
            yield path


# ---
# UTIL FUNCTIONS
# ---


def signal_handler(sig, frame):
    PROCESS.send_signal(sig)


def ensure_list(obj) -> list:
    return obj if isinstance(obj, list) else [obj]


def to_bytes(obj) -> bytes:
    return obj.encode("UTF-8") if isinstance(obj, str) else obj


def to_str(obj) -> bytes:
    return obj.decode("UTF-8") if isinstance(obj, bytes) else obj


# ---
# MAIN ENTRYPOINT
# ---


def main():
    env = dict(os.environ)
    cmd = [TF_CMD] + sys.argv[1:]

    try:
        get_tf_version(env)
        if not TF_VERSION:
            raise ValueError
    except (FileNotFoundError, ValueError) as e:
        print(f"Unable to determine version. See error message for details: {e}")
        exit(1)

    config_override_files = []

    for folder_path in get_folder_paths_that_require_an_override_file():
        config_file_path = get_providers_file_path(folder_path)
        check_override_file(config_file_path)

        providers = determine_provider_aliases()
        create_provider_config_file(config_file_path, providers)
        config_override_files.append(config_file_path)

    # call terraform command if not dry-run or any of the commands
    if not DRY_RUN or not config_override_files:
        try:
            if USE_EXEC:
                run_tf_exec(cmd, env)
            else:
                run_tf_subprocess(cmd, env)
        finally:
            cleanup_override_files(config_override_files)


if __name__ == "__main__":
    main()