auth_proxy.py

import json
import logging
import os
import re
import subprocess
import sys
from functools import cache
from io import BytesIO
from typing import Dict, Optional, Tuple
from urllib.parse import urlparse, urlunparse

import boto3
import requests
from botocore.awsrequest import AWSPreparedRequest, AWSResponse
from botocore.httpchecksum import resolve_checksum_context
from botocore.model import OperationModel
from localstack import config
from localstack import config as localstack_config
from localstack.aws.chain import HandlerChain
from localstack.aws.chain import RequestContext as AwsRequestContext
from localstack.aws.gateway import Gateway
from localstack.aws.protocol.parser import create_parser
from localstack.aws.spec import load_service
from localstack.config import external_service_url
from localstack.constants import AWS_REGION_US_EAST_1, DOCKER_IMAGE_NAME_PRO
from localstack.http import Request
from localstack.http import Response as HttpResponse
from localstack.http.hypercorn import GatewayServer
from localstack.utils.aws.aws_responses import requests_response
from localstack.utils.bootstrap import setup_logging
from localstack.utils.collections import select_attributes
from localstack.utils.container_utils.container_client import PortMappings
from localstack.utils.docker_utils import DOCKER_CLIENT, reserve_available_container_port
from localstack.utils.files import new_tmp_file, save_file
from localstack.utils.functions import run_safe
from localstack.utils.net import get_docker_host_from_container, get_free_tcp_port
from localstack.utils.serving import Server
from localstack.utils.strings import short_uid, to_bytes, to_str, truncate
from localstack_ext.bootstrap.licensingv2 import ENV_LOCALSTACK_API_KEY, ENV_LOCALSTACK_AUTH_TOKEN
from requests import Response

from aws_replicator import config as repl_config
from aws_replicator.client.utils import truncate_content
from aws_replicator.config import HANDLER_PATH_PROXIES
from aws_replicator.shared.models import AddProxyRequest, ProxyConfig

LOG = logging.getLogger(__name__)
LOG.setLevel(logging.INFO)
if config.DEBUG:
    LOG.setLevel(logging.DEBUG)

# TODO make configurable
CLI_PIP_PACKAGE = "localstack-extension-aws-replicator"
# note: enable the line below temporarily for testing:
# CLI_PIP_PACKAGE = "git+https://github.com/localstack/localstack-extensions/@branch#egg=localstack-extension-aws-replicator&subdirectory=aws-replicator"

CONTAINER_NAME_PREFIX = "ls-aws-proxy-"
CONTAINER_CONFIG_FILE = "/tmp/ls.aws.proxy.yml"
CONTAINER_LOG_FILE = "/tmp/ls-aws-proxy.log"

# default bind host if `bind_host` is not specified for the proxy
DEFAULT_BIND_HOST = "127.0.0.1"


