DataOps-platform/tests/security/test_phase3_wp04_data_boundary.py

"""Security boundary tests for the P3-WP04 pull-only edge runtime."""

from __future__ import annotations

import json
from datetime import UTC, datetime, timedelta
from pathlib import Path
from urllib.parse import quote

import pytest
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import NameOID
from flask import Flask

from app.api.data_source.edge_routes import _tls_certificate_sha256
from app.core.edge_gateway.contracts import EdgeContractError, EdgeTaskContract
from app.core.edge_gateway.policy import EdgePolicyError
from app.core.edge_gateway.service import EdgeGatewayAuthenticationError
from app.edge_gateway.agent import EdgeRunnerAdapter
from app.edge_gateway.bootstrap import EdgeBootstrapConfig
from app.edge_gateway.transport import (
    EdgeAuthenticationStopped,
    EdgeTransport,
    EdgeTransportError,
)
from app.runner.api import execute_edge_adapter
from app.runner.nodes import NodeExecutionError

POLICY_DIGEST = "a" * 64


def _escaped_client_certificate():
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "edge-gateway-test")])
    now = datetime.now(UTC)
    certificate = (
        x509.CertificateBuilder()
        .subject_name(name)
        .issuer_name(name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(minutes=1))
        .not_valid_after(now + timedelta(minutes=5))
        .sign(key, hashes.SHA256())
    )
    pem = certificate.public_bytes(serialization.Encoding.PEM).decode()
    return quote(pem, safe=""), certificate.fingerprint(hashes.SHA256()).hex()


def test_backend_accepts_certificate_only_from_explicit_trusted_mtls_proxy():
    app = Flask(__name__)
    app.config["EDGE_MTLS_TRUSTED_PROXY_IPS"] = ("172.31.0.10",)
    escaped, fingerprint = _escaped_client_certificate()
    headers = {
        "X-DataOps-Edge-Client-Cert": escaped,
        "X-DataOps-Edge-Client-Verify": "SUCCESS",
        "X-Edge-Certificate-SHA256": fingerprint,
    }
    with app.test_request_context(
        "/", headers=headers, environ_overrides={"REMOTE_ADDR": "172.31.0.10"}
    ):
        assert _tls_certificate_sha256() == fingerprint
    with app.test_request_context(
        "/", headers=headers, environ_overrides={"REMOTE_ADDR": "203.0.113.7"}
    ), pytest.raises(EdgeGatewayAuthenticationError, match="trusted mTLS proxy"):
        _tls_certificate_sha256()
    with app.test_request_context(
        "/",
        headers={**headers, "X-DataOps-Edge-Client-Cert": "forged"},
        environ_overrides={"REMOTE_ADDR": "172.31.0.10"},
    ), pytest.raises(EdgeGatewayAuthenticationError, match="certificate"):
        _tls_certificate_sha256()


def task(**changes):
    value = {
        "task_id": "task-1",
        "gateway_id": "gateway-1",
        "environment": "production",
        "network_zone": "zone-a",
        "purpose": "governed-inventory",
        "classification": "statistics",
        "task_type": "profile",
        "contract_version": 1,
        "deadline_at": (datetime.now(UTC) + timedelta(minutes=5)).isoformat().replace("+00:00", "Z"),
        "attempt": 1,
        "idempotency_key": "idem-1",
        "policy_digest": POLICY_DIGEST,
    }
    value.update(changes)
    return value


class NeverNetwork:
    def request(self, *args, **kwargs):  # pragma: no cover - must not run
        raise AssertionError("network request must not be attempted")


@pytest.mark.parametrize(
    "url,proxy",
    [
        ("http://control.enterprise.test", None),
        ("https://user@control.enterprise.test", None),
        ("https://control.enterprise.test.evil.test", None),
        ("https://control.enterprise.test", "https://evil.test"),
    ],
)
def test_transport_rejects_unapproved_destinations_before_network(url, proxy):
    with pytest.raises(EdgePolicyError):
        EdgeTransport(
            base_url=url,
            proxy_url=proxy,
            allowed_control_hosts={"control.enterprise.test"},
            allowed_proxy_hosts={"proxy.enterprise.test"},
            gateway_id="gateway-1",
            environment="production",
            network_zone="zone-a",
            generation=1,
            credential="dopg_example-credential",
            certificate_sha256="b" * 64,
            client=NeverNetwork(),
            client_certificate_path="/secure/edge-cert.pem",
            client_private_key_path="/secure/edge-key.pem",
            ca_bundle_path="/secure/ca.pem",
        )


