GTI_GERENCIADOR_FISCAL/fiscal/central/tests/test_agent_distribution.py

"""Agent distribution: releases + secure enrollment.

Covers central/app/agent_distribution.py, enrollment.py, release_storage.py,
rate_limit.py, scripts/publish_agent_release.py and the new admin pages in
web.py. Security tests mirror the mission's checklist: expired/revoked/
second-use/tampered download links, expired/replayed/wrong-installation/
already-enrolled enrollment attempts, path traversal, rate limiting, and
that no secret ever reaches the audit log.
"""
import dataclasses
import io
import json
import sys
import tempfile
import zipfile
from datetime import datetime, timedelta, timezone
from pathlib import Path

import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select

from central.app import config, enrollment, rate_limit, release_storage
from central.app.main import create_app
from central.app.models import Agent, AgentEnrollment, AgentRelease, AuditEvent, Installation, Tenant
from central.scripts.publish_agent_release import publish
from central.tests.test_control_plane import seed, setup_db

VERSION = "1.2.3"


def make_release_zip(version=VERSION) -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as z:
        z.writestr(f"GTIFiscalAgent-{version}/gti-agent.ps1", "# fake install script\n")
        z.writestr(f"GTIFiscalAgent-{version}/VERSION", version)
    return buf.getvalue()


def write_temp_zip(data: bytes, version=VERSION) -> Path:
    path = Path(tempfile.gettempdir()) / f"GTIFiscalAgent-{version}-{id(data)}.zip"
    path.write_bytes(data)
    return path


@pytest.fixture(autouse=True)
def clear_rate_limit():
    rate_limit._attempts.clear()
    yield
    rate_limit._attempts.clear()


@pytest.fixture
def hub(tmp_path, monkeypatch):
    replaced = dataclasses.replace(
        config.settings,
        agent_release_storage_path=str(tmp_path / "releases"),
        agent_download_link_ttl_seconds=1800,
        agent_download_max_uses_default=1,
        agent_enrollment_ttl_seconds=86400,
        base_url="https://gtifiscal.example.gti",
        web_cookie_secure=False,
    )
    monkeypatch.setattr(config, "settings", replaced)
    # publish_agent_release.py is written to also run standalone (`python
    # central/scripts/publish_agent_release.py`), so it imports via a
    # sys.path-inserted "app" package (same convention as register_agent.py)
    # rather than "central.app" — a second, independent copy of every module
    # it touches, including config. Patch that copy's settings too, or the
    # release it publishes lands outside this test's tmp_path.
    if "app.config" in sys.modules:
        monkeypatch.setattr(sys.modules["app.config"], "settings", replaced)
    factory = setup_db()
    seed(factory)  # tenant "Acme" + Installation "srv-01" + admin "operator"/"a-secure-password"
    client = TestClient(create_app(database_probe=lambda: True, session_factory=factory, bootstrap_token=None))
    return client, factory


def login(client):
    response = client.post("/login", data={"username": "operator", "password": "a-secure-password"}, follow_redirects=False)
    assert response.status_code == 303


def csrf(client):
    return client.cookies.get("gti_csrf", "")


def publish_release(factory, version=VERSION):
    with factory() as session:
        release, _created = publish(session, write_temp_zip(make_release_zip(version), version), version=version)
        return release.id


def tenant_id(factory):
    with factory() as session:
        tenant = session.scalar(select(Tenant).where(Tenant.name == "Acme"))
        # /agents/new only offers active GTI_GERENCIADOR-sourced tenants
        # (central/app/client_sync.py) — flip the shared fixture tenant so
        # existing form-based tests keep working under that rule.
        if tenant.source != "GTI_GERENCIADOR" or not tenant.active:
            tenant.source, tenant.external_client_id, tenant.active = "GTI_GERENCIADOR", tenant.id, True
            session.commit()
        return tenant.id


def audits(factory):
    with factory() as session:
        return session.scalars(select(AuditEvent)).all()


def bundle_enrollment_fields(content: bytes) -> dict:
    with zipfile.ZipFile(io.BytesIO(content)) as z:
        [entry] = [n for n in z.namelist() if n.endswith("enrollment.json")]
        return json.loads(z.read(entry))


# ---------- publish_agent_release.py ----------