class AwsProxyHandler:
    """
    A handler for an AWS Handler chain that attempts to forward the request using a specific boto3 session.
    This can be used to proxy incoming requests to real AWS.
    """

    def __init__(self, session: boto3.Session = None):
        self.session = session or boto3.Session()

    def __call__(self, chain: HandlerChain, context: AwsRequestContext, response: HttpResponse):
        # prepare the API invocation parameters
        LOG.info(
            "Received %s.%s = %s",
            context.service.service_name,
            context.operation.name,
            context.service_request,
        )

        # make the actual API call against upstream AWS (will also calculate a new auth signature)
        try:
            aws_response = self._make_aws_api_call(context)
        except Exception:
            LOG.exception(
                "Exception while proxying  %s.%s to AWS",
                context.service.service_name,
                context.operation.name,
            )
            raise

        # tell the handler chain to respond
        LOG.info(
            "AWS Response %s.%s: url=%s status_code=%s, headers=%s, content=%s",
            context.service.service_name,
            context.operation.name,
            aws_response.url,
            aws_response.status_code,
            aws_response.headers,
            aws_response.content,
        )
        chain.respond(aws_response.status_code, aws_response.content, dict(aws_response.headers))

    def _make_aws_api_call(self, context: AwsRequestContext) -> AWSResponse:
        # TODO: reconcile with AwsRequestProxy from localstack, and other forwarder tools
        # create a real AWS client
        client = self.session.client(context.service.service_name, region_name=context.region)
        operation_model = context.operation

        # prepare API request parameters as expected by boto
        api_params = {k: v for k, v in context.service_request.items() if v is not None}

        # this is a stripped down version of botocore's client._make_api_call to immediately get the HTTP
        # response instead of a parsed response.
        request_context = {
            "client_region": client.meta.region_name,
            "client_config": client.meta.config,
            "has_streaming_input": operation_model.has_streaming_input,
            "auth_type": operation_model.auth_type,
        }

        (
            endpoint_url,
            additional_headers,
            properties,
        ) = client._resolve_endpoint_ruleset(operation_model, api_params, request_context)
        if properties:
            # Pass arbitrary endpoint info with the Request
            # for use during construction.
            request_context["endpoint_properties"] = properties

        request_dict = client._convert_to_request_dict(
            api_params=api_params,
            operation_model=operation_model,
            endpoint_url=endpoint_url,
            context=request_context,
            headers=additional_headers,
        )
        resolve_checksum_context(request_dict, operation_model, api_params)

        if operation_model.has_streaming_input:
            request_dict["body"] = request_dict["body"].read()

        self._adjust_request_dict(context.service.service_name, request_dict)

        if operation_model.has_streaming_input:
            request_dict["body"] = BytesIO(request_dict["body"])

        LOG.info("Making AWS request %s", request_dict)
        http, _ = client._endpoint.make_request(operation_model, request_dict)

        http: AWSResponse

        # for some elusive reasons, these header modifications are needed (were part of http2_server)
        http.headers.pop("Date", None)
        http.headers.pop("Server", None)
        if operation_model.has_streaming_output:
            http.headers.pop("Content-Length", None)

        return http

    def _adjust_request_dict(self, service_name: str, request_dict: Dict):
        """Apply minor fixes to the request dict, which seem to be required in the current setup."""
        # TODO: replacing localstack-specific URLs, IDs, etc, should ideally be done in a more generalized
        #  way.

        req_body = request_dict.get("body")

        # TODO: fix for switch between path/host addressing
        # Note: the behavior seems to be different across botocore versions. Seems to be working
        # with 1.29.97 (fix below not required) whereas newer versions like 1.29.151 require the fix.
        if service_name == "s3":
            body_str = run_safe(lambda: to_str(req_body)) or ""

            request_url = request_dict["url"]
            url_parsed = list(urlparse(request_url))
            path_parts = url_parsed[2].strip("/").split("/")
            bucket_subdomain_prefix = f"://{path_parts[0]}.s3."
            if bucket_subdomain_prefix in request_url:
                prefix = f"/{path_parts[0]}"
                url_parsed[2] = url_parsed[2].removeprefix(prefix)
                request_dict["url_path"] = request_dict["url_path"].removeprefix(prefix)
                # replace empty path with "/" (seems required for signature calculation)
                request_dict["url_path"] = request_dict["url_path"] or "/"
                url_parsed[2] = url_parsed[2] or "/"
                # re-construct final URL
                request_dict["url"] = urlunparse(url_parsed)

            # TODO: this custom fix should not be required - investigate and remove!
            if "<CreateBucketConfiguration" in body_str and "LocationConstraint" not in body_str:
                region = request_dict["context"]["client_region"]
                if region == AWS_REGION_US_EAST_1:
                    request_dict["body"] = ""
                else:
                    request_dict["body"] = (
                        '<CreateBucketConfiguration xmlns="http://s3.amazonaws.com/doc/2006-03-01/">'
                        f"<LocationConstraint>{region}</LocationConstraint></CreateBucketConfiguration>"
                    )

        if service_name == "sqs" and isinstance(req_body, dict):
            account_id = self._query_account_id_from_aws()
            if "QueueUrl" in req_body:
                queue_name = req_body["QueueUrl"].split("/")[-1]
                req_body["QueueUrl"] = f"https://queue.amazonaws.com/{account_id}/{queue_name}"
            if "QueueOwnerAWSAccountId" in req_body:
                req_body["QueueOwnerAWSAccountId"] = account_id
        if service_name == "sqs" and request_dict.get("url"):
            req_json = run_safe(lambda: json.loads(body_str)) or {}
            account_id = self._query_account_id_from_aws()
            queue_name = req_json.get("QueueName")
            if account_id and queue_name:
                request_dict["url"] = f"https://queue.amazonaws.com/{account_id}/{queue_name}"
                req_json["QueueOwnerAWSAccountId"] = account_id
                request_dict["body"] = to_bytes(json.dumps(req_json))

    def _fix_headers(self, request: Request, service_name: str):
        if service_name == "s3":
            # fix the Host header, to avoid bucket addressing issues
            host = request.headers.get("Host") or ""
            regex = r"^(https?://)?([0-9.]+|localhost)(:[0-9]+)?"
            if re.match(regex, host):
                request.headers["Host"] = re.sub(regex, r"\1s3.localhost.localstack.cloud", host)
        request.headers.pop("Content-Length", None)
        request.headers.pop("x-localstack-request-url", None)
        request.headers.pop("X-Forwarded-For", None)
        request.headers.pop("X-Localstack-Tgt-Api", None)
        request.headers.pop("X-Moto-Account-Id", None)
        request.headers.pop("Remote-Addr", None)

    @cache
    def _query_account_id_from_aws(self) -> str:
        sts_client = self.session.client("sts")
        result = sts_client.get_caller_identity()
        return result["Account"]


