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 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: PPSClient, qconfig: dict,
                 ops_log_path: str = "worker.log"):
        self.pps = pps
        self.cfg = qconfig
        self.chunk_size = int(qconfig.get("chunk_size", 25))
        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:
            self._conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    id         INTEGER PRIMARY KEY AUTOINCREMENT,
                    action     TEXT NOT NULL,
                    folder     TEXT NOT NULL,
                    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);
                """
            )
            # Migrate older DBs that lack the display columns.
            existing = {r["name"] for r in self._conn.execute("PRAGMA table_info(job_items)")}
            for col in ("subject", "sender", "recipient"):
                if col not in existing:
                    self._conn.execute(f"ALTER TABLE job_items ADD COLUMN {col} 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) -> int:
        """Insert a job. `items` is a list of dicts with keys:
        localguid (required), guid, subject, sender, recipient (all optional)."""
        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, extra, status, total, created_at, updated_at)"
                " VALUES (?,?,?,?,?,?,?)",
                (action, folder, 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 items=%d", job_id, action, folder, 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) -> list[dict]:
        with self._lock:
            rows = self._conn.execute(
                "SELECT id, action, folder, status, processed, total, error, created_at, updated_at"
                " FROM jobs ORDER BY id DESC LIMIT ?",
                (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

    # -------------------------------------------------------------- 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"]
        extra = json.loads(job["extra"] or "{}")
        items = self._items_for(job_id)
        log.info(
            "job #%d start: action=%s folder=%r total=%d", job_id, action, folder, len(items)
        )

        dst = self._destination(action, extra)
        errors: list[str] = []
        processed = 0
        for chunk in _chunks(items, self.chunk_size):
            try:
                self._run_action(action, folder, chunk, extra)
                processed += len(chunk)
                self._set_processed(job_id, processed)
                self._log_ops(job_id, action, folder, dst, chunk, "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, "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, "FAILED", str(exc))

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

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

    def _log_ops(self, job_id: int, action: str, src: str, dst: str | None,
                 chunk: list[dict], 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"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, action: str, folder: str, chunk: list[dict], extra: dict) -> None:
        localguids = [it["localguid"] for it in chunk]
        deleted_folder = self.cfg.get("deleted_folder")

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

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

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

        elif action == "report_release":
            self._report_release(folder, chunk)

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

    def _report_release(self, folder: str, chunk: list[dict]) -> None:
        """Release in place (delivers, keeps message), then move a copy to the report folder.

        Fallback: if release relocated the message to the deleted folder (GUI-like
        behavior), the move from `folder` fails; re-find each message in the deleted
        folder by its stable guid and move it from there instead.
        """
        target = self.cfg.get("report_release_folder")
        localguids = [it["localguid"] for it in chunk]

        # Step 1: release without deletedfolder -> message stays put, localguid stable.
        self.pps.act("release", folder, localguids)

        # Step 2: move the (now released) copy into the report folder.
        try:
            self.pps.act("move", folder, localguids, targetfolder=target)
        except PPSError:
            self._move_from_deleted_by_guid(chunk, target)

    def _move_from_deleted_by_guid(self, chunk: list[dict], target: str) -> None:
        deleted_folder = self.cfg.get("deleted_folder")
        if not deleted_folder:
            raise PPSError("release relocated messages but no deleted_folder configured")
        wanted = {it["guid"]: it for it in chunk if it.get("guid")}
        if not wanted:
            raise PPSError("cannot recover messages: no guids stored for report_release")

        records = self.pps.search(
            deleted_folder,
            self.cfg.get("list_query", "from=*"),
            limit=1000,
            days_back=int(self.cfg.get("default_days_back", 7)),
        )
        found = [
            r["localguid"]
            for r in records
            if r.get("guid") in wanted and r.get("localguid")
        ]
        if not found:
            raise PPSError(
                "released messages not found in deleted folder to move to report folder"
            )
        self.pps.act("move", deleted_folder, found, targetfolder=target)

    # -------------------------------------------------------------- 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 _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 "")