def test_publish_is_idempotent_and_rejects_hash_conflict(hub, tmp_path):
    _client, factory = hub
    with factory() as session:
        zip_bytes = make_release_zip("9.9.9")
        release, created = publish(session, write_temp_zip(zip_bytes, "9.9.9"), version="9.9.9")
        assert created and release.status == "PUBLISHED" and release.version == "9.9.9"
        release2, created2 = publish(session, write_temp_zip(zip_bytes, "9.9.9"), version="9.9.9")
        assert not created2 and release2.id == release.id  # identical bytes: no-op
    with factory() as session:
        with pytest.raises(SystemExit):
            publish(session, write_temp_zip(make_release_zip("9.9.9") + b"\x00", "9.9.9"), version="9.9.9")


def test_publish_records_published_at(hub, tmp_path):
    _client, factory = hub
    with factory() as session:
        release, created = publish(session, write_temp_zip(make_release_zip("9.9.9"), "9.9.9"), version="9.9.9")
        assert created and release.status == "PUBLISHED"
        assert release.published_at is not None  # publishing stamps the moment, not left NULL


def test_published_release_binary_never_stored_in_postgres(hub):
    _client, factory = hub
    release_id = publish_release(factory)
    with factory() as session:
        release = session.get(AgentRelease, release_id)
        assert release.storage_path == Path(release.storage_path).name  # a bare filename, never a path with directories
        path = release_storage.safe_path_for(release.storage_path)
        assert path.exists() and path.read_bytes()  # binary lives on disk, outside the DB


# ---------- path traversal ----------

def test_release_storage_path_traversal_impossible(hub):
    for hostile in ("../../etc/passwd", "sub/dir.zip", "/etc/passwd", "..\\..\\windows\\system32", ""):
        with pytest.raises(ValueError):
            release_storage.safe_path_for(hostile)


# ---------- full happy path: create -> link -> download -> enroll ----------

def test_full_installation_flow_creates_a_working_agent(hub):
    client, factory = hub
    release_id = publish_release(factory)
    login(client)

    created = client.post("/agents/new", data={"csrf_token": csrf(client), "client_id": tenant_id(factory),
                                                 "display_name": "Servidor Principal", "expected_hostname": "SRV-UNIAO",
                                                 "release_id": release_id}, follow_redirects=False)
    assert created.status_code == 303
    installation_uuid = created.headers["location"].rsplit("/", 1)[-1]

    linked = client.post(f"/agents/pending/{installation_uuid}/link", data={"csrf_token": csrf(client)}, follow_redirects=False)
    assert linked.status_code == 303
    download_url = linked.headers["location"].split("link=", 1)[1]
    download_token = download_url.rsplit("/", 1)[-1]

    # A technician on the client machine never has an admin session.
    anon = TestClient(client.app)
    downloaded = anon.get(f"/dl/agent/{download_token}")
    assert downloaded.status_code == 200
    assert downloaded.headers["content-type"] == "application/octet-stream"
    assert "attachment" in downloaded.headers["content-disposition"]

    payload = bundle_enrollment_fields(downloaded.content)
    assert payload["installation_id"] == installation_uuid

    enrolled = anon.post("/v1/agents/enroll", json={
        "installation_id": payload["installation_id"], "enrollment_token": payload["enrollment_token"],
        "agent_version": "1.2.3", "hostname": "SRV-UNIAO", "platform": "windows", "machine_fingerprint": "abc123",
    })
    assert enrolled.status_code == 200
    body = enrolled.json()
    assert body["agent_id"] and body["agent_token"] and body["hostname_mismatch"] is False

    with factory() as session:
        agent = session.get(Agent, body["agent_id"])
        assert agent.enabled and agent.agent_version == "1.2.3"
        installation = session.scalar(select(Installation).where(Installation.agent_id == agent.id))
        assert installation.installation_id == installation_uuid and installation.hostname == "SRV-UNIAO"
        enr = session.scalar(select(AgentEnrollment).where(AgentEnrollment.enrollment_uuid == installation_uuid))
        assert enr.status == "ENROLLED" and enr.enrollment_token_hash is None and enr.agent_id == agent.id

    # Replay: the very same enrollment token can never enroll a second Agent.
    replay = anon.post("/v1/agents/enroll", json={"installation_id": payload["installation_id"], "enrollment_token": payload["enrollment_token"]})
    assert replay.status_code == 401 and replay.json()["detail"] == "ALREADY_ENROLLED"

    # Second use of the (single-use) download link is also rejected.
    redownload = anon.get(f"/dl/agent/{download_token}")
    assert redownload.status_code == 404

    events = {event.action for event in audits(factory)}
    assert {"INSTALLATION_CREATED", "DOWNLOAD_LINK_CREATED", "PACKAGE_DOWNLOADED", "ENROLLMENT_SUCCEEDED", "AGENT_CREDENTIAL_ISSUED", "ENROLLMENT_REJECTED"} <= events

    secrets_leaked = [download_token, payload["enrollment_token"], body["agent_token"]]
    for event in audits(factory):
        haystack = json.dumps({"action": event.action, "target": event.target_id, "context": event.context})
        for secret in secrets_leaked:
            assert secret not in haystack