@pytest.mark.parametrize(
    "operation",
    ["shell", "python", "sql", "raw_export", "http", "../../bin/sh"],
)
def test_runner_adapter_rejects_unknown_or_malicious_task_types(operation):
    adapter = EdgeRunnerAdapter(lambda _node, _request, _cancel: {})
    value = task(task_type=operation)
    with pytest.raises((EdgeContractError, EdgePolicyError)):
        adapter.execute(EdgeTaskContract.from_mapping(value), lambda: False)


@pytest.mark.parametrize(
    "result",
    [
        {"classification": "raw", "payload": {"count": 1}},
        {"classification": "statistics", "payload": {"sql": "SELECT * FROM secret"}},
        {"classification": "statistics", "payload": {"password": "hidden"}},
        {"classification": "statistics", "payload": {"rows": [{"id": 1}]}},
        {"classification": "statistics", "payload": {"count": 1}, "command": "sh"},
    ],
)
def test_runner_adapter_rejects_raw_secret_script_and_unknown_result_fields(result):
    adapter = EdgeRunnerAdapter(lambda _node, _request, _cancel: result)
    with pytest.raises(EdgePolicyError):
        adapter.execute(EdgeTaskContract.from_mapping(task()), lambda: False)


def test_runner_adapter_rejects_recursive_and_cyclic_results():
    cyclic = {}
    cyclic["child"] = cyclic
    adapter = EdgeRunnerAdapter(
        lambda _node, _request, _cancel: {
            "classification": "statistics",
            "payload": cyclic,
        }
    )
    with pytest.raises(EdgePolicyError):
        adapter.execute(EdgeTaskContract.from_mapping(task()), lambda: False)


def test_runner_adapter_rejects_independently_oversized_result():
    adapter = EdgeRunnerAdapter(
        lambda _node, _request, _cancel: {
            "classification": "statistics",
            "payload": {"summary": "x" * 33_000},
        }
    )
    with pytest.raises(EdgePolicyError):
        adapter.execute(EdgeTaskContract.from_mapping(task()), lambda: False)


def test_runner_adapter_rejects_unallowlisted_purpose_before_local_handler():
    called = False

    def execute(_node, _request, _cancel):
        nonlocal called
        called = True
        return {"classification": "statistics", "payload": {"metric_count": 1}}

    adapter = EdgeRunnerAdapter(execute)
    with pytest.raises(EdgePolicyError, match="purpose"):
        adapter.execute(
            EdgeTaskContract.from_mapping(task(purpose="arbitrary-action")),
            lambda: False,
        )
    assert not called


class Response:
    def __init__(self, status_code, content):
        self.status_code = status_code
        self.headers = {"Content-Length": str(len(content))}
        self._content = content

    def iter_content(self, chunk_size):
        yield self._content


class RecordingClient:
    def __init__(self, response):
        self.response = response
        self.calls = []

    def request(self, *args, **kwargs):
        self.calls.append((args, kwargs))
        return self.response


class SequencedClient:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def request(self, *args, **kwargs):
        self.calls.append((args, kwargs))
        return self.responses.pop(0)


def _transport(client):
    return EdgeTransport(
        base_url="https://control.enterprise.test",
        allowed_control_hosts={"control.enterprise.test"},
        gateway_id="gateway-1",
        environment="production",
        network_zone="zone-a",
        generation=1,
        credential="dopg_example-credential",
        certificate_sha256="b" * 64,
        client=client,
        client_certificate_path="/secure/edge-cert.pem",
        client_private_key_path="/secure/edge-key.pem",
        ca_bundle_path="/secure/ca.pem",
    )


def test_transport_uses_exact_bound_headers_and_forbids_redirect_following():
    client = RecordingClient(Response(200, b'{"code":200,"message":"ok","data":{"task":null}}'))
    assert _transport(client).pull_task() is None
    args, kwargs = client.calls[0]
    assert args == (
        "POST",
        "https://control.enterprise.test/api/datasource/edge/gateways/gateway-1/tasks/pull",
    )
    assert kwargs["allow_redirects"] is False
    assert kwargs["headers"] == {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "X-Edge-Credential": "dopg_example-credential",
        "X-Edge-Certificate-SHA256": "b" * 64,
    }


@pytest.mark.parametrize("status", [301, 302, 307, 308])
def test_transport_rejects_all_redirects(status):
    with pytest.raises(EdgeTransportError, match="redirects"):
        _transport(RecordingClient(Response(status, b""))).pull_task()


@pytest.mark.parametrize("status", [401, 403])
def test_transport_stops_on_revoked_or_expired_binding(status):
    with pytest.raises(EdgeAuthenticationStopped):
        _transport(RecordingClient(Response(status, b'{}'))).reconcile()


def test_transport_rejects_oversized_response_before_json_decode():
    with pytest.raises(EdgeTransportError, match="byte limit"):
        _transport(
            RecordingClient(Response(200, b"x" * (EdgeTransport.MAX_RESPONSE_BYTES + 1)))
        ).reconcile()


