DataOps-platform/app/core/edge_gateway/policy.py

"""Default-deny egress and retention policy for enterprise edge events."""

from __future__ import annotations

import ipaddress
import re
import unicodedata
from collections.abc import Mapping
from urllib.parse import urlsplit

from .contracts import EdgeContractError, preflight_json, strict_json_bytes


class EdgePolicyError(ValueError):
    """Raised when content or transport violates the edge boundary."""


EDGE_ONLY = {"raw", "recent_detail", "restricted"}
CONTROL_PLANE_ALLOWED = {
    "desensitized_metadata",
    "statistics",
    "lineage",
    "evidence",
    "health_summary",
    "diagnostic_summary",
}
RETENTION_DAYS = {
    "raw": 0,
    "recent_detail": 365,
    "metadata": 1095,
    "evidence": 2190,
}

DEFAULT_BYTE_LIMITS = {
    "desensitized_metadata": 65_536,
    "statistics": 32_768,
    "lineage": 65_536,
    "evidence": 131_072,
    "health_summary": 8_192,
    "diagnostic_summary": 16_384,
}
DEFAULT_COUNT_LIMITS = {
    "desensitized_metadata": 2_000,
    "statistics": 1_000,
    "lineage": 2_000,
    "evidence": 4_000,
    "health_summary": 256,
    "diagnostic_summary": 512,
}

_SENSITIVE_KEYS = {
    "api_key",
    "apikey",
    "authorization",
    "access_token",
    "cookie",
    "credential",
    "credentials",
    "password",
    "passwd",
    "private_key",
    "query",
    "query_text",
    "raw_record",
    "raw_records",
    "raw_row",
    "raw_rows",
    "secret",
    "sql",
    "sql_text",
    "statement",
    "token",
}
_SENSITIVE_KEY_PARTS = ("password", "passwd", "secret", "private_key", "api_key")
_RAW_ROW_CONTAINER_KEYS = {
    "data",
    "items",
    "item",
    "record",
    "records",
    "result",
    "row",
    "rows",
    "sample",
    "samples",
}
_PRIVATE_KEY = re.compile(r"-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----", re.IGNORECASE)
_BEARER = re.compile(r"\bBearer\s+[A-Za-z0-9._~+/=-]{8,}", re.IGNORECASE)
_CREDENTIAL_ASSIGNMENT = re.compile(
    r"\b(?:[a-z0-9]+[_-])*(?:password|passwd|pwd|secret|token|api[_-]?key|"
    r"authorization|credential)"
    r"\b\s*(?:=|:)\s*[^\s,;]+",
    re.IGNORECASE,
)
_SQL = re.compile(
    r"(?:"
    r"\bSELECT\b.{0,4096}?\bFROM\b"
    r"|\bINSERT\s+INTO\b"
    r"|\bUPDATE\b.{0,4096}?\bSET\b"
    r"|\bDELETE\s+FROM\b"
    r"|\bMERGE(?:\s+INTO)?\b"
    r"|\b(?:CREATE|ALTER|DROP|TRUNCATE)\s+"
    r"(?:TABLE|VIEW|SCHEMA|DATABASE|INDEX|FUNCTION|PROCEDURE)\b"
    r"|\b(?:GRANT|REVOKE)\b.{0,4096}?\b(?:TO|FROM)\b"
    r"|\bCOPY\b"
    r")",
    re.IGNORECASE | re.DOTALL,
)


def _normalized_key(value: str) -> str:
    compatible = unicodedata.normalize("NFKC", value)
    words = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", compatible)
    return re.sub(r"[^a-z0-9]+", "_", words.lower()).strip("_")


def _is_sensitive_key(value: str) -> bool:
    folded = value.replace("_", "")
    return (
        value in _SENSITIVE_KEYS
        or folded
        in {
            "accesstoken",
            "apikey",
            "authorization",
            "clientsecret",
            "credential",
            "credentials",
            "password",
            "passwd",
            "privatekey",
            "rawrecords",
            "rawrows",
            "secret",
            "token",
        }
        or any(part in value for part in _SENSITIVE_KEY_PARTS)
        or value.startswith(("raw_", "recent_detail_", "token_"))
        or value.endswith("_token")
    )