def test_direct_authenticated_download_does_not_need_a_link(hub):
    client, factory = hub
    release_id = publish_release(factory)
    login(client)
    created = client.post("/agents/new", data={"csrf_token": csrf(client), "client_id": tenant_id(factory),
                                                 "display_name": "Loja 2", "release_id": release_id}, follow_redirects=False)
    installation_uuid = created.headers["location"].rsplit("/", 1)[-1]
    downloaded = client.get(f"/agents/pending/{installation_uuid}/download")
    assert downloaded.status_code == 200
    payload = bundle_enrollment_fields(downloaded.content)
    assert payload["installation_id"] == installation_uuid


# ---------- route precedence: /agents/new vs /agents/{agent_id} ----------

def test_get_agents_new_renders_form_not_captured_by_dynamic_route(hub):
    client, factory = hub
    publish_release(factory)
    tenant_id(factory)  # at least one active GTI_GERENCIADOR tenant, or the form shows the empty state instead
    login(client)
    response = client.get("/agents/new")
    assert response.status_code == 200
    assert "text/html" in response.headers["content-type"]
    assert "Nova instalação" in response.text
    assert "Criar instalação" in response.text


def test_get_agents_releases_page_not_captured_by_dynamic_route(hub):
    client, _factory = hub
    login(client)
    response = client.get("/agents/releases")
    assert response.status_code == 200
    assert "text/html" in response.headers["content-type"]
    assert "Versões do Agent" in response.text


def test_get_agent_detail_still_renders_for_valid_numeric_id(hub):
    client, factory = hub
    login(client)
    with factory() as session:
        agent_id = session.scalar(select(Agent.id))
    response = client.get(f"/agents/{agent_id}")
    assert response.status_code == 200
    assert "text/html" in response.headers["content-type"]


def test_get_agents_non_numeric_slug_is_controlled_404_and_never_captures_statics(hub):
    client, _factory = hub
    login(client)
    response = client.get("/agents/abc")
    assert response.status_code == 404
    assert "text/html" in response.headers["content-type"]
    # the legitimate static routes keep working regardless of the slug above
    assert client.get("/agents/new").status_code == 200
    assert client.get("/agents/releases").status_code == 200


# ---------- download-link security ----------

def _pending(factory, release_id, **overrides):
    import uuid as uuid_mod
    with factory() as session:
        enr = AgentEnrollment(enrollment_uuid=str(uuid_mod.uuid4()), tenant_id=tenant_id(factory),
                               display_name="Teste", release_id=release_id, status="PENDING", **overrides)
        session.add(enr); session.commit()
        return enr.enrollment_uuid


def test_expired_download_link_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    uid = _pending(factory, release_id)
    token = enrollment.new_token()
    with factory() as session:
        enr = session.scalar(select(AgentEnrollment).where(AgentEnrollment.enrollment_uuid == uid))
        enr.download_token_hash = enrollment.token_digest(token)
        enr.download_token_expires_at = datetime.now(timezone.utc) - timedelta(seconds=1)
        enr.download_max_uses = 1
        session.commit()
    assert client.get(f"/dl/agent/{token}").status_code == 404


def test_revoked_installation_link_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    login(client)
    created = client.post("/agents/new", data={"csrf_token": csrf(client), "client_id": tenant_id(factory),
                                                 "display_name": "Teste", "release_id": release_id}, follow_redirects=False)
    uid = created.headers["location"].rsplit("/", 1)[-1]
    linked = client.post(f"/agents/pending/{uid}/link", data={"csrf_token": csrf(client)}, follow_redirects=False)
    token = linked.headers["location"].split("link=", 1)[1].rsplit("/", 1)[-1]
    client.post(f"/agents/pending/{uid}/revoke", data={"csrf_token": csrf(client)}, follow_redirects=False)
    assert client.get(f"/dl/agent/{token}").status_code == 404


def test_wrong_or_random_download_token_rejected(hub):
    client, factory = hub
    publish_release(factory)
    assert client.get("/dl/agent/definitely-not-a-real-token").status_code == 404


