egress-proxy 0.1.0__py3-none-any.whl

egress_proxy/__init__.py

@@ -0,0 +1,10 @@
+"""egress-proxy: self-hosted egress proxy with drop-in integrations for Python HTTP clients.
+
+Source:  https://github.com/botifyex/egress-proxy
+Author:  botifyex
+License: Apache-2.0
+"""
+
+from importlib.metadata import version
+
+__version__ = version("egress-proxy")

egress_proxy/client/aiohttp/__init__.py

@@ -0,0 +1,8 @@
+from egress_proxy.extras import require
+
+require("client-aiohttp", "aiohttp", "yarl")
+
+from egress_proxy.client.aiohttp.factory import session
+from egress_proxy.client.aiohttp.middleware import Middleware, Refused
+
+__all__ = ["Middleware", "Refused", "session"]

egress_proxy/client/aiohttp/factory.py

@@ -0,0 +1,17 @@
+from collections.abc import Sequence
+from typing import Any
+
+import aiohttp
+
+from egress_proxy.client.aiohttp.middleware import Middleware
+
+
+def session(
+    *,
+    egress_url: str,
+    egress_key: str,
+    middlewares: Sequence[aiohttp.ClientMiddlewareType] = (),
+    **options: Any,
+) -> aiohttp.ClientSession:
+    egress = Middleware(egress_url=egress_url, egress_key=egress_key)
+    return aiohttp.ClientSession(middlewares=(*middlewares, egress), **options)

egress_proxy/client/aiohttp/middleware.py

@@ -0,0 +1,47 @@
+import aiohttp
+from yarl import URL
+
+from egress_proxy.protocol import Failure, Header
+
+
+class Refused(aiohttp.ClientError):
+    pass
+
+
+ERRORS: dict[Failure, type[aiohttp.ClientError]] = {
+    Failure.UNAUTHORIZED: Refused,
+    Failure.INVALID: Refused,
+    Failure.TIMEOUT: aiohttp.ServerTimeoutError,
+    Failure.UNREACHABLE: aiohttp.ClientConnectionError,
+}
+
+
+class Middleware:
+    def __init__(self, *, egress_url: str, egress_key: str) -> None:
+        self.url = URL(egress_url)
+        self.key = egress_key
+
+    def rewrite(self, request: aiohttp.ClientRequest) -> None:
+        if Header.TARGET in request.headers:
+            return
+        request.headers[Header.TARGET] = str(request.url)
+        request.headers[Header.KEY] = self.key
+        request.headers[aiohttp.hdrs.HOST] = self.url.host_port_subcomponent or ""
+        request.url = self.url
+
+    async def reject(self, response: aiohttp.ClientResponse) -> aiohttp.ClientError:
+        failure = Failure(response.headers[Header.FAILURE])
+        message = await response.text()
+        response.release()
+        return ERRORS[failure](message)
+
+    async def __call__(
+        self,
+        request: aiohttp.ClientRequest,
+        handler: aiohttp.ClientHandlerType,
+    ) -> aiohttp.ClientResponse:
+        self.rewrite(request)
+        response = await handler(request)
+        if Header.FAILURE in response.headers:
+            raise await self.reject(response)
+        return response

egress_proxy/client/httpx/__init__.py

@@ -0,0 +1,9 @@
+from egress_proxy.extras import require
+
+require("client-httpx", "httpx")
+
+from egress_proxy.client.httpx.client import AsyncClient, Client
+from egress_proxy.client.httpx.route import Refused
+from egress_proxy.client.httpx.transport import AsyncTransport, Transport
+
+__all__ = ["AsyncClient", "AsyncTransport", "Client", "Refused", "Transport"]

egress_proxy/client/httpx/client.py

@@ -0,0 +1,31 @@
+from typing import Any
+
+import httpx
+
+from egress_proxy.client.httpx.transport import AsyncTransport, Transport
+
+
+class Client(httpx.Client):
+    def __init__(
+        self,
+        *,
+        egress_url: str,
+        egress_key: str,
+        egress_transport: httpx.BaseTransport | None = None,
+        **options: Any,
+    ) -> None:
+        transport = Transport(egress_url=egress_url, egress_key=egress_key, egress_transport=egress_transport)
+        super().__init__(transport=transport, **options)
+
+
+class AsyncClient(httpx.AsyncClient):
+    def __init__(
+        self,
+        *,
+        egress_url: str,
+        egress_key: str,
+        egress_transport: httpx.AsyncBaseTransport | None = None,
+        **options: Any,
+    ) -> None:
+        transport = AsyncTransport(egress_url=egress_url, egress_key=egress_key, egress_transport=egress_transport)
+        super().__init__(transport=transport, **options)

