proofpoint-quarantine-portal/worker.py

"""Background job queue for slow PPS quarantine actions.

The frontend fires an action and returns immediately; the actual PPS API calls
(which can be very slow, especially delete) run here on a daemon thread backed by
SQLite so queued work survives an app restart.

One JobQueue instance owns a single SQLite connection (guarded by a lock) shared
between Flask request threads and the worker thread.
"""

from __future__ import annotations

import json
import logging
import sqlite3
import threading
import time
import traceback
from datetime import datetime, timezone
from logging.handlers import RotatingFileHandler
from typing import Callable, Mapping

import pipeline
from pps_client import PPSClient, PPSError

log = logging.getLogger("pps.worker")

# Dedicated operations/audit log — one line per message acted on. Written to its own
# file (see _setup_ops_log) and NOT propagated to the root/stdout logger.
ops_log = logging.getLogger("pps.ops")
ops_log.propagate = False


def _setup_ops_log(path: str) -> None:
    """Attach a rotating file handler to the ops logger (idempotent)."""
    for h in ops_log.handlers:
        if isinstance(h, RotatingFileHandler) and getattr(h, "baseFilename", "").endswith(path.split("/")[-1]):
            return  # already configured
    handler = RotatingFileHandler(path, maxBytes=5_000_000, backupCount=5, encoding="utf-8")
    handler.setFormatter(logging.Formatter("%(asctime)s %(message)s"))
    ops_log.addHandler(handler)
    ops_log.setLevel(logging.INFO)

# Actions understood from the frontend.
VALID_ACTIONS = {"release", "report_release", "delete", "move"}

_POLL_SECONDS = 1.0


def _now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")