def _is_raw_row_container(key: str, value: object) -> bool:
    return (
        key in _RAW_ROW_CONTAINER_KEYS
        or key.endswith(("_rows", "_records"))
    ) and (
        isinstance(value, Mapping)
        or (
            isinstance(value, (list, tuple))
            and any(isinstance(child, Mapping) for child in value)
        )
    )


def _contains_obfuscated_sql(value: str) -> bool:
    compatible = unicodedata.normalize("NFKC", value)
    without_comments = re.sub(
        r"/\*.*?\*/|--[^\r\n]*", "", compatible, flags=re.DOTALL
    )
    normalized = re.sub(r"\s+", " ", without_comments).upper()
    return bool(
        re.search(
            r"(?:\bSELECT\b.{0,4096}\bFROM\b|\bINSERT\s+INTO\b|"
            r"\bUPDATE\b.{0,4096}\bSET\b|\bDELETE\s+FROM\b|"
            r"\bMERGE(?:\s+INTO)?\b|\b(?:CREATE|ALTER|DROP|TRUNCATE)\s+"
            r"(?:TABLE|VIEW|SCHEMA|DATABASE|INDEX|FUNCTION|PROCEDURE)\b|"
            r"\b(?:GRANT|REVOKE)\b.{0,4096}\b(?:TO|FROM)\b|\bCOPY\b|"
            r"\bEXEC(?:UTE)?\b\s+[A-Z_\[][A-Z0-9_.$\[\]]*|"
            r"\bCALL\b\s+[A-Z_][A-Z0-9_.$]*\s*\()",
            normalized,
        )
    )


def _approved_hosts(values: set[str] | frozenset[str], label: str) -> frozenset[str]:
    if not isinstance(values, (set, frozenset)):
        raise ValueError(f"{label} must be an explicit set")
    approved: set[str] = set()
    for value in values:
        if (
            not isinstance(value, str)
            or not value
            or value != value.lower()
            or value.endswith(".")
            or ":" in value
            or "/" in value
            or "@" in value
        ):
            raise ValueError(f"{label} contains an invalid host")
        try:
            value.encode("ascii")
            ipaddress.ip_address(value)
        except UnicodeEncodeError as exc:
            raise ValueError(f"{label} contains an invalid host") from exc
        except ValueError:
            pass
        else:
            raise ValueError(f"{label} cannot contain an IP literal")
        if not re.fullmatch(
            r"(?=.{1,253}\Z)(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)*"
            r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?",
            value,
        ):
            raise ValueError(f"{label} contains an invalid host")
        approved.add(value)
    return frozenset(approved)


def _approved_origins(
    values: set[str] | frozenset[str] | None,
    hosts: frozenset[str],
    label: str,
) -> frozenset[str]:
    if values is None:
        return frozenset(f"https://{host}" for host in hosts)
    if not isinstance(values, (set, frozenset)):
        raise ValueError(f"{label} must be an explicit set")
    origins: set[str] = set()
    for value in values:
        if not isinstance(value, str):
            raise ValueError(f"{label} contains an invalid origin")
        try:
            parsed = urlsplit(value)
            port = parsed.port
        except ValueError as exc:
            raise ValueError(f"{label} contains an invalid origin") from exc
        if (
            parsed.scheme != "https"
            or parsed.hostname not in hosts
            or parsed.username is not None
            or parsed.password is not None
            or parsed.path not in {"", "/"}
            or parsed.query
            or parsed.fragment
            or port is None and parsed.netloc.endswith(":")
        ):
            raise ValueError(f"{label} contains an invalid origin")
        effective_port = port or 443
        canonical = (
            f"https://{parsed.hostname}"
            if effective_port == 443
            else f"https://{parsed.hostname}:{effective_port}"
        )
        if value.rstrip("/") != canonical:
            raise ValueError(f"{label} contains a non-canonical origin")
        origins.add(canonical)
    return frozenset(origins)