egress_proxy/client/httpx/route.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+
+import httpx
+
+from egress_proxy.protocol import Failure, Header
+
+
+class Refused(httpx.TransportError):
+    pass
+
+
+ERRORS: dict[Failure, type[httpx.TransportError]] = {
+    Failure.UNAUTHORIZED: Refused,
+    Failure.INVALID: Refused,
+    Failure.TIMEOUT: httpx.TimeoutException,
+    Failure.UNREACHABLE: httpx.ConnectError,
+}
+
+
+@dataclass(frozen=True)
+class Route:
+    url: httpx.URL
+    key: str
+
+    def headers(self, request: httpx.Request) -> httpx.Headers:
+        headers = request.headers.copy()
+        headers["host"] = self.url.netloc.decode("ascii")
+        headers[Header.KEY] = self.key
+        headers[Header.TARGET] = str(request.url)
+        return headers
+
+    def rewrite(self, request: httpx.Request) -> httpx.Request:
+        return httpx.Request(
+            request.method,
+            self.url,
+            headers=self.headers(request),
+            stream=request.stream,
+            extensions=request.extensions,
+        )
+
+    def failed(self, response: httpx.Response) -> bool:
+        return Header.FAILURE in response.headers
+
+    def error(self, request: httpx.Request, response: httpx.Response) -> httpx.TransportError:
+        failure = Failure(response.headers[Header.FAILURE])
+        return ERRORS[failure](response.text, request=request)

egress_proxy/client/httpx/transport.py

@@ -0,0 +1,47 @@
+import httpx
+
+from egress_proxy.client.httpx.route import Route
+
+
+class Transport(httpx.BaseTransport):
+    def __init__(
+        self,
+        *,
+        egress_url: str,
+        egress_key: str,
+        egress_transport: httpx.BaseTransport | None = None,
+    ) -> None:
+        self.route = Route(httpx.URL(egress_url), egress_key)
+        self.inner = egress_transport or httpx.HTTPTransport()
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        response = self.inner.handle_request(self.route.rewrite(request))
+        if self.route.failed(response):
+            response.read()
+            raise self.route.error(request, response)
+        return response
+
+    def close(self) -> None:
+        self.inner.close()
+
+
+class AsyncTransport(httpx.AsyncBaseTransport):
+    def __init__(
+        self,
+        *,
+        egress_url: str,
+        egress_key: str,
+        egress_transport: httpx.AsyncBaseTransport | None = None,
+    ) -> None:
+        self.route = Route(httpx.URL(egress_url), egress_key)
+        self.inner = egress_transport or httpx.AsyncHTTPTransport()
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        response = await self.inner.handle_async_request(self.route.rewrite(request))
+        if self.route.failed(response):
+            await response.aread()
+            raise self.route.error(request, response)
+        return response
+
+    async def aclose(self) -> None:
+        await self.inner.aclose()

egress_proxy/extras.py

@@ -0,0 +1,13 @@
+from importlib import import_module
+
+
+class Missing(ModuleNotFoundError):
+    pass
+
+
+def require(extra: str, *libraries: str) -> None:
+    for library in libraries:
+        try:
+            import_module(library)
+        except ModuleNotFoundError as error:
+            raise Missing(f"{library} is not installed: pip install 'egress-proxy[{extra}]'") from error

egress_proxy/protocol.py

@@ -0,0 +1,14 @@
+from enum import StrEnum
+
+
+class Header(StrEnum):
+    KEY = "x-egress-key"
+    TARGET = "x-egress-target"
+    FAILURE = "x-egress-failure"
+
+
+class Failure(StrEnum):
+    UNAUTHORIZED = "unauthorized"
+    INVALID = "invalid"
+    TIMEOUT = "timeout"
+    UNREACHABLE = "unreachable"

egress_proxy/server/__init__.py

@@ -0,0 +1,3 @@
+from egress_proxy.extras import require
+
+require("server", "aiohttp", "certifi", "httptools", "pydantic_settings", "uvicorn", "yarl")

egress_proxy/server/addresses.py

