feat: migrate audit_logger to PyMySQL MariaDB

Replace SQLite AuditLogger with MariaDB-backed version using PyMySQL. Keep identical public interface (log, query, count, stats_today). Constructor accepts pre-existing connection (tests) or env-var credentials (production). addon.py reads KVM_MITM_DB_HOST/USER/PASS/NAME env vars; guards against init failure. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-04 02:42:54 +00:00
parent dc09306655
commit a75c2696cc
3 changed files with 208 additions and 143 deletions
@@ -11,6 +11,7 @@ import asyncio
 import hashlib
 import json
 import logging
+import os
 import sys
 from pathlib import Path

@@ -42,7 +43,16 @@ _RETRY_BACKOFF = (0.5, 1.0, 2.0)

 _DOMAINS_FILE = Path(__file__).parent / "ai_domains.txt"
 _STATE_FILE = Path("/var/lib/kvm-privacy/state.json")
-_audit = AuditLogger()
+try:
+    _audit: AuditLogger | None = AuditLogger(
+        host=os.environ.get("KVM_MITM_DB_HOST", "localhost"),
+        user=os.environ.get("KVM_MITM_DB_USER", "kvm_mitm"),
+        password=os.environ.get("KVM_MITM_DB_PASS", ""),
+        database=os.environ.get("KVM_MITM_DB_NAME", "kvm"),
+    )
+except Exception as _e:  # pragma: no cover
+    logger.warning("[privacy-gw] AuditLogger init failed: %s — audit disabled", _e)
+    _audit = None
 _KVM_AUDIT_URL = "http://127.0.0.1:8080/api/internal/privacy-event"
 _http_client: httpx.AsyncClient | None = None

@@ -124,20 +134,22 @@ class PrivacyGatewayAddon:
                result = await scan_and_redact(f.filename, f.data)
            except Exception as e:
                ctx.log.warn(f"[privacy-gw] scan failed {f.filename}: {e}")
-                _audit.log(
-                    host, {}, "allow", f.data, total_files,
-                    client_ip=client_ip, request_url=request_url,
-                    filename=f.filename, file_size=len(f.data),
-                )
+                if _audit is not None:
+                    _audit.log(
+                        host, {}, "allow", f.data, total_files,
+                        client_ip=client_ip, request_url=request_url,
+                        filename=f.filename, file_size=len(f.data),
+                    )
                continue

            if mode == "audit":
                # Log only — never modify the request
-                _audit.log(
-                    host, result.pii_types, "allow", f.data, total_files,
-                    client_ip=client_ip, request_url=request_url,
-                    filename=f.filename, file_size=len(f.data),
-                )
+                if _audit is not None:
+                    _audit.log(
+                        host, result.pii_types, "allow", f.data, total_files,
+                        client_ip=client_ip, request_url=request_url,
+                        filename=f.filename, file_size=len(f.data),
+                    )
                if result.pii_found:
                    ctx.log.info(f"[privacy-gw] audit: {f.filename} has PII {result.pii_types}")

@@ -149,11 +161,12 @@ class PrivacyGatewayAddon:
                else:
                    action = "allow"

