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 pps_client import PPSClient, PPSError

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

# 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):
        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()

    # -------------------------------------------------------------- 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,
                    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);
                """
            )
            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 {'localguid':..., 'guid':...} dicts."""
        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) VALUES (?,?,?,?)",
                [(job_id, it["localguid"], it.get("guid"), folder) for it in items],
            )
            self._conn.commit()
        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)
        )

        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)
            except PPSError as exc:
                log.error("job #%d chunk failed: %s", job_id, exc)
                errors.append(f"chunk failed: {exc}")
            except Exception as exc:  # noqa: BLE001
                log.exception("job #%d chunk raised", job_id)
                errors.append(f"chunk failed: {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)

    # -------------------------------------------------------------- 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 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]