@@ -0,0 +1,55 @@
+import ipaddress
+import socket
+from typing import TypeAlias
+
+from aiohttp.abc import AbstractResolver, ResolveResult
+from aiohttp.resolver import DefaultResolver
+
+Address: TypeAlias = ipaddress.IPv4Address | ipaddress.IPv6Address
+
+NAT64 = ipaddress.ip_network("64:ff9b::/96")
+
+
+class Local(OSError):
+    pass
+
+
+def unwrapped(address: ipaddress.IPv6Address) -> list[Address]:
+    inner: list[Address | None] = [address.ipv4_mapped, address.sixtofour, *(address.teredo or ())]
+    if address in NAT64:
+        inner.append(ipaddress.IPv4Address(int(address) & 0xFFFFFFFF))
+    return [item for item in inner if item is not None]
+
+
+def embedded(address: Address) -> list[Address]:
+    if isinstance(address, ipaddress.IPv4Address):
+        return [address]
+    return [address, *unwrapped(address)]
+
+
+def public(address: Address) -> bool:
+    return all(item.is_global and not item.is_multicast for item in embedded(address))
+
+
+def literal(host: str) -> Address | None:
+    try:
+        return ipaddress.ip_address(host.strip("[]"))
+    except ValueError:
+        return None
+
+
+class Guard(AbstractResolver):
+    def __init__(self, inner: AbstractResolver | None = None) -> None:
+        self.inner = inner or DefaultResolver()
+
+    async def resolve(
+        self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET
+    ) -> list[ResolveResult]:
+        results = await self.inner.resolve(host, port, family)
+        allowed = [result for result in results if public(ipaddress.ip_address(result["host"]))]
+        if not allowed:
+            raise Local(f"{host} resolves only to local addresses")
+        return allowed
+
+    async def close(self) -> None:
+        await self.inner.close()

egress_proxy/server/application.py

@@ -0,0 +1,53 @@
+from egress_proxy.server.asgi import Receive, Scope, Send
+from egress_proxy.server.forward import Forward
+from egress_proxy.server.routes import balancer
+from egress_proxy.server.settings import Proxy
+from egress_proxy.server.upstreams import configured
+
+
+class Unstarted(RuntimeError):
+    pass
+
+
+class Application:
+    def __init__(self, proxy: Proxy) -> None:
+        self.proxy = proxy
+        self.upstreams = configured(proxy)
+        self.forward: Forward | None = None
+
+    async def startup(self) -> None:
+        self.forward = Forward(balancer(self.proxy, self.upstreams), self.proxy)
+
+    async def shutdown(self) -> None:
+        if self.forward is not None:
+            await self.forward.balancer.close()
+            self.forward = None
+
+    async def lifespan(self, receive: Receive, send: Send) -> None:
+        while True:
+            message = await receive()
+            if message["type"] == "lifespan.startup":
+                await self.startup()
+                await send({"type": "lifespan.startup.complete"})
+                continue
+            await self.shutdown()
+            await send({"type": "lifespan.shutdown.complete"})
+            return
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "lifespan":
+            await self.lifespan(receive, send)
+            return
+        if scope["type"] != "http":
+            return
+        if self.forward is None:
+            raise Unstarted("application must complete lifespan startup before serving requests")
+        await self.forward.handle(scope, receive, send)
+
+
+def application(proxy: Proxy) -> Application:
+    return Application(proxy)
+
+
+def create() -> Application:
+    return application(Proxy())

egress_proxy/server/asgi.py

@@ -0,0 +1,7 @@
+from collections.abc import Awaitable, Callable, MutableMapping
+from typing import Any, TypeAlias
+
+Scope: TypeAlias = MutableMapping[str, Any]
+Message: TypeAlias = MutableMapping[str, Any]
+Receive: TypeAlias = Callable[[], Awaitable[Message]]
+Send: TypeAlias = Callable[[Message], Awaitable[None]]

egress_proxy/server/forward.py

