Files
AWatch-rus/aw-server/dlp-policy-engine/policy_storage.py
T

368 lines
14 KiB
Python

from __future__ import annotations
import hashlib
import json
import sqlite3
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator
def utc_now() -> str:
return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
def canonical_policy_json(policy: dict[str, Any]) -> str:
return json.dumps(policy, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def checksum_policy(policy: dict[str, Any]) -> str:
return hashlib.sha256(canonical_policy_json(policy).encode("utf-8")).hexdigest()
class PolicyStorage:
def __init__(self, db_path: str) -> None:
self.db_path = Path(db_path)
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self._init_schema()
@contextmanager
def connect(self) -> Iterator[sqlite3.Connection]:
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
try:
yield conn
conn.commit()
finally:
conn.close()
def _init_schema(self) -> None:
with self.connect() as conn:
conn.executescript(
"""
PRAGMA journal_mode=WAL;
CREATE TABLE IF NOT EXISTS policies (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE,
description TEXT,
status TEXT NOT NULL DEFAULT 'draft',
is_active INTEGER NOT NULL DEFAULT 0,
current_version INTEGER NOT NULL DEFAULT 1,
checksum TEXT NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS policy_versions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
policy_id INTEGER NOT NULL,
version INTEGER NOT NULL,
policy_json TEXT NOT NULL,
checksum TEXT NOT NULL,
created_at TEXT NOT NULL,
created_by TEXT,
rollback_of_version INTEGER,
FOREIGN KEY(policy_id) REFERENCES policies(id),
UNIQUE(policy_id, version)
);
CREATE TABLE IF NOT EXISTS policy_audit (
id INTEGER PRIMARY KEY AUTOINCREMENT,
policy_id INTEGER,
action TEXT NOT NULL,
actor TEXT,
comment TEXT,
details_json TEXT,
created_at TEXT NOT NULL,
FOREIGN KEY(policy_id) REFERENCES policies(id)
);
CREATE INDEX IF NOT EXISTS idx_policies_active ON policies(is_active);
CREATE INDEX IF NOT EXISTS idx_policy_versions_policy ON policy_versions(policy_id, version DESC);
CREATE INDEX IF NOT EXISTS idx_policy_audit_policy ON policy_audit(policy_id, id DESC);
"""
)
cols = [r["name"] for r in conn.execute("PRAGMA table_info(policies)").fetchall()]
if "status" not in cols:
conn.execute("ALTER TABLE policies ADD COLUMN status TEXT NOT NULL DEFAULT 'draft'")
def _audit(
self,
conn: sqlite3.Connection,
policy_id: int | None,
action: str,
actor: str | None,
comment: str | None = None,
details: dict[str, Any] | None = None,
) -> None:
conn.execute(
"""
INSERT INTO policy_audit(policy_id, action, actor, comment, details_json, created_at)
VALUES(?, ?, ?, ?, ?, ?)
""",
(
policy_id,
action,
actor,
comment,
canonical_policy_json(details) if details is not None else None,
utc_now(),
),
)
def list_policies(self) -> list[dict[str, Any]]:
with self.connect() as conn:
rows = conn.execute(
"""
SELECT id, name, description, status, is_active, current_version, checksum, created_at, updated_at
FROM policies
ORDER BY is_active DESC, updated_at DESC, id DESC
"""
).fetchall()
return [dict(row) for row in rows]
def get_policy(self, policy_id: int) -> dict[str, Any] | None:
with self.connect() as conn:
policy_row = conn.execute(
"""
SELECT id, name, description, status, is_active, current_version, checksum, created_at, updated_at
FROM policies
WHERE id = ?
""",
(policy_id,),
).fetchone()
if not policy_row:
return None
version_row = conn.execute(
"""
SELECT version, policy_json, checksum, created_at, created_by
FROM policy_versions
WHERE policy_id = ? AND version = ?
""",
(policy_id, policy_row["current_version"]),
).fetchone()
if not version_row:
return None
result = dict(policy_row)
result["policy"] = json.loads(version_row["policy_json"])
result["version_created_at"] = version_row["created_at"]
result["version_created_by"] = version_row["created_by"]
return result
def get_active_policy(self) -> dict[str, Any] | None:
with self.connect() as conn:
row = conn.execute("SELECT id FROM policies WHERE is_active = 1 ORDER BY updated_at DESC LIMIT 1").fetchone()
if not row:
return None
return self.get_policy(int(row["id"]))
def create_policy(self, name: str, description: str | None, policy: dict[str, Any], activate: bool, actor: str | None) -> dict[str, Any]:
checksum = checksum_policy(policy)
now = utc_now()
status = "deployed" if activate else "draft"
policy_json = canonical_policy_json(policy)
with self.connect() as conn:
if activate:
conn.execute("UPDATE policies SET is_active = 0")
cursor = conn.execute(
"""
INSERT INTO policies(name, description, status, is_active, current_version, checksum, created_at, updated_at)
VALUES(?, ?, ?, ?, 1, ?, ?, ?)
""",
(name, description, status, 1 if activate else 0, checksum, now, now),
)
policy_id = int(cursor.lastrowid)
conn.execute(
"""
INSERT INTO policy_versions(policy_id, version, policy_json, checksum, created_at, created_by, rollback_of_version)
VALUES(?, 1, ?, ?, ?, ?, NULL)
""",
(policy_id, policy_json, checksum, now, actor),
)
self._audit(conn, policy_id, "create", actor, details={"activate": activate, "status": status})
return self.get_policy(policy_id) # type: ignore[return-value]
def update_policy(
self,
policy_id: int,
name: str | None,
description: str | None,
policy: dict[str, Any] | None,
activate: bool,
actor: str | None,
) -> dict[str, Any] | None:
current = self.get_policy(policy_id)
if not current:
return None
with self.connect() as conn:
new_name = name if name is not None else current["name"]
new_description = description if description is not None else current["description"]
new_version = int(current["current_version"])
new_checksum = current["checksum"]
new_status = current.get("status", "draft")
if policy is not None:
new_version += 1
new_checksum = checksum_policy(policy)
new_status = "draft"
policy_json = canonical_policy_json(policy)
conn.execute(
"""
INSERT INTO policy_versions(policy_id, version, policy_json, checksum, created_at, created_by, rollback_of_version)
VALUES(?, ?, ?, ?, ?, ?, NULL)
""",
(policy_id, new_version, policy_json, new_checksum, utc_now(), actor),
)
if activate:
conn.execute("UPDATE policies SET is_active = 0")
new_status = "deployed"
conn.execute(
"""
UPDATE policies
SET name = ?, description = ?, status = ?, is_active = ?, current_version = ?, checksum = ?, updated_at = ?
WHERE id = ?
""",
(
new_name,
new_description,
new_status,
1 if activate else current["is_active"],
new_version,
new_checksum,
utc_now(),
policy_id,
),
)
self._audit(conn, policy_id, "update", actor, details={"activate": activate, "status": new_status})
return self.get_policy(policy_id)
def activate_policy(self, policy_id: int, actor: str | None) -> dict[str, Any] | None:
current = self.get_policy(policy_id)
if not current:
return None
if current.get("status") != "approved":
raise ValueError("policy must be approved before deploy")
with self.connect() as conn:
conn.execute("UPDATE policies SET is_active = 0")
conn.execute(
"UPDATE policies SET status = 'deployed', is_active = 1, updated_at = ? WHERE id = ?",
(utc_now(), policy_id),
)
self._audit(conn, policy_id, "deploy", actor)
return self.get_policy(policy_id)
def rollback_active_policy(self, actor: str | None) -> dict[str, Any] | None:
active = self.get_active_policy()
if not active:
return None
with self.connect() as conn:
rows = conn.execute(
"""
SELECT version, policy_json
FROM policy_versions
WHERE policy_id = ?
ORDER BY version DESC
LIMIT 2
""",
(active["id"],),
).fetchall()
if len(rows) < 2:
return active
previous_version = int(rows[1]["version"])
previous_policy = json.loads(rows[1]["policy_json"])
rollback_version = int(active["current_version"]) + 1
rollback_checksum = checksum_policy(previous_policy)
now = utc_now()
conn.execute(
"""
INSERT INTO policy_versions(policy_id, version, policy_json, checksum, created_at, created_by, rollback_of_version)
VALUES(?, ?, ?, ?, ?, ?, ?)
""",
(
active["id"],
rollback_version,
canonical_policy_json(previous_policy),
rollback_checksum,
now,
actor,
previous_version,
),
)
conn.execute(
"""
UPDATE policies
SET status = 'draft', current_version = ?, checksum = ?, updated_at = ?
WHERE id = ?
""",
(rollback_version, rollback_checksum, now, active["id"]),
)
self._audit(conn, int(active["id"]), "rollback", actor, details={"rollback_to": previous_version})
return self.get_policy(int(active["id"]))
def delete_policy(self, policy_id: int) -> bool:
current = self.get_policy(policy_id)
if not current:
return False
if current["is_active"]:
raise ValueError("cannot delete active policy")
with self.connect() as conn:
self._audit(conn, policy_id, "delete", None)
conn.execute("DELETE FROM policy_versions WHERE policy_id = ?", (policy_id,))
conn.execute("DELETE FROM policies WHERE id = ?", (policy_id,))
return True
def set_policy_status(self, policy_id: int, status: str, actor: str | None, comment: str | None = None) -> dict[str, Any] | None:
current = self.get_policy(policy_id)
if not current:
return None
allowed = {"draft", "pending_approval", "approved", "deployed"}
if status not in allowed:
raise ValueError(f"unsupported status: {status}")
with self.connect() as conn:
conn.execute(
"UPDATE policies SET status = ?, updated_at = ? WHERE id = ?",
(status, utc_now(), policy_id),
)
self._audit(conn, policy_id, "status_change", actor, comment=comment, details={"status": status})
return self.get_policy(policy_id)
def list_audit(self, policy_id: int | None = None, limit: int = 200) -> list[dict[str, Any]]:
query = """
SELECT id, policy_id, action, actor, comment, details_json, created_at
FROM policy_audit
"""
params: tuple[Any, ...]
if policy_id is None:
query += " ORDER BY id DESC LIMIT ?"
params = (limit,)
else:
query += " WHERE policy_id = ? ORDER BY id DESC LIMIT ?"
params = (policy_id, limit)
with self.connect() as conn:
rows = conn.execute(query, params).fetchall()
items: list[dict[str, Any]] = []
for row in rows:
d = dict(row)
if d.get("details_json"):
try:
d["details"] = json.loads(d["details_json"])
except Exception:
d["details"] = None
else:
d["details"] = None
d.pop("details_json", None)
items.append(d)
return items