def test_transport_cancel_probe_uses_bounded_reconcile_contract():
    client = RecordingClient(
            Response(
                200,
                b'{"code":200,"message":"ok","data":{"cancelled_task_ids":["task-1"],"cancel_next_cursor":null,"release_offers":[],"release_next_cursor":null,"release_baseline":null}}',
            )
    )
    assert _transport(client).cancel_requested("task-1") is True


def test_transport_reconcile_drains_all_stable_cursor_pages():
    first_ids = [f"task-{index:02d}" for index in range(50)]
    first = {
        "code": 200,
        "message": "ok",
        "data": {
            "cancelled_task_ids": first_ids,
            "cancel_next_cursor": first_ids[-1],
            "release_offers": [],
            "release_next_cursor": None,
            "release_baseline": None,
        },
    }
    second = {
        "code": 200,
        "message": "ok",
        "data": {
            "cancelled_task_ids": ["task-50"],
            "cancel_next_cursor": None,
            "release_offers": [],
            "release_next_cursor": None,
            "release_baseline": None,
        },
    }
    client = SequencedClient(
        [
            Response(200, json.dumps(first, separators=(",", ":")).encode()),
            Response(200, json.dumps(second, separators=(",", ":")).encode()),
        ]
    )
    reconciled = _transport(client).reconcile()
    assert reconciled == {
        "cancelled_task_ids": [*first_ids, "task-50"],
        "release_offers": [],
    }
    second_payload = json.loads(client.calls[1][1]["data"])
    assert second_payload["cancel_cursor"] == first_ids[-1]


def _mtls_files(tmp_path: Path):
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "edge-gateway-1")])
    now = datetime.now(UTC)
    cert = (
        x509.CertificateBuilder()
        .subject_name(name)
        .issuer_name(name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(minutes=1))
        .not_valid_after(now + timedelta(days=1))
        .sign(key, hashes.SHA256())
    )
    cert_path = tmp_path / "edge-cert.pem"
    key_path = tmp_path / "edge-key.pem"
    ca_path = tmp_path / "enterprise-ca.pem"
    cert_pem = cert.public_bytes(serialization.Encoding.PEM)
    cert_path.write_bytes(cert_pem)
    ca_path.write_bytes(cert_pem)
    key_path.write_bytes(
        key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )
    )
    key_path.chmod(0o600)
    return cert_path, key_path, ca_path, cert.fingerprint(hashes.SHA256()).hex()


def test_bootstrap_requires_real_mtls_material_and_matching_fingerprint(tmp_path):
    cert, key, ca, fingerprint = _mtls_files(tmp_path)
    value = {
        "queue_path": str(tmp_path / "edge.sqlite3"),
        "artifact_root": str(tmp_path),
        "gateway_id": "gateway-1",
        "credential": "dopg_local-test-credential",
        "certificate_sha256": fingerprint,
        "client_certificate_path": str(cert),
        "client_private_key_path": str(key),
        "ca_bundle_path": str(ca),
        "generation": 1,
        "environment": "production",
        "network_zone": "zone-a",
        "policy_digest": "a" * 64,
        "control_url": "https://control.enterprise.test",
        "proxy_url": None,
        "allowed_control_hosts": ["control.enterprise.test"],
        "allowed_proxy_hosts": [],
        "trusted_release_keys": {"release-key-1": "01" * 32},
        "trusted_task_keys": {"task-key-1": "02" * 32},
        "task_authority_clock_skew_seconds": 60,
        "version": "3.0.0",
    }
    config = EdgeBootstrapConfig.from_mapping(value)
    assert config.certificate_sha256 == fingerprint
    with pytest.raises(ValueError, match="fingerprint"):
        EdgeBootstrapConfig.from_mapping({**value, "certificate_sha256": "f" * 64})


@pytest.mark.parametrize(
    "url,host",
    [
        ("https://control.enterprise.test:444", "control.enterprise.test"),
        ("https://127.0.0.1", "127.0.0.1"),
        ("https://éxample.test", "éxample.test"),
    ],
)
def test_transport_rejects_unapproved_port_ip_literal_and_unicode_host(url, host):
    with pytest.raises((ValueError, EdgePolicyError)):
        EdgeTransport(
            base_url=url,
            allowed_control_hosts={host},
            gateway_id="gateway-1",
            environment="production",
            network_zone="zone-a",
            generation=1,
            credential="dopg_example-credential",
            certificate_sha256="b" * 64,
            client=NeverNetwork(),
            client_certificate_path="/secure/edge-cert.pem",
            client_private_key_path="/secure/edge-key.pem",
            ca_bundle_path="/secure/ca.pem",
        )


