GTI_GERENCIADOR_FISCAL/central/app/agent_distribution.py

"""Public, pre-authentication routes for Agent distribution.

GET /dl/agent/{token} and POST /v1/agents/enroll are reachable by anyone —
same trust level as GET /login. Every other Agent-distribution route
(create pending installation, generate a link, direct authenticated
download, list/revoke releases) is a control-panel page and lives in
web.py, gated the same way as the rest of the panel (admin session + CSRF).

The enrollment token embedded in a downloaded bundle is generated fresh at
download time and its plaintext is never written to storage — only its
SHA-256 hash persists, exactly like the download token that gated getting
it. Neither is ever the Agent's own permanent Bearer token, which is
minted only once, inside a successful POST /v1/agents/enroll.
"""
from __future__ import annotations

import secrets
from datetime import datetime, timedelta, timezone

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import Response
from sqlalchemy import select

from . import config, enrollment, release_storage
from .audit import record_audit
from .auth import token_hash as agent_token_hash
from .models import Agent, AgentEnrollment, AgentRelease, Installation
from .rate_limit import blocked, record_attempt


def client_ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


def _aware(value: datetime | None) -> datetime | None:
    """SQLite (used by the test suite) drops tzinfo on round-trip even for a
    DateTime(timezone=True) column; Postgres never does. Normalize either way
    before comparing against datetime.now(timezone.utc) — same pattern as
    human_auth.py's session-expiry check."""
    if value is None or value.tzinfo is not None:
        return value
    return value.replace(tzinfo=timezone.utc)


def build_and_serve(session, enr: AgentEnrollment, request: Request, *, actor_admin_id: int | None = None) -> Response:
    """Mint a fresh enrollment token, build the per-installation bundle and
    return it as a download response. Shared by the temp-link route and the
    admin's direct "Baixar Agent" button in web.py."""
    release = session.get(AgentRelease, enr.release_id)
    if release is None or release.status != "PUBLISHED":
        raise enrollment.EnrollmentRejected("RELEASE_UNAVAILABLE")
    path = release_storage.safe_path_for(release.storage_path)
    if not path.exists():
        raise enrollment.EnrollmentRejected("RELEASE_FILE_MISSING")
    token = enrollment.new_token()
    now = datetime.now(timezone.utc)
    enr.enrollment_token_hash = enrollment.token_digest(token)
    enr.enrollment_token_expires_at = now + timedelta(seconds=config.settings.agent_enrollment_ttl_seconds)
    enr.enrollment_used_at = None
    enr.downloaded_at = now
    if enr.status not in ("ENROLLED", "REVOKED"):
        enr.status = "DOWNLOADED"
    plan = enrollment.BundlePlan(hub_url=config.settings.base_url, enrollment_uuid=enr.enrollment_uuid, enrollment_token=token)
    data = enrollment.build_bundle(path.read_bytes(), plan)
    record_audit(session, admin_user_id=actor_admin_id, tenant_id=enr.tenant_id, action="PACKAGE_DOWNLOADED",
                 target_type="agent_enrollment", target_id=enr.enrollment_uuid,
                 context={"ip": client_ip(request), "release_version": release.version, "release_id": release.id})
    session.commit()
    filename = f"GTIFiscalAgent-{release.version}-{enr.enrollment_uuid[:8]}.zip"
    return Response(content=data, media_type="application/octet-stream",
                     headers={"Content-Disposition": f'attachment; filename="{filename}"',
                              "X-Content-Type-Options": "nosniff", "Cache-Control": "no-store"})