class AwsProxyGateway(Gateway):
    """
    A handler chain that receives AWS requests, and proxies them transparently to upstream AWS using real
    credentials. It de-constructs the incoming request, and creates a new request signed with the AWS
    credentials configured in the environment.
    """

    def __init__(self) -> None:
        from localstack.aws import handlers

        super().__init__(
            request_handlers=[
                handlers.parse_service_name,
                handlers.content_decoder,
                handlers.add_region_from_header,
                handlers.add_account_id,
                handlers.parse_service_request,
                AwsProxyHandler(),
            ],
            exception_handlers=[
                handlers.log_exception,
                handlers.handle_internal_failure,
            ],
            context_class=AwsRequestContext,
        )


class AuthProxyAWS(Server):
    def __init__(self, config: ProxyConfig, port: int = None):
        self.config = config
        port = port or get_free_tcp_port()
        super().__init__(port=port)

    def do_run(self):
        self.register_in_instance()

        bind_host = self.config.get("bind_host") or DEFAULT_BIND_HOST
        srv = GatewayServer(AwsProxyGateway(), localstack_config.HostAndPort(bind_host, self.port))
        srv.start()
        srv.join()
        # proxy = run_server(port=self.port, bind_addresses=[bind_host], handler=self.proxy_request)
        # proxy.join()

    def proxy_request(self, request: Request, data: bytes) -> Response:
        parsed = self._extract_region_and_service(request.headers)
        if not parsed:
            return requests_response("", status_code=400)
        region_name, service_name = parsed
        query_string = to_str(request.query_string or "")

        LOG.debug(
            "Proxying request to %s (%s): %s %s %s",
            service_name,
            region_name,
            request.method,
            request.path,
            query_string,
        )

        request = Request(
            body=data,
            method=request.method,
            headers=request.headers,
            path=request.path,
            query_string=query_string,
        )
        session = boto3.Session()
        client = session.client(service_name, region_name=region_name)

        # fix headers (e.g., "Host") and create client
        self._fix_headers(request, service_name)

        # create request and request dict
        operation_model, aws_request, request_dict = self._parse_aws_request(
            request, service_name, region_name=region_name, client=client
        )

        # adjust request dict and fix certain edge cases in the request
        self._adjust_request_dict(service_name, request_dict)

        headers_truncated = {k: truncate(to_str(v)) for k, v in dict(aws_request.headers).items()}
        LOG.debug(
            "Sending request for service %s to AWS: %s %s - %s - %s",
            service_name,
            request.method,
            aws_request.url,
            truncate_content(request_dict.get("body"), max_length=500),
            headers_truncated,
        )
        try:
            # send request to upstream AWS
            result = client._endpoint.make_request(operation_model, request_dict)

            # create response object - TODO: to be replaced with localstack.http.Response over time
            response = requests_response(
                result[0].content,
                status_code=result[0].status_code,
                headers=dict(result[0].headers),
            )

            LOG.debug(
                "Received response for service %s from AWS: %s - %s",
                service_name,
                response.status_code,
                truncate_content(response.content, max_length=500),
            )
            return response
        except Exception as e:
            if LOG.isEnabledFor(logging.DEBUG):
                LOG.exception("Error when making request to AWS service %s: %s", service_name, e)
            return requests_response("", status_code=400)

    def register_in_instance(self):
        port = getattr(self, "port", None)
        if not port:
            raise Exception("Proxy currently not running")
        url = f"{external_service_url()}{HANDLER_PATH_PROXIES}"
        data = AddProxyRequest(port=port, config=self.config)
        LOG.debug("Registering new proxy in main container via: %s", url)
        try:
            response = requests.post(url, json=data)
            assert response.ok
            return response
        except Exception:
            LOG.warning(
                "Unable to register auth proxy - is LocalStack running with the extension enabled?"
            )
            raise

    def _parse_aws_request(
        self, request: Request, service_name: str, region_name: str, client
    ) -> Tuple[OperationModel, AWSPreparedRequest, Dict]:
        parser = create_parser(load_service(service_name))
        operation_model, parsed_request = parser.parse(request)
        request_context = {
            "client_region": region_name,
            "has_streaming_input": operation_model.has_streaming_input,
            "auth_type": operation_model.auth_type,
            "client_config": client.meta.config,
        }
        parsed_request = {} if parsed_request is None else parsed_request
        parsed_request = {k: v for k, v in parsed_request.items() if v is not None}

        # get endpoint info
        endpoint_info = client._resolve_endpoint_ruleset(
            operation_model, parsed_request, request_context
        )
        # switch for https://github.com/boto/botocore/commit/826b78c54dd87b9da368e9ab6017d8c4823b28c1
        if len(endpoint_info) == 3:
            endpoint_url, additional_headers, properties = endpoint_info
            if properties:
                request_context["endpoint_properties"] = properties
        else:
            endpoint_url, additional_headers = endpoint_info

        # create request dict
        request_dict = client._convert_to_request_dict(
            parsed_request,
            operation_model,
            endpoint_url=endpoint_url,
            context=request_context,
            headers=additional_headers,
        )

        # TODO: fix for switch between path/host addressing
        # Note: the behavior seems to be different across botocore versions. Seems to be working
        # with 1.29.97 (fix below not required) whereas newer versions like 1.29.151 require the fix.
        if service_name == "s3":
            request_url = request_dict["url"]
            url_parsed = list(urlparse(request_url))
            path_parts = url_parsed[2].strip("/").split("/")
            bucket_subdomain_prefix = f"://{path_parts[0]}.s3."
            if bucket_subdomain_prefix in request_url:
                prefix = f"/{path_parts[0]}"
                url_parsed[2] = url_parsed[2].removeprefix(prefix)
                request_dict["url_path"] = request_dict["url_path"].removeprefix(prefix)
                # replace empty path with "/" (seems required for signature calculation)
                request_dict["url_path"] = request_dict["url_path"] or "/"
                url_parsed[2] = url_parsed[2] or "/"
                # re-construct final URL
                request_dict["url"] = urlunparse(url_parsed)

        aws_request = client._endpoint.create_request(request_dict, operation_model)

        return operation_model, aws_request, request_dict

    def _adjust_request_dict(self, service_name: str, request_dict: Dict):
        """Apply minor fixes to the request dict, which seem to be required in the current setup."""
        # TODO: replacing localstack-specific URLs, IDs, etc, should ideally be done in a more generalized
        #  way.
        req_body = request_dict.get("body")

        if service_name == "s3":
            body_str = run_safe(lambda: to_str(req_body)) or ""

            # TODO: this custom fix should not be required - investigate and remove!
            if "<CreateBucketConfiguration" in body_str and "LocationConstraint" not in body_str:
                region = request_dict["context"]["client_region"]
                if region == AWS_REGION_US_EAST_1:
                    request_dict["body"] = ""
                else:
                    request_dict["body"] = (
                        '<CreateBucketConfiguration xmlns="http://s3.amazonaws.com/doc/2006-03-01/">'
                        f"<LocationConstraint>{region}</LocationConstraint></CreateBucketConfiguration>"
                    )

        if service_name == "sqs" and isinstance(req_body, dict):
            account_id = self._query_account_id_from_aws()
            if "QueueUrl" in req_body:
                queue_name = req_body["QueueUrl"].split("/")[-1]
                req_body["QueueUrl"] = f"https://queue.amazonaws.com/{account_id}/{queue_name}"
            if "QueueOwnerAWSAccountId" in req_body:
                req_body["QueueOwnerAWSAccountId"] = account_id
        if service_name == "sqs" and request_dict.get("url"):
            req_json = run_safe(lambda: json.loads(body_str)) or {}
            account_id = self._query_account_id_from_aws()
            queue_name = req_json.get("QueueName")
            if account_id and queue_name:
                request_dict["url"] = f"https://queue.amazonaws.com/{account_id}/{queue_name}"
                req_json["QueueOwnerAWSAccountId"] = account_id
                request_dict["body"] = to_bytes(json.dumps(req_json))

    def _fix_headers(self, request: Request, service_name: str):
        if service_name == "s3":
            # fix the Host header, to avoid bucket addressing issues
            host = request.headers.get("Host") or ""
            regex = r"^(https?://)?([0-9.]+|localhost)(:[0-9]+)?"
            if re.match(regex, host):
                request.headers["Host"] = re.sub(regex, r"\1s3.localhost.localstack.cloud", host)
        request.headers.pop("Content-Length", None)
        request.headers.pop("x-localstack-request-url", None)
        request.headers.pop("X-Forwarded-For", None)
        request.headers.pop("X-Localstack-Tgt-Api", None)
        request.headers.pop("X-Moto-Account-Id", None)
        request.headers.pop("Remote-Addr", None)

    def _extract_region_and_service(self, headers) -> Optional[Tuple[str, str]]:
        auth_header = headers.pop("Authorization", "")
        parts = auth_header.split("Credential=", maxsplit=1)
        if len(parts) < 2:
            return
        parts = parts[1].split("/")
        if len(parts) < 5:
            return
        return parts[2], parts[3]

    @cache
    def _query_account_id_from_aws(self) -> str:
        session = boto3.Session()
        sts_client = session.client("sts")
        result = sts_client.get_caller_identity()
        return result["Account"]


