admin / Syanpse-Vanguard

"""Enrollment + reliable shipper. Registers once (persisting identity), then
loops: pull a batch from the disk buffer, zstd-compress, POST to /v1/ingest/bulk
with mTLS (prod) or bearer token (dev), and trim the buffer only after the
server acks. Network loss just grows the buffer; it drains on reconnect."""
from __future__ import annotations

import json
import logging
import time
from datetime import datetime, timezone

import httpx
import zstandard

from agents.common.buffer import DiskBuffer
from agents.common.config import AgentConfig, load_state, save_state

log = logging.getLogger("sv.agent.sender")
_ZSTD = zstandard.ZstdCompressor(level=3)


class Sender:
    def __init__(self, cfg: AgentConfig, buffer: DiskBuffer):
        self.cfg = cfg
        self.buffer = buffer
        self.state = load_state(cfg.state_path)
        self._client = self._build_client()

    def _build_client(self) -> httpx.Client:
        kwargs: dict = {"timeout": 30.0, "base_url": self.cfg.ingest_url}
        if self.cfg.mtls_enabled:
            kwargs["cert"] = (self.cfg.client_cert, self.cfg.client_key)
            kwargs["verify"] = self.cfg.ca_cert or self.cfg.verify_tls
        else:
            kwargs["verify"] = self.cfg.verify_tls
        return httpx.Client(**kwargs)

    # --- enrollment ---------------------------------------------------------
    def ensure_enrolled(self) -> None:
        if self.state.get("agent_id"):
            return
        resp = self._client.post(
            "/v1/agents/register",
            headers={"Authorization": f"Bearer {self.cfg.enrollment_token}"},
            json={
                "hostname": self.cfg.hostname,
                "platform": self.cfg.platform,
                "agent_version": self.cfg.agent_version,
                "channels": self.cfg.channels,
            },
        )
        resp.raise_for_status()
        data = resp.json()
        self.state["agent_id"] = data["agent_id"]
        if data.get("ingest_token"):
            self.state["ingest_token"] = data["ingest_token"]
        save_state(self.cfg.state_path, self.state)
        log.info("enrolled as agent_id=%s", self.state["agent_id"])

    # --- shipping -----------------------------------------------------------
    def flush_once(self) -> int:
        max_id, lines = self.buffer.next_batch(self.cfg.batch_size)
        if not lines:
            return 0
        body = _ZSTD.compress(("\n".join(lines)).encode())
        bundle = {
            "agent_id": self.state["agent_id"],
            "seq": max_id,
            "codec": "zstd",
            "event_count": len(lines),
            "sent_at": datetime.now(timezone.utc).isoformat(),
        }
        headers = {"X-SV-Bundle": json.dumps(bundle),
                   "Content-Type": "application/octet-stream"}
        if not self.cfg.mtls_enabled:
            headers["Authorization"] = f"Bearer {self.state['ingest_token']}"

        resp = self._client.post("/v1/ingest/bulk", headers=headers, content=body)
        resp.raise_for_status()
        ack = resp.json()
        self.buffer.trim(ack["last_seq"])           # only trim what the server accepted
        return ack["accepted"]

    def run_forever(self) -> None:
        self.ensure_enrolled()
        while True:
            try:
                sent = self.flush_once()
                if sent:
                    log.info("shipped %d events (%d pending)", sent, self.buffer.pending())
            except Exception as exc:                # keep buffering on any failure
                log.warning("flush failed (buffering): %s", exc)
            time.sleep(self.cfg.flush_interval_s)