def router(session_factory):
    routes = APIRouter()

    @routes.get("/dl/agent/{token}")
    def download_by_link(token: str, request: Request):
        ip = client_ip(request)
        if blocked("download", ip):
            raise HTTPException(429, "too many attempts")
        with session_factory() as session:
            candidates = session.scalars(select(AgentEnrollment).where(AgentEnrollment.download_token_hash.is_not(None))).all()
            enr = next((c for c in candidates if enrollment.token_matches(token, c.download_token_hash)), None)
            now = datetime.now(timezone.utc)
            try:
                if enr is None:
                    raise enrollment.EnrollmentRejected("INVALID_LINK")
                if enr.status == "REVOKED":
                    raise enrollment.EnrollmentRejected("REVOKED")
                if not enr.download_token_expires_at or _aware(enr.download_token_expires_at) < now:
                    raise enrollment.EnrollmentRejected("EXPIRED_LINK")
                if enr.download_uses >= enr.download_max_uses:
                    raise enrollment.EnrollmentRejected("LINK_ALREADY_USED")
                enr.download_uses += 1
                if enr.download_uses >= enr.download_max_uses:
                    enr.download_token_hash = None  # the link's last remaining use is spent: it dies immediately
                return build_and_serve(session, enr, request)
            except enrollment.EnrollmentRejected as rejected:
                session.rollback()
                record_attempt("download", ip)
                record_audit(session, admin_user_id=None, tenant_id=enr.tenant_id if enr else None, action="DOWNLOAD_LINK_REJECTED",
                              target_type="agent_enrollment", target_id=enr.enrollment_uuid if enr else "-",
                              context={"ip": ip, "reason": rejected.reason}, severity="WARNING")
                session.commit()
                raise HTTPException(404, "invalid or expired link")

    @routes.post("/v1/agents/enroll")
    async def enroll(request: Request):
        ip = client_ip(request)
        if blocked("enroll", ip):
            raise HTTPException(429, "too many attempts")
        try:
            payload = await request.json()
        except Exception:
            payload = {}
        installation_id = str(payload.get("installation_id") or "")
        supplied_token = payload.get("enrollment_token")
        with session_factory() as session:
            enr = None
            try:
                if not enrollment.valid_uuid(installation_id):
                    raise enrollment.EnrollmentRejected("INVALID_INSTALLATION")
                enr = session.scalar(select(AgentEnrollment).where(AgentEnrollment.enrollment_uuid == installation_id))
                if enr is None:
                    raise enrollment.EnrollmentRejected("INVALID_INSTALLATION")
                if enr.status == "REVOKED":
                    raise enrollment.EnrollmentRejected("REVOKED")
                if enr.status == "ENROLLED":
                    raise enrollment.EnrollmentRejected("ALREADY_ENROLLED")
                now = datetime.now(timezone.utc)
                if not enr.enrollment_token_expires_at or _aware(enr.enrollment_token_expires_at) < now:
                    raise enrollment.EnrollmentRejected("EXPIRED")
                if not enrollment.token_matches(supplied_token, enr.enrollment_token_hash):
                    raise enrollment.EnrollmentRejected("INVALID_TOKEN")
                release = session.get(AgentRelease, enr.release_id)
                if release is None or release.status != "PUBLISHED":
                    raise enrollment.EnrollmentRejected("RELEASE_UNAVAILABLE")

                hostname = str(payload.get("hostname") or "")[:255]
                mismatch = bool(enr.expected_hostname) and bool(hostname) and hostname.lower() != enr.expected_hostname.lower()

                plain_token = secrets.token_urlsafe(32)
                agent = Agent(tenant_id=enr.tenant_id, name=enr.display_name, token_hash=agent_token_hash(plain_token),
                               enabled=True, agent_version=str(payload.get("agent_version") or "")[:32] or None)
                session.add(agent)
                session.flush()
                session.add(Installation(agent_id=agent.id, tenant_id=enr.tenant_id, installation_id=enr.enrollment_uuid,
                                          display_name=enr.display_name, hostname=hostname or None))

                enr.status = "ENROLLED"
                enr.enrollment_used_at = now
                enr.enrollment_token_hash = None  # replay-proof even inside its TTL window
                enr.hostname_seen = hostname or None
                enr.machine_fingerprint_seen = str(payload.get("machine_fingerprint") or "")[:128] or None
                enr.hostname_mismatch = mismatch
                enr.agent_id = agent.id

                record_audit(session, admin_user_id=None, tenant_id=enr.tenant_id, action="ENROLLMENT_SUCCEEDED",
                             target_type="agent_enrollment", target_id=enr.enrollment_uuid,
                             context={"ip": ip, "hostname": hostname, "hostname_mismatch": mismatch, "agent_id": agent.id})
                record_audit(session, admin_user_id=None, tenant_id=enr.tenant_id, action="AGENT_CREDENTIAL_ISSUED",
                             target_type="agent", target_id=str(agent.id), context={"ip": ip, "installation": enr.enrollment_uuid})
                session.commit()
                return {"agent_id": agent.id, "agent_token": plain_token, "tenant_id": enr.tenant_id,
                        "hub_url": config.settings.base_url, "hostname_mismatch": mismatch}
            except enrollment.EnrollmentRejected as rejected:
                session.rollback()
                record_attempt("enroll", ip)
                record_audit(session, admin_user_id=None, tenant_id=enr.tenant_id if enr else None, action="ENROLLMENT_REJECTED",
                             target_type="agent_enrollment", target_id=installation_id or "-",
                             context={"ip": ip, "reason": rejected.reason}, severity="WARNING")
                session.commit()
                raise HTTPException(401, rejected.reason)

    return routes