def start_aws_auth_proxy(config: ProxyConfig, port: int = None) -> AuthProxyAWS:
    setup_logging()
    proxy = AuthProxyAWS(config, port=port)
    proxy.start()
    return proxy


def start_aws_auth_proxy_in_container(
    config: ProxyConfig, env_vars: dict = None, port: int = None, quiet: bool = False
):
    """
    Run the auth proxy in a separate local container. This can help in cases where users
    are running into version/dependency issues on their host machines.
    """
    # TODO: Currently running a container and installing the extension on the fly - we
    #  should consider building pre-baked images for the extension in the future. Also,
    #  the new packaged CLI binary can help us gain more stability over time...

    logging.getLogger("localstack.utils.container_utils.docker_cmd_client").setLevel(logging.INFO)
    logging.getLogger("localstack.utils.docker_utils").setLevel(logging.INFO)
    logging.getLogger("localstack.utils.run").setLevel(logging.INFO)

    print("Proxy container is starting up...")

    # determine port mapping
    localstack_config.PORTS_CHECK_DOCKER_IMAGE = DOCKER_IMAGE_NAME_PRO
    port = port or reserve_available_container_port()
    ports = PortMappings()
    ports.add(port, port)

    # create container
    container_name = f"{CONTAINER_NAME_PREFIX}{short_uid()}"
    image_name = DOCKER_IMAGE_NAME_PRO
    DOCKER_CLIENT.create_container(
        image_name,
        name=container_name,
        entrypoint="",
        command=["bash", "-c", f"touch {CONTAINER_LOG_FILE}; tail -f {CONTAINER_LOG_FILE}"],
        ports=ports,
        additional_flags=repl_config.PROXY_DOCKER_FLAGS,
    )

    # start container in detached mode
    DOCKER_CLIENT.start_container(container_name, attach=False)

    # install extension CLI package
    venv_activate = ". .venv/bin/activate"
    command = [
        "bash",
        "-c",
        f"{venv_activate}; pip install --upgrade --no-deps '{CLI_PIP_PACKAGE}'",
    ]
    DOCKER_CLIENT.exec_in_container(container_name, command=command)

    # create config file in container
    config_file_host = new_tmp_file()
    save_file(config_file_host, json.dumps(config))
    DOCKER_CLIENT.copy_into_container(
        container_name, config_file_host, container_path=CONTAINER_CONFIG_FILE
    )

    # prepare environment variables
    env_var_names = [
        "DEBUG",
        "AWS_SECRET_ACCESS_KEY",
        "AWS_ACCESS_KEY_ID",
        "AWS_SESSION_TOKEN",
        "AWS_DEFAULT_REGION",
        ENV_LOCALSTACK_API_KEY,
        ENV_LOCALSTACK_AUTH_TOKEN,
    ]
    env_vars = env_vars or os.environ
    env_vars = select_attributes(dict(env_vars), env_var_names)

    # Determine target hostname - we make the host configurable via PROXY_LOCALSTACK_HOST,
    #  and if not configured then use get_docker_host_from_container() as a fallback.
    target_host = repl_config.PROXY_LOCALSTACK_HOST
    if not repl_config.PROXY_LOCALSTACK_HOST:
        target_host = get_docker_host_from_container()
    env_vars["LOCALSTACK_HOST"] = target_host

    try:
        print("Proxy container is ready.")
        command = f"{venv_activate}; localstack aws proxy -c {CONTAINER_CONFIG_FILE} -p {port} --host 0.0.0.0 > {CONTAINER_LOG_FILE} 2>&1"
        if quiet:
            DOCKER_CLIENT.exec_in_container(
                container_name, command=["bash", "-c", command], env_vars=env_vars, interactive=True
            )
        else:
            env_vars_list = []
            for key, value in env_vars.items():
                env_vars_list += ["-e", f"{key}={value}"]
            # note: using docker command directly, as our Docker client doesn't fully support log piping yet
            command = [
                "docker",
                "exec",
                "-it",
                *env_vars_list,
                container_name,
                "bash",
                "-c",
                command,
            ]
            subprocess.run(command, stdout=sys.stdout, stderr=sys.stderr)
    except KeyboardInterrupt:
        pass
    except Exception as e:
        LOG.info("Error: %s", e)
        if isinstance(e, subprocess.CalledProcessError):
            LOG.info("Error in called process - output: %s\n%s", e.stdout, e.stderr)
    finally:
        try:
            if repl_config.CLEANUP_PROXY_CONTAINERS:
                DOCKER_CLIENT.remove_container(container_name, force=True)
        except Exception as e:
            if "already in progress" not in str(e):
                raise