Investment Plans workspace
Open raw ↗
"""AI-engine adapter interface + the outbound request guard (GOV-B7.3 / PIL-WA-3).

Every adapter implements: name, chat(messages) -> reply text. The engine endpoint and
API key are read from ENVIRONMENT VARIABLES ONLY — never from a file the server sends
to a browser, never hardcoded. guarded_request() is the ONLY way an adapter talks to
the network: http/https only, host resolved once and every resolved IP checked against
the forbidden set (loopback/private/link-local/reserved/multicast), the connection then
PINNED to the validated IP (anti DNS-rebinding: DNS is consulted once, not twice), TLS
still verifies the hostname, and redirects are refused (a redirect can bounce inward).
"""
import http.client
import ipaddress
import os
import socket
import ssl
import urllib.parse


class GuardError(ValueError):
    """An outbound URL failed the safety guard."""


_FORBIDDEN = lambda ip: (ip.is_loopback or ip.is_private or ip.is_link_local
                         or ip.is_reserved or ip.is_multicast or ip.is_unspecified)


def guard_url(url: str) -> tuple[str, str, int]:
    """Validate an outbound URL -> (hostname, validated_ip, port). Raises GuardError."""
    parts = urllib.parse.urlsplit(url)
    if parts.scheme not in ("http", "https"):
        raise GuardError(f"scheme '{parts.scheme}' not allowed (http/https only)")
    host = parts.hostname or ""
    if not host:
        raise GuardError("no host in URL")
    port = parts.port or (443 if parts.scheme == "https" else 80)
    try:
        infos = socket.getaddrinfo(host, port, proto=socket.IPPROTO_TCP)
    except socket.gaierror as exc:
        raise GuardError(f"host does not resolve: {host}") from exc
    for info in infos:                      # EVERY resolved address must be public
        ip = ipaddress.ip_address(info[4][0])
        if _FORBIDDEN(ip):
            raise GuardError(f"host resolves to a forbidden address: {ip} "
                             "(localhost/loopback/private/reserved refused)")
    # resolve once, pin that IP for the connection (DNS is not consulted again -> no rebinding)
    return host, infos[0][4][0], port


def guarded_post_json(url: str, headers: dict, payload: bytes, timeout: int = 60) -> dict:
    """POST bytes to a guarded URL, pin the validated IP, refuse redirects, return parsed JSON."""
    hostname, ip, port = guard_url(url)
    scheme = urllib.parse.urlsplit(url).scheme
    ctx = ssl.create_default_context()
    if scheme == "https":
        conn = http.client.HTTPSConnection(hostname, port, timeout=timeout, context=ctx)
        conn._conn_address = ip            # consumed by PinnedConnection below
    else:
        conn = http.client.HTTPConnection(hostname, port, timeout=timeout)
        conn._conn_address = ip
    # pin: create_connection to the validated IP, TLS SNI/verification still by hostname
    _pin(conn, ip)
    try:
        conn.request("POST", urllib.parse.urlsplit(url).path or "/", body=payload, headers=headers)
        resp = conn.getresponse()
        if 300 <= resp.status < 400:
            raise GuardError(f"redirect to {resp.getheader('Location')!r} refused (a redirect can bounce inward)")
        if resp.status != 200:
            raise GuardError(f"engine returned HTTP {resp.status}: {resp.read(200).decode('utf-8', 'replace')}")
        return json_loads(resp.read().decode("utf-8"))
    finally:
        conn.close()


def _pin(conn, ip: str) -> None:
    orig_connect = conn.connect

    def connect():
        sock = socket.create_connection((ip, conn.port), timeout=conn.timeout)
        if isinstance(conn, http.client.HTTPSConnection):
            sock = conn._context.wrap_socket(sock, server_hostname=conn.host)  # SNI + cert by hostname
        conn.sock = sock

    conn.connect = connect


def json_loads(text: str) -> dict:
    import json
    return json.loads(text)


class BaseAdapter:
    """One AI engine behind one interface. Subclasses set NAME and implement chat()."""

    NAME = "base"

    def __init__(self, env: dict | None = None):
        self.env = env if env is not None else os.environ

    @property
    def api_key(self) -> str:
        raise NotImplementedError

    def chat(self, messages: list[dict]) -> str:
        raise NotImplementedError