class StreamingResponse:
    def __init__(self, chunks, *, content_length=None):
        self.status_code = 200
        self.headers = {}
        if content_length is not None:
            self.headers["Content-Length"] = str(content_length)
        self._chunks = chunks

    @property
    def content(self):  # pragma: no cover - accessing this is the bug
        raise AssertionError("streaming transport must not buffer response.content")

    def iter_content(self, chunk_size):
        assert chunk_size <= 65_536
        yield from self._chunks


class CloseTrackingResponse(StreamingResponse):
    def __init__(self, chunks, *, status_code=200, content_length=None, read_error=None):
        super().__init__(chunks, content_length=content_length)
        self.status_code = status_code
        self.read_error = read_error
        self.close_count = 0

    def iter_content(self, chunk_size):
        if self.read_error is not None:
            raise self.read_error
        yield from super().iter_content(chunk_size)

    def close(self):
        self.close_count += 1


@pytest.mark.parametrize(
    "response,expected_exception",
    [
        (
            CloseTrackingResponse(
                [b'{"code":200,"message":"ok","data":{"task":null}}']
            ),
            None,
        ),
        (CloseTrackingResponse([], status_code=302), EdgeTransportError),
        (CloseTrackingResponse([b"{}"], status_code=401), EdgeAuthenticationStopped),
        (CloseTrackingResponse([b"{}"], status_code=500), EdgeTransportError),
        (
            CloseTrackingResponse([], content_length=EdgeTransport.MAX_RESPONSE_BYTES + 1),
            EdgeTransportError,
        ),
        (
            CloseTrackingResponse(
                [b"x" * (EdgeTransport.MAX_RESPONSE_BYTES + 1)], content_length=1
            ),
            EdgeTransportError,
        ),
        (
            CloseTrackingResponse([], read_error=TimeoutError("slow stream")),
            EdgeTransportError,
        ),
    ],
)
def test_transport_closes_streamed_response_exactly_once_on_every_exit(
    response, expected_exception
):
    transport = _transport(RecordingClient(response))
    if expected_exception is None:
        assert transport.pull_task() is None
    else:
        with pytest.raises(expected_exception):
            transport.pull_task()
    assert response.close_count == 1


def test_transport_streams_with_mtls_and_rejects_deceptive_chunk_overflow():
    body = b'{"code":200,"message":"ok","data":{"task":null}}'
    client = RecordingClient(StreamingResponse([body], content_length=len(body)))
    transport = EdgeTransport(
        base_url="https://control.enterprise.test",
        allowed_control_hosts={"control.enterprise.test"},
        gateway_id="gateway-1",
        environment="production",
        network_zone="zone-a",
        generation=1,
        credential="dopg_example-credential",
        certificate_sha256="b" * 64,
        client=client,
        client_certificate_path="/secure/edge-cert.pem",
        client_private_key_path="/secure/edge-key.pem",
        ca_bundle_path="/secure/ca.pem",
    )
    assert transport.pull_task() is None
    _, kwargs = client.calls[0]
    assert kwargs["stream"] is True
    assert kwargs["cert"] == ("/secure/edge-cert.pem", "/secure/edge-key.pem")
    assert kwargs["verify"] == "/secure/ca.pem"
    assert kwargs["timeout"] == (5, 20)

    overflow = RecordingClient(
        StreamingResponse(
            [b"x" * 700_000, b"y" * 700_000],
            content_length=1,
        )
    )
    bad = EdgeTransport(
        base_url="https://control.enterprise.test",
        allowed_control_hosts={"control.enterprise.test"},
        gateway_id="gateway-1",
        environment="production",
        network_zone="zone-a",
        generation=1,
        credential="dopg_example-credential",
        certificate_sha256="b" * 64,
        client=overflow,
        client_certificate_path="/secure/edge-cert.pem",
        client_private_key_path="/secure/edge-key.pem",
        ca_bundle_path="/secure/ca.pem",
    )
    with pytest.raises(EdgeTransportError, match="byte limit"):
        bad.pull_task()


def test_runner_bridge_has_exact_local_only_schema():
    request = {
        "task_id": "task-1",
        "operation": "profile",
        "purpose": "inventory",
        "classification": "statistics",
        "environment": "production",
        "network_zone": "zone-a",
        "idempotency_key": "idem-1",
        "deadline_at": "2030-01-01T00:00:00Z",
    }
    handlers = {"edge.profile": lambda value, _cancel: {"seen": sorted(value)}}
    assert execute_edge_adapter(handlers, "edge.profile", request, lambda: False) == {
        "seen": sorted(request)
    }
    with pytest.raises(NodeExecutionError):
        execute_edge_adapter(handlers, "python", request, lambda: False)
    with pytest.raises(NodeExecutionError):
        execute_edge_adapter(
            handlers,
            "edge.profile",
            {**request, "sql": "SELECT 1"},
            lambda: False,
        )