@@ -0,0 +1,159 @@
+import hmac
+from collections.abc import AsyncIterator
+from typing import TypeAlias
+
+import aiohttp
+from aiohttp_socks import ProxyConnectionError, ProxyError, ProxyTimeoutError
+from yarl import URL
+
+from egress_proxy.protocol import Failure, Header
+from egress_proxy.server.addresses import Local, literal, public
+from egress_proxy.server.asgi import Receive, Scope, Send
+from egress_proxy.server.headers import REQUEST, RESPONSE, Pairs, filtered
+from egress_proxy.server.rejection import Rejection
+from egress_proxy.server.routes import Balancer, Route
+from egress_proxy.server.settings import Proxy
+
+Headers: TypeAlias = dict[bytes, bytes]
+
+SCHEMES = frozenset({"http", "https"})
+
+METHODS = frozenset({"GET", "HEAD", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"})
+
+KEY = Header.KEY.encode()
+
+TARGET = Header.TARGET.encode()
+
+UNREACHABLE = (ProxyConnectionError, ProxyTimeoutError)
+
+REFUSED = (*UNREACHABLE, ProxyError)
+
+
+def decoded(pairs: Pairs) -> list[tuple[str, str]]:
+    return [(name.decode("latin-1"), value.decode("latin-1")) for name, value in pairs]
+
+
+async def chunks(receive: Receive) -> AsyncIterator[bytes]:
+    while True:
+        message = await receive()
+        if message["type"] == "http.disconnect":
+            return
+        if message.get("body"):
+            yield message["body"]
+        if not message.get("more_body", False):
+            return
+
+
+class Stream:
+    def __init__(self, receive: Receive, present: bool) -> None:
+        self.receive = receive
+        self.present = present
+        self.started = False
+
+    async def chunks(self) -> AsyncIterator[bytes]:
+        self.started = True
+        async for chunk in chunks(self.receive):
+            yield chunk
+
+    def body(self) -> AsyncIterator[bytes] | None:
+        return self.chunks() if self.present else None
+
+
+class Forward:
+    def __init__(self, balancer: Balancer, proxy: Proxy) -> None:
+        self.balancer = balancer
+        self.key = proxy.key.get_secret_value().encode()
+        self.local = proxy.allow_local
+
+    def authorize(self, headers: Headers) -> None:
+        if not hmac.compare_digest(headers.get(KEY, b""), self.key):
+            raise Rejection(Failure.UNAUTHORIZED, "egress key is missing or invalid")
+
+    def method(self, scope: Scope) -> str:
+        method: str = scope["method"]
+        if method not in METHODS:
+            raise Rejection(Failure.INVALID, f"method {method} is not supported")
+        return method
+
+    def target(self, headers: Headers) -> URL:
+        value = headers.get(TARGET)
+        if value is None:
+            raise Rejection(Failure.INVALID, f"'{Header.TARGET}' header is required")
+        text = value.decode("latin-1")
+        try:
+            url = URL(text, encoded=True)
+        except ValueError as error:
+            raise Rejection(Failure.INVALID, f"'{text}' is not a valid url") from error
+        if url.scheme not in SCHEMES or not url.host:
+            raise Rejection(Failure.INVALID, f"'{text}' must be an absolute http or https url")
+        self.reachable(url)
+        return url
+
+    def reachable(self, url: URL) -> None:
+        address = literal(url.host or "")
+        if not self.local and address is not None and not public(address):
+            raise Rejection(Failure.INVALID, f"{url} points to a local address")
+
+    def unreachable(self, url: URL, error: aiohttp.ClientError) -> Rejection:
+        if isinstance(getattr(error, "os_error", None), Local):
+            return Rejection(Failure.INVALID, f"{url} resolves only to local addresses")
+        return Rejection(Failure.UNREACHABLE, f"{url} is unreachable")
+
+    def stream(self, headers: Headers, receive: Receive) -> Stream:
+        return Stream(receive, b"content-length" in headers or b"transfer-encoding" in headers)
+
+    async def send(
+        self,
+        route: Route,
+        method: str,
+        url: URL,
+        headers: list[tuple[str, str]],
+        stream: Stream,
+    ) -> aiohttp.ClientResponse:
+        try:
+            return await route.client.request(method, url, headers=headers, data=stream.body(), allow_redirects=False)
+        except TimeoutError as error:
+            raise Rejection(Failure.TIMEOUT, f"{url} timed out") from error
+        except aiohttp.ClientError as error:
+            raise self.unreachable(url, error) from error
+
+    def failed(self, route: Route, url: URL, error: Exception) -> Rejection:
+        if isinstance(error, UNREACHABLE):
+            self.balancer.rest(route)
+            return Rejection(Failure.UNREACHABLE, f"upstream proxy {route.name} is unreachable")
+        return Rejection(Failure.UNREACHABLE, f"{url} is unreachable through upstream proxy {route.name}")
+
+    async def request(self, scope: Scope, receive: Receive) -> aiohttp.ClientResponse:
+        headers: Headers = dict(scope["headers"])
+        self.authorize(headers)
+        method = self.method(scope)
+        url = self.target(headers)
+        forwarded = decoded(filtered(scope["headers"], REQUEST))
+        stream = self.stream(headers, receive)
+        failure = Rejection(Failure.UNREACHABLE, f"{url} is unreachable")
+        for route in self.balancer.candidates():
+            try:
+                return await self.send(route, method, url, forwarded, stream)
+            except REFUSED as error:
+                failure = self.failed(route, url, error)
+            if stream.started:
+                break
+        raise failure
+
+    async def relay(self, upstream: aiohttp.ClientResponse, send: Send) -> None:
+        headers = filtered(list(upstream.raw_headers), RESPONSE)
+        await send({"type": "http.response.start", "status": upstream.status, "headers": headers})
+        async for chunk in upstream.content.iter_any():
+            await send({"type": "http.response.body", "body": chunk, "more_body": True})
+        await send({"type": "http.response.body", "body": b""})
+
+    async def handle(self, scope: Scope, receive: Receive, send: Send) -> None:
+        try:
+            upstream = await self.request(scope, receive)
+        except Rejection as rejection:
+            await rejection.send(send)
+            return
+        try:
+            await self.relay(upstream, send)
+        finally:
+            upstream.release()

egress_proxy/server/headers.py

@@ -0,0 +1,133 @@
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TypeAlias
+
+from egress_proxy.protocol import Header
+
+Pairs: TypeAlias = list[tuple[bytes, bytes]]
+
+HOP = frozenset(
+    {
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "proxy-connection",
+        "te",
+        "trailer",
+        "transfer-encoding",
+        "upgrade",
+        "xonnection",
+        "xroxy-connection",
+    }
+)
+
+UNSAFE = frozenset({"proxy"})
+
+ADDRESSES = (
+    "forwarded",
+    "forward-for",
+    "client-ip",
+    "clientip",
+    "real-ip",
+    "remote-addr",
+    "remote-ip",
+    "connecting-ip",
+    "user-ip",
+    "originating-ip",
+    "original-ip",
+    "socketip",
+)
+
+ADDRESSED = frozenset(
+    {
+        "x-ip",
+        "x-from-ip",
+        "source-ip",
+        "coming-from",
+        "x-coming-from",
+        "connect-via-ip",
+        "x-delegate-remote-host",
+        "x-imforwards",
+        "x-proxymesh-ip",
+        "cf-pseudo-ipv4",
+    }
+)
+
+MARKERS = frozenset(
+    {
+        "via",
+        "x-bluecoat-via",
+        "x-proxy-id",
+        "cdn-loop",
+        "cf-ew-via",
+        "cf-ipcountry",
+        "cf-ray",
+        "cf-visitor",
+        "fastly-ff",
+        "akamai-origin-hop",
+        "x-akamai-edgescape",
+        "x-azure-fdid",
+        "x-azure-ref",
+        "x-amzn-trace-id",
+        "x-host",
+        "x-original-host",
+        "x-backend-host",
+        "x-gateway-host",
+        "x-url-scheme",
+        "front-end-https",
+    }
+)
+
+EDGES = ("x-envoy-", "cloudfront-viewer-")
+
+PRIVATE = frozenset(header.value for header in Header)
+
+
+@dataclass(frozen=True)
+class Policy:
+    exact: frozenset[str]
+    fragments: tuple[str, ...] = ()
+    prefixes: tuple[str, ...] = ()
+
+    def excluded(self, name: str) -> bool:
+        if name in self.exact or name.startswith(self.prefixes):
+            return True
+        return any(fragment in name for fragment in self.fragments)
+
+
+REQUEST = Policy(
+    exact=HOP | UNSAFE | ADDRESSED | MARKERS | PRIVATE | {"host"},
+    fragments=ADDRESSES,
+    prefixes=EDGES,
+)
+
+RESPONSE = Policy(exact=HOP | PRIVATE)
+
+
+def canonical(text: str) -> str:
+    return text.strip().lower().replace("_", "-")
+
+
+@lru_cache(maxsize=4096)
+def normalized(name: bytes) -> str:
+    return canonical(name.decode("latin-1"))
+
+
+@lru_cache(maxsize=4096)
+def blocked(policy: Policy, name: bytes) -> bool:
+    return policy.excluded(normalized(name))
+
+
+def declared(pairs: Pairs) -> frozenset[str]:
+    values = (value.decode("latin-1") for name, value in pairs if normalized(name) == "connection")
+    return frozenset(canonical(token) for value in values for token in value.split(","))
+
+
+def kept(name: bytes, policy: Policy, hop: frozenset[str]) -> bool:
+    return not blocked(policy, name) and normalized(name) not in hop
+
+
+def filtered(pairs: Pairs, policy: Policy) -> Pairs:
+    hop = declared(pairs)
+    return [(name, value) for name, value in pairs if kept(name, policy, hop)]