def _limit_map(
    defaults: Mapping[str, int], overrides: Mapping[str, int] | None, label: str
) -> dict[str, int]:
    result = dict(defaults)
    if overrides is None:
        return result
    if not isinstance(overrides, Mapping):
        raise ValueError(f"{label} must be a mapping")
    unknown = set(overrides) - CONTROL_PLANE_ALLOWED
    if unknown:
        raise ValueError(f"{label} contains an unknown classification")
    for classification, limit in overrides.items():
        if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1:
            raise ValueError(f"{label} values must be positive integers")
        result[classification] = limit
    return result


class EdgeEgressPolicy:
    def __init__(
        self,
        *,
        allowed_control_hosts: set[str] | frozenset[str],
        allowed_proxy_hosts: set[str] | frozenset[str] | None = None,
        allowed_control_origins: set[str] | frozenset[str] | None = None,
        allowed_proxy_origins: set[str] | frozenset[str] | None = None,
        byte_limits: Mapping[str, int] | None = None,
        count_limits: Mapping[str, int] | None = None,
        max_depth: int = 12,
        max_string_bytes: int = 4_096,
    ) -> None:
        self.allowed_control_hosts = _approved_hosts(
            allowed_control_hosts, "allowed_control_hosts"
        )
        self.allowed_proxy_hosts = _approved_hosts(
            allowed_proxy_hosts or set(), "allowed_proxy_hosts"
        )
        self.allowed_control_origins = _approved_origins(
            allowed_control_origins,
            self.allowed_control_hosts,
            "allowed_control_origins",
        )
        self.allowed_proxy_origins = _approved_origins(
            allowed_proxy_origins,
            self.allowed_proxy_hosts,
            "allowed_proxy_origins",
        )
        self.byte_limits = _limit_map(
            DEFAULT_BYTE_LIMITS, byte_limits, "byte_limits"
        )
        self.count_limits = _limit_map(
            DEFAULT_COUNT_LIMITS, count_limits, "count_limits"
        )
        if isinstance(max_depth, bool) or not isinstance(max_depth, int) or max_depth < 1:
            raise ValueError("max_depth must be a positive integer")
        if (
            isinstance(max_string_bytes, bool)
            or not isinstance(max_string_bytes, int)
            or max_string_bytes < 1
        ):
            raise ValueError("max_string_bytes must be a positive integer")
        self.max_depth = max_depth
        self.max_string_bytes = max_string_bytes

    def approve_event(self, event: Mapping[str, object]) -> dict[str, object]:
        if not isinstance(event, Mapping):
            raise EdgePolicyError("event must be a mapping")
        unknown = set(event) - {
            "event_id",
            "task_id",
            "gateway_id",
            "environment",
            "network_zone",
            "purpose",
            "classification",
            "contract_version",
            "occurred_at",
            "attempt",
            "idempotency_key",
            "policy_digest",
            "payload",
        }
        if unknown:
            raise EdgePolicyError("event contains unapproved properties")
        classification = event.get("classification")
        if classification in EDGE_ONLY:
            raise EdgePolicyError("classification is edge-only")
        if classification not in CONTROL_PLANE_ALLOWED:
            raise EdgePolicyError("classification is not control-plane approved")
        payload = event.get("payload")
        if not isinstance(payload, Mapping):
            raise EdgePolicyError("approved event payload must be a mapping")
        self.validate_approved_payload(str(classification), payload)
        try:
            return dict(event)
        except (TypeError, ValueError) as exc:  # pragma: no cover - Mapping guard
            raise EdgePolicyError("event is invalid") from exc

    def validate_approved_payload(
        self, classification: str, payload: Mapping[str, object]
    ) -> None:
        if classification not in CONTROL_PLANE_ALLOWED:
            raise EdgePolicyError("classification is not control-plane approved")
        try:
            preflight_json(
                payload,
                max_depth=self.max_depth,
                max_nodes=self.count_limits[classification] + 1,
                max_string_bytes=self.max_string_bytes,
                max_total_bytes=self.byte_limits[classification],
            )
            encoded = strict_json_bytes(payload)
        except EdgeContractError as exc:
            detail = str(exc)
            if "node count" in detail:
                detail = "payload exceeds its independent item count limit"
            elif "string exceeds" in detail:
                detail = "payload string exceeds the string limit"
            elif "depth" in detail:
                detail = "payload exceeds the maximum depth"
            elif "byte limit" in detail:
                detail = "payload exceeds its independent byte limit"
            raise EdgePolicyError(detail) from exc
        if len(encoded) > self.byte_limits[classification]:
            raise EdgePolicyError("payload exceeds its independent byte limit")
        item_count = self._count_items(payload)
        if item_count > self.count_limits[classification]:
            raise EdgePolicyError("payload exceeds its independent item count limit")
        self._inspect(payload, depth=1)

    def _count_items(self, value: object) -> int:
        if isinstance(value, Mapping):
            return len(value) + sum(self._count_items(item) for item in value.values())
        if isinstance(value, (list, tuple)):
            return len(value) + sum(self._count_items(item) for item in value)
        return 0

    def _inspect(self, value: object, *, depth: int) -> None:
        if depth > self.max_depth:
            raise EdgePolicyError("payload exceeds the maximum depth")
        if isinstance(value, Mapping):
            for key, item in value.items():
                if not isinstance(key, str):
                    raise EdgePolicyError("payload contains a non-string key")
                normalized = _normalized_key(key)
                if _is_sensitive_key(normalized):
                    raise EdgePolicyError("payload contains sensitive content")
                if _is_raw_row_container(normalized, item):
                    raise EdgePolicyError("payload contains sensitive raw-row content")
                self._inspect(item, depth=depth + 1)
            return
        if isinstance(value, (list, tuple)):
            for item in value:
                self._inspect(item, depth=depth + 1)
            return
        if isinstance(value, str):
            compatible = unicodedata.normalize("NFKC", value)
            if len(value.encode("utf-8")) > self.max_string_bytes:
                raise EdgePolicyError("payload string exceeds the string limit")
            if (
                _PRIVATE_KEY.search(compatible)
                or _BEARER.search(compatible)
                or _CREDENTIAL_ASSIGNMENT.search(compatible)
                or _SQL.search(compatible)
                or _contains_obfuscated_sql(compatible)
            ):
                raise EdgePolicyError("payload contains sensitive content")

    def validate_destination(
        self, control_url: str, *, proxy_url: str | None = None
    ) -> None:
        self._validate_url(
            control_url,
            approved_hosts=self.allowed_control_hosts,
            approved_origins=self.allowed_control_origins,
            label="destination",
        )
        if proxy_url is not None:
            self._validate_url(
                proxy_url,
                approved_hosts=self.allowed_proxy_hosts,
                approved_origins=self.allowed_proxy_origins,
                label="proxy",
            )

    @staticmethod
    def _validate_url(
        url: str,
        *,
        approved_hosts: frozenset[str],
        approved_origins: frozenset[str],
        label: str,
    ) -> None:
        if not isinstance(url, str):
            raise EdgePolicyError(f"{label} is not approved")
        try:
            parsed = urlsplit(url)
            port = parsed.port
        except ValueError as exc:
            raise EdgePolicyError(f"{label} is not approved") from exc
        if (
            parsed.scheme != "https"
            or not parsed.hostname
            or parsed.hostname not in approved_hosts
            or parsed.username is not None
            or parsed.password is not None
            or parsed.fragment
            or port is None and parsed.netloc.endswith(":")
        ):
            raise EdgePolicyError(f"{label} is not approved")
        effective_port = port or 443
        origin = (
            f"https://{parsed.hostname}"
            if effective_port == 443
            else f"https://{parsed.hostname}:{effective_port}"
        )
        if origin not in approved_origins:
            raise EdgePolicyError(f"{label} origin is not approved")

    @staticmethod
    def retention_days(classification: str) -> int:
        try:
            return RETENTION_DAYS[classification]
        except (KeyError, TypeError) as exc:
            raise EdgePolicyError("retention classification is invalid") from exc