extension.py

import logging

from localstack import config
from localstack.aws.chain import CompositeHandler
from localstack.extensions.api import http
from localstack.extensions.patterns.webapp import WebAppExtension
import typing as t

from localstack.services.internal import get_internal_apis

from aws_proxy.server.aws_request_forwarder import AwsProxyHandler
from aws_proxy.server.request_handler import RequestHandler, WebApp

LOG = logging.getLogger(__name__)


class AwsProxyExtension(WebAppExtension):
    name = "aws-proxy"

    def __init__(self):
        super().__init__(template_package_path=None)

    def on_extension_load(self):
        if config.GATEWAY_SERVER == "twisted":
            LOG.warning(
                "AWS Proxy: The aws-proxy extension currently requires hypercorn as "
                "gateway server. Please start localstack with GATEWAY_SERVER=hypercorn"
            )

    def update_gateway_routes(self, router: http.Router[http.RouteHandler]):
        super().update_gateway_routes(router)

        LOG.info("AWS Proxy: adding routes to activate extension")
        get_internal_apis().add(RequestHandler())

    def collect_routes(self, routes: list[t.Any]):
        routes.append(WebApp())

    def update_request_handlers(self, handlers: CompositeHandler):
        LOG.debug("AWS Proxy: adding AWS proxy handler to the request chain")
        handlers.handlers.append(AwsProxyHandler())