-                _audit.log(
-                    host, result.pii_types, action, f.data, total_files,
-                    client_ip=client_ip, request_url=request_url,
-                    filename=f.filename, file_size=len(f.data),
-                )
+                if _audit is not None:
+                    _audit.log(
+                        host, result.pii_types, action, f.data, total_files,
+                        client_ip=client_ip, request_url=request_url,
+                        filename=f.filename, file_size=len(f.data),
+                    )
                await _post_to_kvm_audit(
                    domain=host, pii_types=result.pii_types, action=action,
                    doc_hash=hashlib.sha256(f.data).hexdigest(), file_count=total_files,
@@ -1,82 +1,76 @@
-"""审计日志 —— SQLite append-only，只记 PII 类型/哈希，不存原文。
-
-作为本地冗余备份，即使 KVM MySQL 不可达也不丢数据。
-"""
+"""审计日志 —— MariaDB kvm.privacy_audit_log，只记 PII 类型/哈希，不存原文。"""
 from __future__ import annotations

 import hashlib
 import json
-import sqlite3
+import os
 import threading
-from datetime import datetime
-from pathlib import Path
+from datetime import datetime, date
+
+import pymysql
+import pymysql.cursors
+

-_DB_PATH = Path("/var/lib/kvm-privacy/audit.db")
 _SCHEMA = """
-CREATE TABLE IF NOT EXISTS audit_log (
-    id        INTEGER PRIMARY KEY AUTOINCREMENT,
-    ts        TEXT    NOT NULL,
-    domain    TEXT    NOT NULL,
-    pii_types TEXT,           -- JSON {"id_card":2,"phone":1}
-    action    TEXT    NOT NULL, -- auto_redact | blocked | bypass
-    doc_hash  TEXT,           -- sha256(original_bytes) — 不存原文
-    file_count INTEGER DEFAULT 0
-);
-CREATE INDEX IF NOT EXISTS idx_ts     ON audit_log(ts);
-CREATE INDEX IF NOT EXISTS idx_domain ON audit_log(domain);
+CREATE TABLE IF NOT EXISTS privacy_audit_log (
+    id                 INT AUTO_INCREMENT PRIMARY KEY,
+    ts                 DATETIME(3) NOT NULL DEFAULT NOW(3),
+    domain             VARCHAR(255) NOT NULL,
+    pii_types          JSON,
+    action             VARCHAR(20)  NOT NULL,
+    doc_hash           VARCHAR(64),
+    file_count         INT          DEFAULT 0,
+    client_ip          VARCHAR(45)  DEFAULT '',
+    request_url        VARCHAR(2048) DEFAULT '',
+    filename           VARCHAR(255) DEFAULT '',
+    file_size          INT          DEFAULT 0,
+    redacted_pii_count INT          DEFAULT 0,
+    INDEX idx_ts     (ts),
+    INDEX idx_domain (domain),
+    INDEX idx_action (action)
+)
 """

-# 新增字段（幂等迁移，已有列时忽略）
-_MIGRATIONS = [
-    "ALTER TABLE audit_log ADD COLUMN client_ip TEXT DEFAULT ''",
-    "ALTER TABLE audit_log ADD COLUMN request_url TEXT DEFAULT ''",
-    "ALTER TABLE audit_log ADD COLUMN filename TEXT DEFAULT ''",
-    "ALTER TABLE audit_log ADD COLUMN file_size INTEGER DEFAULT 0",
-    "ALTER TABLE audit_log ADD COLUMN redacted_pii_count INTEGER DEFAULT 0",
-]
-

 class AuditLogger:
-    def __init__(self, db_path: Path = _DB_PATH) -> None:
-        db_path.parent.mkdir(parents=True, exist_ok=True)
-        self._db_path = str(db_path)
+    """MariaDB-backed audit logger for the privacy gateway."""
+
+    def __init__(
+        self,
+        conn: "pymysql.connections.Connection | None" = None,
+        *,
+        host: str = "localhost",
+        user: str = "kvm_mitm",
+        password: str = "",
+        database: str = "kvm",
+    ) -> None:
        self._lock = threading.Lock()
-        # Persistent connection with WAL mode for concurrent read/write
-        self._conn = self._create_conn()
-        self._conn.executescript(_SCHEMA)
-        for stmt in _MIGRATIONS:
-            try:
-                self._conn.execute(stmt)
-            except sqlite3.OperationalError:
-                pass  # column already exists
-        # Additional indexes for dashboard queries
-        for idx_stmt in (
-            "CREATE INDEX IF NOT EXISTS idx_action    ON audit_log(action)",
-            "CREATE INDEX IF NOT EXISTS idx_client_ip ON audit_log(client_ip)",
-        ):
-            self._conn.execute(idx_stmt)
+        if conn is not None:
+            self._conn = conn
+            self._owns_conn = False
+        else:
+            self._conn = pymysql.connect(
+                host=host, user=user, password=password,
+                database=database, charset="utf8mb4", autocommit=False,
+            )
+            self._owns_conn = True
+        self._ensure_schema()
+
+    def _ensure_schema(self) -> None:
+        with self._conn.cursor() as cur:
+            cur.execute(_SCHEMA)
        self._conn.commit()

-    def _create_conn(self) -> sqlite3.Connection:
-        conn = sqlite3.connect(self._db_path, check_same_thread=False)
-        conn.row_factory = sqlite3.Row
-        conn.execute("PRAGMA journal_mode=WAL")
-        return conn
-
-    def _get_conn(self) -> sqlite3.Connection:
-        """Return persistent connection, reconnecting if closed."""
-        try:
-            self._conn.execute("SELECT 1")
-        except (sqlite3.ProgrammingError, sqlite3.OperationalError):
-            self._conn = self._create_conn()
-        return self._conn
+    def _cursor(self) -> "pymysql.cursors.DictCursor":
+        self._conn.ping(reconnect=True)
+        return self._conn.cursor(pymysql.cursors.DictCursor)

    def log(
        self,
        domain: str,
-        pii_types: dict[str, int],
+        pii_types: "dict[str, int]",
        action: str,
-        raw_bytes: bytes | None = None,
+        raw_bytes: "bytes | None" = None,
        file_count: int = 0,
        *,
        client_ip: str = "",
@@ -86,98 +80,83 @@ class AuditLogger:
    ) -> None:
        doc_hash = hashlib.sha256(raw_bytes).hexdigest() if raw_bytes else None
        redacted_pii_count = sum(pii_types.values()) if pii_types else 0
-        with self._lock:
-            conn = self._get_conn()
-            conn.execute(
-                "INSERT INTO audit_log"
+        with self._lock, self._cursor() as cur:
+            cur.execute(
+                "INSERT INTO privacy_audit_log "
                "(ts, domain, pii_types, action, doc_hash, file_count,"
                " client_ip, request_url, filename, file_size, redacted_pii_count) "
-                "VALUES (?,?,?,?,?,?,?,?,?,?,?)",
+                "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)",
                (
-                    datetime.utcnow().isoformat(),
+                    datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
                    domain,
                    json.dumps(pii_types) if pii_types else None,
-                    action,
-                    doc_hash,
-                    file_count,
-                    client_ip,
-                    request_url,
-                    filename,
-                    file_size,
+                    action, doc_hash, file_count,
+                    client_ip, request_url, filename, file_size,
                    redacted_pii_count,
                ),
            )
-            conn.commit()
+            self._conn.commit()

    def query(
-        self, limit: int = 50, offset: int = 0, domain: str | None = None
-    ) -> list[dict]:
-        sql = "SELECT * FROM audit_log"
+        self, limit: int = 50, offset: int = 0, domain: "str | None" = None
+    ) -> "list[dict]":
+        sql = "SELECT * FROM privacy_audit_log"
        params: list = []
        if domain:
-            sql += " WHERE domain = ?"
+            sql += " WHERE domain = %s"
            params.append(domain)
-        sql += " ORDER BY id DESC LIMIT ? OFFSET ?"
+        sql += " ORDER BY id DESC LIMIT %s OFFSET %s"
        params.extend([limit, offset])
-        with self._lock:
-            rows = self._get_conn().execute(sql, params).fetchall()
-        return [dict(r) for r in rows]
+        with self._lock, self._cursor() as cur:
+            cur.execute(sql, params)
+            return list(cur.fetchall())

-    def count(self, domain: str | None = None) -> int:
-        """Return total number of audit log entries."""
-        sql = "SELECT COUNT(*) FROM audit_log"
+    def count(self, domain: "str | None" = None) -> int:
+        sql = "SELECT COUNT(*) AS n FROM privacy_audit_log"
        params: list = []
        if domain:
-            sql += " WHERE domain = ?"
+            sql += " WHERE domain = %s"
            params.append(domain)
-        with self._lock:
-            row = self._get_conn().execute(sql, params).fetchone()
-        return row[0] if row else 0
+        with self._lock, self._cursor() as cur:
+            cur.execute(sql, params)
+            row = cur.fetchone()
+        return row["n"] if row else 0

    def stats_today(self) -> dict:
        """返回今日统计：总量 + 按类型/动作/域名 分组。"""
-        today = datetime.utcnow().date().isoformat()
-        with self._lock:
-            conn = self._get_conn()
-            # 总量
-            row = conn.execute(
+        today = date.today().strftime("%Y-%m-%d")
+        with self._lock, self._cursor() as cur:
+            cur.execute(
                "SELECT COUNT(*) AS requests, "
                "COALESCE(SUM(file_count),0) AS files, "
                "COALESCE(SUM(redacted_pii_count),0) AS pii_total "
-                "FROM audit_log WHERE ts >= ?",
+                "FROM privacy_audit_log WHERE ts >= %s",
                (today,),
-            ).fetchone()
-            summary = {
-                "requests": row["requests"] or 0,
-                "files": row["files"] or 0,
-                "pii_total": row["pii_total"] or 0,
-            }
+            )
+            row = cur.fetchone()
+            summary = {k: int(row[k] or 0) for k in ("requests", "files", "pii_total")}

-            # 按动作分组
-            by_action: dict[str, int] = {}
-            for r in conn.execute(
-                "SELECT action, COUNT(*) AS cnt FROM audit_log "
-                "WHERE ts >= ? GROUP BY action",
+            cur.execute(
+                "SELECT action, COUNT(*) AS cnt FROM privacy_audit_log "
+                "WHERE ts >= %s GROUP BY action",
                (today,),
-            ).fetchall():
-                by_action[r["action"]] = r["cnt"]
+            )
+            by_action = {r["action"]: r["cnt"] for r in cur.fetchall()}

-            # 按域名 Top-10
-            by_domain: dict[str, int] = {}
-            for r in conn.execute(
-                "SELECT domain, COUNT(*) AS cnt FROM audit_log "
-                "WHERE ts >= ? GROUP BY domain ORDER BY cnt DESC LIMIT 10",
+            cur.execute(
+                "SELECT domain, COUNT(*) AS cnt FROM privacy_audit_log "
+                "WHERE ts >= %s GROUP BY domain ORDER BY cnt DESC LIMIT 10",
                (today,),
-            ).fetchall():
-                by_domain[r["domain"]] = r["cnt"]
+            )
+            by_domain = {r["domain"]: r["cnt"] for r in cur.fetchall()}

-            # 按 PII 类型聚合（需解析 JSON）
-            pii_by_type: dict[str, int] = {}
-            for r in conn.execute(
-                "SELECT pii_types FROM audit_log "
-                "WHERE ts >= ? AND pii_types IS NOT NULL",
+            cur.execute(
+                "SELECT pii_types FROM privacy_audit_log "
+                "WHERE ts >= %s AND pii_types IS NOT NULL",
                (today,),
-            ).fetchall():
+            )
+            pii_by_type: "dict[str, int]" = {}
+            for r in cur.fetchall():
                try:
                    types = json.loads(r["pii_types"])
                    for k, v in types.items():
@@ -185,9 +164,9 @@ class AuditLogger:
                except (json.JSONDecodeError, TypeError):
                    pass

-        return {
-            **summary,
-            "by_action": by_action,
-            "by_domain": by_domain,
-            "pii_by_type": pii_by_type,
-        }
+        return {**summary, "by_action": by_action,
+                "by_domain": by_domain, "pii_by_type": pii_by_type}
+
+    def close(self) -> None:
+        if self._owns_conn:
+            self._conn.close()
@@ -0,0 +1,73 @@
+"""Integration tests for AuditLogger with MariaDB kvm_test."""
+import json
+import pytest
+import pymysql
+
+from privacy_gateway.audit_logger import AuditLogger
+
+DB_CFG = dict(host="localhost", user="kvm_mitm", password="test_mitm_pass",
+              database="kvm_test", charset="utf8mb4")
+
+
+@pytest.fixture()
+def db_conn():
+    conn = pymysql.connect(**DB_CFG)
+    # Create table fresh
+    with conn.cursor() as cur:
+        cur.execute("DROP TABLE IF EXISTS privacy_audit_log")
+        cur.execute("""
+            CREATE TABLE privacy_audit_log (
+                id INT AUTO_INCREMENT PRIMARY KEY,
+                ts DATETIME(3) NOT NULL DEFAULT NOW(3),
+                domain VARCHAR(255) NOT NULL,
+                pii_types JSON,
+                action VARCHAR(20) NOT NULL,
+                doc_hash VARCHAR(64),
+                file_count INT DEFAULT 0,
+                client_ip VARCHAR(45) DEFAULT '',
+                request_url VARCHAR(2048) DEFAULT '',
+                filename VARCHAR(255) DEFAULT '',
+                file_size INT DEFAULT 0,
+                redacted_pii_count INT DEFAULT 0
+            )
+        """)
+    conn.commit()
+    yield conn
+    with conn.cursor() as cur:
+        cur.execute("DROP TABLE IF EXISTS privacy_audit_log")
+    conn.commit()
+    conn.close()
+
+
+@pytest.fixture()
+def logger(db_conn):
+    return AuditLogger(db_conn)
+
+
+def test_log_and_query(logger):
+    logger.log("example.com", {"id_card": 2}, "allow",
+               raw_bytes=b"data", file_count=1,
+               client_ip="1.2.3.4", filename="a.pdf", file_size=100)
+    rows = logger.query(limit=10)
+    assert len(rows) == 1
+    assert rows[0]["domain"] == "example.com"
+    assert rows[0]["action"] == "allow"
+    assert rows[0]["file_count"] == 1
+
+
+def test_count(logger):
+    logger.log("a.com", {}, "scan_failed")
+    logger.log("b.com", {"phone": 1}, "block")
+    assert logger.count() == 2
+    assert logger.count(domain="a.com") == 1
+
+
+def test_stats_today(logger):
+    logger.log("x.com", {"phone": 1, "id_card": 2}, "block",
+               file_count=3, file_size=500)
+    stats = logger.stats_today()
+    assert stats["requests"] == 1
+    assert stats["files"] == 3
+    assert stats["pii_total"] == 3
+    assert stats["by_action"]["block"] == 1
+    assert stats["by_domain"]["x.com"] == 1