devplacepy/devplacepy/services/audit/store.py

# retoor <retoor@molodetz.nl>

import logging
from datetime import datetime, timezone
from typing import Optional

from devplacepy.database import db, get_table
from devplacepy.utils import generate_uid

logger = logging.getLogger(__name__)

AUDIT_TABLE = "audit_log"
LINKS_TABLE = "audit_log_links"

EVENT_COLUMNS = (
    "uid",
    "created_at",
    "event_key",
    "category",
    "actor_kind",
    "actor_uid",
    "actor_username",
    "actor_role",
    "origin",
    "via_agent",
    "request_method",
    "request_path",
    "actor_ip",
    "user_agent",
    "target_type",
    "target_uid",
    "target_label",
    "old_value",
    "new_value",
    "summary",
    "metadata",
    "result",
)

LINK_COLUMNS = (
    "uid",
    "audit_uid",
    "relation",
    "object_type",
    "object_uid",
    "object_label",
    "created_at",
)


def now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _empty_event() -> dict:
    row: dict = {column: None for column in EVENT_COLUMNS}
    row["via_agent"] = 0
    return row


def _empty_link() -> dict:
    return {column: None for column in LINK_COLUMNS}


def ensure_tables() -> None:
    try:
        if AUDIT_TABLE not in db.tables:
            sentinel = _empty_event()
            sentinel["uid"] = "__audit_sentinel__"
            sentinel["created_at"] = now()
            sentinel["event_key"] = "__sentinel__"
            sentinel["category"] = "system"
            sentinel["actor_kind"] = "system"
            sentinel["actor_role"] = "system"
            sentinel["origin"] = "system"
            sentinel["result"] = "success"
            table = get_table(AUDIT_TABLE)
            table.insert(sentinel)
            table.delete(uid="__audit_sentinel__")
        if LINKS_TABLE not in db.tables:
            sentinel = _empty_link()
            sentinel["uid"] = "__audit_link_sentinel__"
            sentinel["audit_uid"] = "__audit_sentinel__"
            sentinel["relation"] = "actor"
            sentinel["object_type"] = "user"
            sentinel["object_uid"] = "__sentinel__"
            sentinel["created_at"] = now()
            table = get_table(LINKS_TABLE)
            table.insert(sentinel)
            table.delete(uid="__audit_link_sentinel__")
    except Exception as exc:
        logger.warning("audit ensure_tables failed: %s", exc)


def insert_event(row: dict) -> str:
    record = _empty_event()
    record.update({key: value for key, value in row.items() if key in EVENT_COLUMNS})
    if not record.get("uid"):
        record["uid"] = generate_uid()
    if not record.get("created_at"):
        record["created_at"] = now()
    if record.get("via_agent") is None:
        record["via_agent"] = 0
    get_table(AUDIT_TABLE).insert(record, ensure=False)
    return record["uid"]


def insert_links(audit_uid: str, links: list[dict]) -> int:
    if not links:
        return 0
    stamp = now()
    records = []
    for link in links:
        record = _empty_link()
        record.update(
            {key: value for key, value in link.items() if key in LINK_COLUMNS}
        )
        record["uid"] = generate_uid()
        record["audit_uid"] = audit_uid
        record["created_at"] = stamp
        if not record.get("object_uid"):
            continue
        records.append(record)
    if records:
        get_table(LINKS_TABLE).insert_many(records, ensure=False)
    return len(records)


def get_event(uid: str) -> Optional[dict]:
    if AUDIT_TABLE not in db.tables:
        return None
    return get_table(AUDIT_TABLE).find_one(uid=uid)


def get_links(audit_uid: str) -> list[dict]:
    if LINKS_TABLE not in db.tables:
        return []
    return list(
        get_table(LINKS_TABLE).find(audit_uid=audit_uid, order_by=["created_at"])
    )


def sweep(cutoff: str, batch_size: int = 5000, max_batches: int = 100) -> tuple[int, int]:
    if AUDIT_TABLE not in db.tables:
        return (0, 0)
    removed_links = 0
    removed_events = 0
    for _ in range(max_batches):
        rows = list(
            db.query(
                f"SELECT uid FROM {AUDIT_TABLE} WHERE created_at < :cutoff "
                f"ORDER BY created_at LIMIT :limit",
                cutoff=cutoff,
                limit=batch_size,
            )
        )
        uids = [row["uid"] for row in rows]
        if not uids:
            break
        placeholders, params = _in_clause(uids)
        with db:
            if LINKS_TABLE in db.tables:
                link_rows = list(
                    db.query(
                        f"SELECT COUNT(*) AS c FROM {LINKS_TABLE} "
                        f"WHERE audit_uid IN ({placeholders})",
                        **params,
                    )
                )
                removed_links += int(link_rows[0]["c"]) if link_rows else 0
                db.query(
                    f"DELETE FROM {LINKS_TABLE} WHERE audit_uid IN ({placeholders})",
                    **params,
                )
            db.query(
                f"DELETE FROM {AUDIT_TABLE} WHERE uid IN ({placeholders})", **params
            )
        removed_events += len(uids)
        if len(uids) < batch_size:
            break
    return (removed_links, removed_events)


def _in_clause(uids: list[str], prefix: str = "u") -> tuple[str, dict]:
    placeholders = ", ".join(f":{prefix}{index}" for index in range(len(uids)))
    params = {f"{prefix}{index}": uid for index, uid in enumerate(uids)}
    return placeholders, params