ctms/backend/app/services/security_access_log_writer.py

"""底层安全访问日志后台写入器"""

from __future__ import annotations

import asyncio
import logging
import uuid
from datetime import datetime, timezone
from typing import Any

from app.db.session import SessionLocal
from app.models.security_access_log import SecurityAccessLog
from app.services.ip_location import resolve_ip_location
from app.services.security_events import classify_security_event

logger = logging.getLogger("ctms.security_access_log_writer")

BATCH_SIZE = 100
FLUSH_INTERVAL = 3.0
QUEUE_MAX_SIZE = 20000
WRITE_ATTEMPTS = 2
_QUEUE_STOP = object()


class SecurityAccessLogWriter:
    def __init__(self) -> None:
        self._queue: asyncio.Queue[dict | object] = asyncio.Queue(maxsize=QUEUE_MAX_SIZE)
        self._task: asyncio.Task | None = None
        self._stopping = False
        self._started_at: datetime | None = None
        self._last_success_at: datetime | None = None
        self._last_error_at: datetime | None = None
        self._last_error_type: str | None = None
        self._accepted_entries = 0
        self._written_entries = 0
        self._dropped_entries = 0
        self._failed_batches = 0
        self._failed_entries = 0
        self._retry_count = 0

    def enqueue(self, entry: dict) -> None:
        if self._stopping:
            self._dropped_entries += 1
            logger.warning("Security access log writer is stopping, dropping entry")
            return
        try:
            self._queue.put_nowait(entry)
            self._accepted_entries += 1
        except asyncio.QueueFull:
            self._dropped_entries += 1
            logger.warning("Security access log queue full, dropping entry")

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        self._stopping = False
        self._started_at = datetime.now(timezone.utc)
        self._task = asyncio.create_task(self._flush_loop())
        logger.info("SecurityAccessLogWriter started")

    async def stop(self) -> None:
        if self._task and not self._task.done():
            self._stopping = True
            await self._queue.put(_QUEUE_STOP)
            await self._task
        self._task = None
        logger.info("SecurityAccessLogWriter stopped")

    async def _flush_loop(self) -> None:
        while True:
            batch, should_stop = await self._collect_batch()
            if batch:
                await self._write_batch(batch)
            if should_stop:
                break

    async def _collect_batch(self) -> tuple[list[dict], bool]:
        batch: list[dict] = []
        try:
            first = await asyncio.wait_for(self._queue.get(), timeout=FLUSH_INTERVAL)
            if first is _QUEUE_STOP:
                return batch, True
            batch.append(first)
        except asyncio.TimeoutError:
            return batch, False

        while len(batch) < BATCH_SIZE:
            try:
                item = self._queue.get_nowait()
                if item is _QUEUE_STOP:
                    return batch, True
                batch.append(item)
            except asyncio.QueueEmpty:
                break
        return batch, False

    async def _write_batch(self, batch: list[dict]) -> bool:
        for attempt in range(WRITE_ATTEMPTS):
            try:
                async with SessionLocal() as session:
                    for entry in batch:
                        classification = classify_security_event(
                            path=entry["path"],
                            status_code=entry["status_code"],
                            auth_status=entry["auth_status"],
                            ip_location=resolve_ip_location(entry.get("client_ip")),
                        )
                        session.add(
                            SecurityAccessLog(
                                id=uuid.uuid4(),
                                method=entry["method"],
                                path=entry["path"],
                                status_code=entry["status_code"],
                                elapsed_ms=entry["elapsed_ms"],
                                client_ip=entry.get("client_ip"),
                                user_agent=entry.get("user_agent"),
                                client_type=entry.get("client_type"),
                                client_version=entry.get("client_version"),
                                client_platform=entry.get("client_platform"),
                                build_channel=entry.get("build_channel"),
                                build_commit=entry.get("build_commit"),
                                request_headers=entry.get("request_headers"),
                                request_snapshot=entry.get("request_snapshot"),
                                request_id=entry.get("request_id"),
                                category=classification["category"],
                                severity=classification["severity"],
                                auth_status=entry["auth_status"],
                                user_identifier=entry.get("user_identifier"),
                                created_at=entry.get("created_at", datetime.now(timezone.utc)),
                            )
                        )
                    await session.commit()
            except Exception as exc:
                self._last_error_at = datetime.now(timezone.utc)
                self._last_error_type = type(exc).__name__
                if attempt + 1 < WRITE_ATTEMPTS:
                    self._retry_count += 1
                    logger.warning(
                        "Security access log batch write failed; retrying (%d entries)",
                        len(batch),
                    )
                    await asyncio.sleep(0)
                    continue
                self._failed_batches += 1
                self._failed_entries += len(batch)
                logger.exception(
                    "Failed to write security access log batch after retries (%d entries)",
                    len(batch),
                )
                return False
            self._written_entries += len(batch)
            self._last_success_at = datetime.now(timezone.utc)
            return True
        return False

    def stats(self) -> dict[str, Any]:
        return {
            "running": bool(self._task and not self._task.done()),
            "stopping": self._stopping,
            "queue_size": self._queue.qsize(),
            "queue_capacity": self._queue.maxsize,
            "accepted_entries": self._accepted_entries,
            "written_entries": self._written_entries,
            "dropped_entries": self._dropped_entries,
            "failed_batches": self._failed_batches,
            "failed_entries": self._failed_entries,
            "retry_count": self._retry_count,
            "started_at": self._started_at.isoformat() if self._started_at else None,
            "last_success_at": self._last_success_at.isoformat() if self._last_success_at else None,
            "last_error_at": self._last_error_at.isoformat() if self._last_error_at else None,
            "last_error_type": self._last_error_type,
        }


_writer: SecurityAccessLogWriter | None = None


def get_security_log_writer() -> SecurityAccessLogWriter | None:
    return _writer


async def start_security_log_writer() -> SecurityAccessLogWriter:
    global _writer
    _writer = SecurityAccessLogWriter()
    await _writer.start()
    return _writer


async def stop_security_log_writer() -> None:
    global _writer
    if _writer:
        await _writer.stop()
        _writer = None