def test_release_revoked_blocks_download(hub):
    client, factory = hub
    release_id = publish_release(factory)
    login(client)
    created = client.post("/agents/new", data={"csrf_token": csrf(client), "client_id": tenant_id(factory),
                                                 "display_name": "Teste", "release_id": release_id}, follow_redirects=False)
    uid = created.headers["location"].rsplit("/", 1)[-1]
    linked = client.post(f"/agents/pending/{uid}/link", data={"csrf_token": csrf(client)}, follow_redirects=False)
    token = linked.headers["location"].split("link=", 1)[1].rsplit("/", 1)[-1]
    client.post(f"/agents/releases/{release_id}/revoke", data={"csrf_token": csrf(client)}, follow_redirects=False)
    assert client.get(f"/dl/agent/{token}").status_code == 404


# ---------- enrollment security ----------

def _enrollable(factory, release_id, **overrides):
    token = enrollment.new_token()
    uid = _pending(factory, release_id)
    with factory() as session:
        enr = session.scalar(select(AgentEnrollment).where(AgentEnrollment.enrollment_uuid == uid))
        enr.enrollment_token_hash = enrollment.token_digest(token)
        enr.enrollment_token_expires_at = overrides.pop("expires_at", datetime.now(timezone.utc) + timedelta(hours=1))
        for key, value in overrides.items():
            setattr(enr, key, value)
        session.commit()
    return uid, token


def test_expired_enrollment_token_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    uid, token = _enrollable(factory, release_id, expires_at=datetime.now(timezone.utc) - timedelta(seconds=1))
    response = client.post("/v1/agents/enroll", json={"installation_id": uid, "enrollment_token": token})
    assert response.status_code == 401 and response.json()["detail"] == "EXPIRED"


def test_wrong_installation_uuid_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    _uid, token = _enrollable(factory, release_id)
    import uuid as uuid_mod
    response = client.post("/v1/agents/enroll", json={"installation_id": str(uuid_mod.uuid4()), "enrollment_token": token})
    assert response.status_code == 401 and response.json()["detail"] == "INVALID_INSTALLATION"


def test_malformed_installation_id_rejected(hub):
    client, _factory = hub
    response = client.post("/v1/agents/enroll", json={"installation_id": "not-a-uuid", "enrollment_token": "x"})
    assert response.status_code == 401 and response.json()["detail"] == "INVALID_INSTALLATION"


def test_random_enrollment_token_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    uid, _token = _enrollable(factory, release_id)
    response = client.post("/v1/agents/enroll", json={"installation_id": uid, "enrollment_token": enrollment.new_token()})
    assert response.status_code == 401 and response.json()["detail"] == "INVALID_TOKEN"


def test_revoked_installation_enroll_rejected(hub):
    client, factory = hub
    release_id = publish_release(factory)
    uid, token = _enrollable(factory, release_id, status="REVOKED")
    response = client.post("/v1/agents/enroll", json={"installation_id": uid, "enrollment_token": token})
    assert response.status_code == 401 and response.json()["detail"] == "REVOKED"


def test_hostname_mismatch_is_flagged_but_not_blocked(hub):
    client, factory = hub
    release_id = publish_release(factory)
    uid, token = _enrollable(factory, release_id, expected_hostname="SRV-EXPECTED")
    response = client.post("/v1/agents/enroll", json={"installation_id": uid, "enrollment_token": token, "hostname": "SRV-OUTRO"})
    assert response.status_code == 200 and response.json()["hostname_mismatch"] is True
    with factory() as session:
        enr = session.scalar(select(AgentEnrollment).where(AgentEnrollment.enrollment_uuid == uid))
        assert enr.hostname_mismatch is True and enr.status == "ENROLLED"


# ---------- rate limiting ----------

def test_enroll_rate_limited_after_repeated_invalid_attempts(hub):
    client, _factory = hub
    for _ in range(rate_limit.MAX_ATTEMPTS):
        response = client.post("/v1/agents/enroll", json={"installation_id": "not-a-uuid", "enrollment_token": "x"})
        assert response.status_code == 401
    blocked_response = client.post("/v1/agents/enroll", json={"installation_id": "not-a-uuid", "enrollment_token": "x"})
    assert blocked_response.status_code == 429


def test_rate_limit_buckets_are_independent_per_ip(hub):
    # Unit-level proof that the limiter itself keys on (bucket, ip) rather than a
    # shared global counter — the end-to-end TestClient can't vary the caller's
    # apparent IP (that requires the real ProxyHeadersMiddleware, wired up only
    # under uvicorn --proxy-headers in production, per FORWARDED_ALLOW_IPS).
    for _ in range(rate_limit.MAX_ATTEMPTS):
        rate_limit.record_attempt("enroll", "198.51.100.1")
    assert rate_limit.blocked("enroll", "198.51.100.1")
    assert not rate_limit.blocked("enroll", "198.51.100.2")