class JobQueue:
    def __init__(self, db_path: str,
                 pps_provider: Callable[[], PPSClient],
                 cfg_provider: Callable[[], Mapping],
                 ops_log_path: str = "worker.log"):
        # Providers, not frozen values: each job reads a fresh, consistent snapshot at
        # start (see AGENTS.md "snapshot-per-job"), so a live admin edit affects the
        # NEXT job, never a job mid-flight.
        self._pps = pps_provider
        self._cfg = cfg_provider
        self._lock = threading.Lock()
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._init_db()
        self._recover()
        self._stop = threading.Event()
        _setup_ops_log(ops_log_path)
        self.ops_log_path = ops_log_path

    # -------------------------------------------------------------- schema/setup

    def _init_db(self) -> None:
        with self._lock:
            # WAL + busy_timeout so the prefs store (a second connection to this same
            # file) never collides with per-chunk progress writes. WAL is a persistent
            # property of the file; busy_timeout must be set on every connection.
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.execute("PRAGMA busy_timeout=5000")
            self._conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    id         INTEGER PRIMARY KEY AUTOINCREMENT,
                    action     TEXT NOT NULL,
                    folder     TEXT NOT NULL,
                    user       TEXT,
                    extra      TEXT NOT NULL DEFAULT '{}',
                    status     TEXT NOT NULL DEFAULT 'pending',
                    processed  INTEGER NOT NULL DEFAULT 0,
                    total      INTEGER NOT NULL DEFAULT 0,
                    error      TEXT,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS job_items (
                    job_id    INTEGER NOT NULL,
                    localguid TEXT NOT NULL,
                    guid      TEXT,
                    folder    TEXT NOT NULL,
                    subject   TEXT,
                    sender    TEXT,
                    recipient TEXT,
                    FOREIGN KEY (job_id) REFERENCES jobs(id)
                );
                CREATE INDEX IF NOT EXISTS idx_job_items_folder ON job_items(folder);
                CREATE INDEX IF NOT EXISTS idx_job_items_job ON job_items(job_id);
                CREATE INDEX IF NOT EXISTS idx_jobs_user ON jobs(user);
                """
            )
            # Migrate older DBs that lack later-added columns.
            item_cols = {r["name"] for r in self._conn.execute("PRAGMA table_info(job_items)")}
            for col in ("subject", "sender", "recipient"):
                if col not in item_cols:
                    self._conn.execute(f"ALTER TABLE job_items ADD COLUMN {col} TEXT")
            job_cols = {r["name"] for r in self._conn.execute("PRAGMA table_info(jobs)")}
            if "user" not in job_cols:
                self._conn.execute("ALTER TABLE jobs ADD COLUMN user TEXT")
            self._conn.commit()

    def _recover(self) -> None:
        # Any job left 'running' when the process died is requeued.
        with self._lock:
            cur = self._conn.execute(
                "UPDATE jobs SET status='pending', updated_at=? WHERE status='running'",
                (_now(),),
            )
            self._conn.commit()
        if cur.rowcount:
            log.warning("requeued %d job(s) left running from a previous run", cur.rowcount)

    # -------------------------------------------------------------- public API

    def enqueue(self, action: str, folder: str, items: list[dict], extra: dict,
                user: str | None = None) -> int:
        """Insert a job. `items` is a list of dicts with keys:
        localguid (required), guid, subject, sender, recipient (all optional).
        `user` is the acting user's email, recorded for attribution/visibility."""
        if action not in VALID_ACTIONS:
            raise ValueError(f"unknown action: {action}")
        now = _now()
        with self._lock:
            cur = self._conn.execute(
                "INSERT INTO jobs (action, folder, user, extra, status, total, created_at, updated_at)"
                " VALUES (?,?,?,?,?,?,?,?)",
                (action, folder, user, json.dumps(extra or {}), "pending", len(items), now, now),
            )
            job_id = cur.lastrowid
            self._conn.executemany(
                "INSERT INTO job_items (job_id, localguid, guid, folder, subject, sender, recipient)"
                " VALUES (?,?,?,?,?,?,?)",
                [
                    (job_id, it["localguid"], it.get("guid"), folder,
                     it.get("subject"), it.get("sender"), it.get("recipient"))
                    for it in items
                ],
            )
            self._conn.commit()
        log.info("enqueued job #%d: action=%s folder=%r user=%s items=%d",
                 job_id, action, folder, user, len(items))
        return job_id

    def acted_localguids(self, folder: str) -> set[str]:
        """localguids in `folder` that belong to a non-failed job (hide them from the list)."""
        with self._lock:
            rows = self._conn.execute(
                "SELECT ji.localguid FROM job_items ji JOIN jobs j ON j.id = ji.job_id"
                " WHERE ji.folder = ? AND j.status != 'failed'",
                (folder,),
            ).fetchall()
        return {r["localguid"] for r in rows}

    def recent_jobs(self, limit: int = 25, user: str | None = None) -> list[dict]:
        """Recent jobs, newest first. Pass `user` to scope to one owner (None = all)."""
        cols = ("id, action, folder, user, status, processed, total, error,"
                " created_at, updated_at")
        with self._lock:
            if user is None:
                rows = self._conn.execute(
                    f"SELECT {cols} FROM jobs ORDER BY id DESC LIMIT ?", (limit,)
                ).fetchall()
            else:
                rows = self._conn.execute(
                    f"SELECT {cols} FROM jobs WHERE user = ? ORDER BY id DESC LIMIT ?",
                    (user, limit),
                ).fetchall()
        return [dict(r) for r in rows]

    def has_active_jobs(self) -> bool:
        with self._lock:
            row = self._conn.execute(
                "SELECT 1 FROM jobs WHERE status IN ('pending','running') LIMIT 1"
            ).fetchone()
        return row is not None

    def retry(self, job_id: int) -> bool:
        """Requeue a failed job. Returns False if it isn't a failed job."""
        with self._lock:
            cur = self._conn.execute(
                "UPDATE jobs SET status='pending', processed=0, error=NULL, updated_at=?"
                " WHERE id=? AND status='failed'",
                (_now(), job_id),
            )
            self._conn.commit()
        return cur.rowcount > 0

    # -------------------------------------------------------------- worker loop

    def start(self) -> None:
        thread = threading.Thread(target=self._run_loop, name="job-worker", daemon=True)
        thread.start()

    def stop(self) -> None:
        self._stop.set()

    def _run_loop(self) -> None:
        while not self._stop.is_set():
            job = self._claim_next()
            if job is None:
                time.sleep(_POLL_SECONDS)
                continue
            try:
                self._process(job)
            except Exception:  # noqa: BLE001 - never let the worker thread die
                self._mark_failed(job["id"], traceback.format_exc())

    def _claim_next(self) -> dict | None:
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM jobs WHERE status='pending' ORDER BY id ASC LIMIT 1"
            ).fetchone()
            if row is None:
                return None
            self._conn.execute(
                "UPDATE jobs SET status='running', updated_at=? WHERE id=?",
                (_now(), row["id"]),
            )
            self._conn.commit()
            return dict(row)

    def _process(self, job: dict) -> None:
        job_id = job["id"]
        action = job["action"]
        folder = job["folder"]
        user = job["user"] if "user" in job.keys() else None
        extra = json.loads(job["extra"] or "{}")
        items = self._items_for(job_id)

        # Snapshot config + client ONCE at job start — a mid-job admin edit affects the
        # next job, not this one (see AGENTS.md "snapshot-per-job").
        cfg = self._cfg()
        pps = self._pps()
        chunk_size = int(cfg.get("chunk_size", 25))
        log.info(
            "job #%d start: action=%s folder=%r user=%s total=%d",
            job_id, action, folder, user, len(items),
        )

        dst = _destination(action, extra, cfg)
        errors: list[str] = []
        processed = 0
        started = time.monotonic()
        for chunk in _chunks(items, chunk_size):
            try:
                self._run_action(pps, cfg, action, folder, chunk, extra)
                processed += len(chunk)
                self._set_processed(job_id, processed)
                self._log_ops(job_id, action, folder, dst, chunk, user, "ok", None)
            except PPSError as exc:
                log.error("job #%d chunk failed: %s", job_id, exc)
                errors.append(f"chunk failed: {exc}")
                self._log_ops(job_id, action, folder, dst, chunk, user, "FAILED", str(exc))
            except Exception as exc:  # noqa: BLE001
                log.exception("job #%d chunk raised", job_id)
                errors.append(f"chunk failed: {exc}")
                self._log_ops(job_id, action, folder, dst, chunk, user, "FAILED", str(exc))

        took = _fmt_duration(time.monotonic() - started)
        if errors:
            log.error("job #%d FAILED: %d/%d processed, %d chunk error(s), took %s",
                      job_id, processed, len(items), len(errors), took)
            self._mark_failed(job_id, "\n".join(errors))
        else:
            log.info("job #%d done: %d/%d processed, took %s",
                     job_id, processed, len(items), took)
            self._mark_done(job_id)

    def _log_ops(self, job_id: int, action: str, src: str, dst: str | None,
                 chunk: list[dict], user: str | None,
                 result: str, error: str | None) -> None:
        """Write one audit line per message in the chunk to the operations log."""
        for it in chunk:
            fields = [
                f"job=#{job_id}",
                f"user={user or '-'}",
                f"action={action}",
                f"result={result}",
                f"src={src!r}",
                f"dst={dst!r}",
                f"localguid={it.get('localguid')}",
                f"guid={it.get('guid')}",
                f"from={_clip(it.get('sender'))!r}",
                f"rcpt={_clip(it.get('recipient'))!r}",
                f"subject={_clip(it.get('subject'))!r}",
            ]
            if error:
                fields.append(f"error={_clip(error, 300)!r}")
            ops_log.info(" ".join(fields))

    # -------------------------------------------------------------- action logic

    def _run_action(self, pps: PPSClient, cfg: Mapping, action: str,
                    folder: str, chunk: list[dict], extra: dict) -> None:
        localguids = [it["localguid"] for it in chunk]
        deleted_folder = cfg.get("deleted_folder")

        if action == "release":
            pps.act("release", folder, localguids, deletedfolder=deleted_folder)

        elif action == "delete":
            pps.act("delete", folder, localguids, deletedfolder=deleted_folder)

        elif action == "move":
            target = extra.get("targetfolder")
            if not target:
                raise ValueError("move action requires targetfolder")
            pps.act("move", folder, localguids, targetfolder=target)

        elif action == "report_release":
            # Data-driven pipeline (release -> move -> delete, configurable subset).
            pipeline.run_pipeline(pps, cfg, folder, chunk)

        else:
            raise ValueError(f"unknown action: {action}")

    # -------------------------------------------------------------- db helpers

    def _items_for(self, job_id: int) -> list[dict]:
        with self._lock:
            rows = self._conn.execute(
                "SELECT localguid, guid, subject, sender, recipient"
                " FROM job_items WHERE job_id=?", (job_id,)
            ).fetchall()
        return [dict(r) for r in rows]

    def _set_processed(self, job_id: int, processed: int) -> None:
        with self._lock:
            self._conn.execute(
                "UPDATE jobs SET processed=?, updated_at=? WHERE id=?",
                (processed, _now(), job_id),
            )
            self._conn.commit()

    def _mark_done(self, job_id: int) -> None:
        with self._lock:
            self._conn.execute(
                "UPDATE jobs SET status='done', updated_at=? WHERE id=?",
                (_now(), job_id),
            )
            self._conn.commit()

    def _mark_failed(self, job_id: int, error: str) -> None:
        with self._lock:
            self._conn.execute(
                "UPDATE jobs SET status='failed', error=?, updated_at=? WHERE id=?",
                (error, _now(), job_id),
            )
            self._conn.commit()


def _destination(action: str, extra: dict, cfg: Mapping) -> str | None:
    """The folder messages end up in, for the ops log."""
    if action == "move":
        return extra.get("targetfolder")
    if action == "report_release":
        return pipeline.pipeline_destination(cfg)
    if action in ("release", "delete"):
        return cfg.get("deleted_folder")
    return None


def _chunks(seq: list, size: int):
    for i in range(0, len(seq), size):
        yield seq[i : i + size]


def _clip(value, length: int = 120) -> str:
    """Collapse newlines and truncate a value for a single-line log field."""
    s = " ".join(str(value or "").split())
    return s[:length] + ("…" if len(s) > length else "")


def _fmt_duration(seconds: float) -> str:
    """Human-readable elapsed time, e.g. '0.42s', '12.3s', '3m 07s', '1h 04m'."""
    if seconds < 60:
        return f"{seconds:.2f}s"
    if seconds < 3600:
        m, s = divmod(int(seconds), 60)
        return f"{m}m {s:02d}s"
    h, rem = divmod(int(seconds), 3600)
    m = rem // 60
    return f"{h}h {m:02d}m"