admin / Syanpse-Vanguard

"""Agent registry — enrollment, identity binding, revocation, replay bookkeeping."""
from __future__ import annotations

import hashlib
import secrets
import uuid
from datetime import datetime, timezone

from sv.config import DEFAULT_TENANT_ID
from sv.models.agent import AgentRegisterRequest
from sv.storage.base import Query, StorageEngine


def _now() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


def hash_token(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


async def enroll(engine: StorageEngine, req: AgentRegisterRequest,
                 cert_cn: str | None, ip: str | None = None) -> tuple[str, str | None]:
    """Create an agent row. Returns (agent_id, ingest_token). In mTLS mode the
    token is None and identity is the cert CN; in dev mode a bearer token is
    issued and only its hash is stored."""
    import json

    agent_id = str(uuid.uuid4())
    ingest_token = None if cert_cn else secrets.token_urlsafe(32)
    await engine.run(Query(
        "INSERT INTO agents (id, tenant_id, hostname, platform, agent_version, "
        "channels, cert_cn, ingest_token_hash, last_seq, last_ip, revoked, enrolled_ts) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, ?, FALSE, ?)",
        [agent_id, str(DEFAULT_TENANT_ID), req.hostname, req.platform,
         req.agent_version, json.dumps(req.channels), cert_cn,
         hash_token(ingest_token) if ingest_token else None, ip, _now()],
    ))
    return agent_id, ingest_token


async def get_agent(engine: StorageEngine, agent_id: str) -> dict | None:
    return await engine.fetchone(Query(
        "SELECT id, cert_cn, ingest_token_hash, last_seq, revoked "
        "FROM agents WHERE id = ?", [agent_id],
    ))


async def find_by_cert_cn(engine: StorageEngine, cert_cn: str) -> dict | None:
    return await engine.fetchone(Query(
        "SELECT id, revoked FROM agents WHERE cert_cn = ?", [cert_cn],
    ))


async def touch_seq(engine: StorageEngine, agent_id: str, seq: int,
                    ip: str | None = None) -> None:
    if ip:
        await engine.run(Query(
            "UPDATE agents SET last_seq = ?, last_seen_ts = ?, last_ip = ? WHERE id = ?",
            [seq, _now(), ip, agent_id],
        ))
    else:
        await engine.run(Query(
            "UPDATE agents SET last_seq = ?, last_seen_ts = ? WHERE id = ?",
            [seq, _now(), agent_id],
        ))


async def revoke(engine: StorageEngine, agent_id: str) -> None:
    await engine.run(Query(
        "UPDATE agents SET revoked = TRUE WHERE id = ?", [agent_id],
    ))