proofpoint-quarantine-portal/pipeline.py

"""Configurable Report & Release pipeline.

"Report & Release" is a sequence of steps run per message chunk, chosen in the admin
panel by checkboxes. The order is FIXED: move → release. Config only toggles which steps
are included (a subsequence of ALLOWED_STEPS).

Deleting is deliberately NOT part of this pipeline — Delete is its own action/button, so
Report & Release can never destroy mail as a side effect.

Why the order and the re-reads are load-bearing (see AGENTS.md):

* `move` runs FIRST, straight out of the quarantine folder, using the localguids the job
  was enqueued with. The message reaches its destination folder before anything else
  touches it.
* `release` then runs where the message NOW is, and MUST NOT pass `deletedfolder` —
  that would relocate it again and invalidate every localguid.
* localguids are NEVER carried from one step to the next. They are folder-local, and
  PPS's backend is eventually consistent: after a relocating step the cached handles are
  stale, and acting on stale handles is exactly the race this pipeline kept hitting.
  Every step after the first waits `step_delay_seconds` and then re-reads where the
  messages actually are, by their stable `guid` (`locate_messages`).

The worker drives one step at a time via `run_step` and parks the job on the queue
during the delay, so a waiting pipeline never blocks jobs queued behind it
(`worker.JobQueue._process`).

Each step acts on the message's CURRENT location. Adding a future action = one
`@step(...)` function + an entry in ALLOWED_STEPS; no worker or frontend change.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Callable, Mapping

from pps_client import PPSError

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

# Canonical order. Also the validation whitelist (config_store imports this).
ALLOWED_STEPS: tuple[str, ...] = ("move", "release")

# Seconds to wait between two steps on the same message when config omits the setting.
# PPS's backend is eventually-consistent: acting again too soon (e.g. releasing a
# just-moved message) can hit a stale view. The delay lets the backend catch up — and
# only after it has passed do we re-read where the messages are.
DEFAULT_STEP_DELAY = 60


@dataclass(frozen=True)
class StepContext:
    pps: Any
    cfg: Mapping[str, Any]           # the quarantine config section
    folder: str                      # where the messages are RIGHT NOW
    localguids: list[str]            # freshly read handles, valid in `folder`
    chunk: list[dict]                # job_items rows: localguid, guid, subject, ...


@dataclass(frozen=True)
class StepResult:
    folder: str                      # where the messages are after this step
    dst: str | None                  # folder this step lands messages in (for ops log)


StepFn = Callable[[StepContext], StepResult]
STEPS: dict[str, StepFn] = {}


def step(name: str):
    def register(fn: StepFn) -> StepFn:
        STEPS[name] = fn
        return fn
    return register


@step("move")
def _move(ctx: StepContext) -> StepResult:
    target = ctx.cfg["report_release"]["move_target"]
    if not target:
        raise ValueError("report_release 'move' step requires move_target")
    try:
        ctx.pps.act("move", ctx.folder, ctx.localguids, targetfolder=target)
    except PPSError as exc:
        # As the first step, `move` runs on the localguids captured when the job was
        # enqueued — which may have gone stale in the meantime (or the message may
        # already have been relocated). Re-read the messages' real location and retry
        # once; a genuine failure raises again from `locate_messages` or the retry.
        log.warning("move from %r failed (%s) — re-reading message locations", ctx.folder, exc)
        folder, localguids = locate_messages(ctx.pps, ctx.cfg, ctx.folder, ctx.chunk)
        ctx.pps.act("move", folder, localguids, targetfolder=target)
    return StepResult(folder=target, dst=target)


@step("release")
def _release(ctx: StepContext) -> StepResult:
    # NO deletedfolder — this is the localguid-stability rule. Never add it here.
    ctx.pps.act("release", ctx.folder, ctx.localguids)
    return StepResult(folder=ctx.folder, dst=None)


# ---------------------------------------------------------------- driving the steps

def job_steps(cfg: Mapping[str, Any]) -> list[str]:
    """The configured steps, in canonical order."""
    return list(cfg["report_release"]["steps"])


def step_delay(cfg: Mapping[str, Any]) -> int:
    """Seconds to wait between steps (0 disables the wait)."""
    return int(cfg["report_release"].get("step_delay_seconds", DEFAULT_STEP_DELAY))


def run_step(pps, cfg: Mapping[str, Any], name: str, folder: str, chunk: list[dict],
             *, refresh: bool) -> str:
    """Run one pipeline step for one chunk; return the folder the messages end up in.

    `refresh=True` (every step after the first) discards the job's cached localguids and
    re-reads the messages' current handles in `folder` by their stable guid. The caller
    is responsible for having waited `step_delay` before asking for a refresh.
    """
    if refresh:
        folder, localguids = locate_messages(pps, cfg, folder, chunk)
    else:
        localguids = [it["localguid"] for it in chunk]
    result = STEPS[name](StepContext(pps=pps, cfg=cfg, folder=folder,
                                     localguids=localguids, chunk=chunk))
    return result.folder


def step_destination(cfg: Mapping[str, Any], name: str, folder: str) -> str | None:
    """Folder the messages are in after `name`, computed WITHOUT touching PPS."""
    if name == "move":
        return cfg["report_release"]["move_target"]
    return folder


def locate_messages(pps, cfg: Mapping[str, Any], folder: str,
                    chunk: list[dict]) -> tuple[str, list[str]]:
    """Re-read where `chunk`'s messages are right now: `(folder, localguids)`.

    Looks in `folder` first; if nothing matches there and a `deleted_folder` is
    configured, looks there too — some deployments' actions relocate messages to the
    deleted folder behind our back. Raises PPSError if the messages are nowhere.
    """
    if not any(it.get("guid") for it in chunk):
        raise PPSError("cannot locate messages: no guids stored for this chunk")
    candidates = [folder]
    deleted = cfg.get("deleted_folder")
    if deleted and deleted != folder:
        candidates.append(deleted)
    for i, candidate in enumerate(candidates):
        try:
            return candidate, refind_localguids(pps, cfg, candidate, chunk)
        except PPSError:
            if i == len(candidates) - 1:
                raise
            log.warning("messages not in %r — looking in %r", candidate, candidates[i + 1])
    raise AssertionError("unreachable")


def refind_localguids(pps, cfg: Mapping[str, Any], folder: str, chunk: list[dict]) -> list[str]:
    """Recover current localguids in `folder` by matching the stable guid.

    Raises PPSError if nothing matches.
    """
    wanted = {it["guid"]: it for it in chunk if it.get("guid")}
    if not wanted:
        raise PPSError("cannot recover messages: no guids stored for this chunk")
    records = pps.search(
        folder,
        cfg.get("list_query", "from=*"),
        limit=1000,
        days_back=int(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(f"messages not found in {folder!r} to continue the pipeline")
    log.info("re-read %d/%d message(s) in %r", len(found), len(wanted), folder)
    return found


def pipeline_destination(cfg: Mapping[str, Any]) -> str | None:
    """Final folder after all steps, computed WITHOUT touching PPS (for the ops log)."""
    folder = None
    for name in cfg["report_release"]["steps"]:
        if name == "move":
            folder = cfg["report_release"]["move_target"]
    return folder


def describe_pipeline(cfg: Mapping[str, Any]) -> str:
    """Human sentence for the admin panel, e.g.
    'moved to "Debugging - Josef" → released in place'."""
    rr = cfg.get("report_release", {})
    phrases = []
    for name in rr.get("steps", []):
        if name == "move":
            phrases.append(f'moved to "{rr.get("move_target", "?")}"')
        elif name == "release":
            phrases.append("released in place")
    return " → ".join(phrases) if phrases else "(no steps configured)"