proofpoint-quarantine-portal/app.py

"""PPS Quarantine Manager — Flask frontend + API.

Serves the single-page UI and proxies to the Proofpoint Quarantine Search REST API.
Slow actions (release/move/delete) are handed to the background JobQueue and return
immediately; the UI polls /api/jobs for progress.
"""

from __future__ import annotations

import email
import logging
import re
import tomllib
from email import policy
from functools import wraps
from pathlib import Path

from flask import (
    Flask,
    abort,
    jsonify,
    redirect,
    render_template,
    request,
    session,
    url_for,
)

from pps_client import PPSClient, PPSError
from worker import VALID_ACTIONS, JobQueue

CONFIG_PATH = Path(__file__).with_name("config.toml")


def load_config() -> dict:
    if not CONFIG_PATH.exists():
        raise SystemExit(
            f"Missing {CONFIG_PATH.name}. Copy config.example.toml to config.toml and edit it."
        )
    with CONFIG_PATH.open("rb") as fh:
        return tomllib.load(fh)


config = load_config()
_pps_cfg = config["pps"]
_q_cfg = config["quarantine"]
_app_cfg = config["app"]
_auth_cfg = config["auth"]

logging.basicConfig(
    level=getattr(logging, str(_app_cfg.get("log_level", "INFO")).upper(), logging.INFO),
    format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
)
log = logging.getLogger("pps.app")

pps = PPSClient(
    base_url=_pps_cfg["base_url"],
    username=_pps_cfg["username"],
    password=_pps_cfg["password"],
    verify_tls=_pps_cfg.get("verify_tls", False),
    timeout=int(_pps_cfg.get("timeout", 120)),
    client_cert=_pps_cfg.get("client_cert") or None,
    client_key=_pps_cfg.get("client_key") or None,
)

queue = JobQueue(
    _app_cfg.get("db_path", "jobs.db"),
    pps,
    _q_cfg,
    ops_log_path=_app_cfg.get("worker_log", "worker.log"),
)
queue.start()

app = Flask(__name__)
app.secret_key = _app_cfg["secret_key"]


# ------------------------------------------------------------------ auth

def login_required(view):
    @wraps(view)
    def wrapped(*args, **kwargs):
        if not session.get("user"):
            if request.path.startswith("/api/"):
                abort(401)
            return redirect(url_for("login", next=request.path))
        return view(*args, **kwargs)

    return wrapped


@app.route("/login", methods=["GET", "POST"])
def login():
    error = None
    if request.method == "POST":
        username = request.form.get("username", "")
        password = request.form.get("password", "")
        if username == _auth_cfg["username"] and password == _auth_cfg["password"]:
            session["user"] = username
            return redirect(request.args.get("next") or url_for("index"))
        error = "Invalid credentials"
    return render_template("login.html", error=error)


@app.route("/logout")
def logout():
    session.clear()
    return redirect(url_for("login"))


# ------------------------------------------------------------------ pages

@app.route("/")
@login_required
def index():
    return render_template("index.html")


# ------------------------------------------------------------------ JSON API

@app.route("/api/config")
@login_required
def api_config():
    return jsonify(
        {
            "folders": _q_cfg.get("folders", []),
            "default_folder": _q_cfg.get("default_folder", "Quarantine"),
            "default_limit": int(_q_cfg.get("default_limit", 200)),
            "report_release_folder": _q_cfg.get("report_release_folder"),
            "deleted_folder": _q_cfg.get("deleted_folder"),
        }
    )


@app.route("/api/messages")
@login_required
def api_messages():
    """Return one page of a folder, newest-first.

    The client pages through the whole folder by passing `before` (the oldest date
    it has seen) as a cursor, since the PPS API caps each response at 1000 rows and
    offers no offset paging. `page.has_more` and `page.oldest` drive that loop.
    Messages with a pending/running/done action are filtered out server-side.
    """
    folder = request.args.get("folder") or _q_cfg.get("default_folder", "Quarantine")
    page_size = _clamp_limit(request.args.get("limit"))
    before = request.args.get("before") or None
    try:
        records = pps.search(
            folder,
            _q_cfg.get("list_query", "from=*"),
            limit=page_size,
            days_back=int(_q_cfg.get("default_days_back", 7)),
            enddate=before,
        )
    except PPSError as exc:
        return _pps_error_response(exc, f"search folder={folder!r} before={before!r}")

    raw_count = len(records)
    hidden = queue.acted_localguids(folder)
    messages = [
        {
            "date": r.get("date"),
            "from": r.get("from"),
            "rcpts": r.get("rcpts", []),
            "subject": r.get("subject"),
            "guid": r.get("guid"),
            "localguid": r.get("localguid"),
            "size": r.get("size"),
        }
        for r in records
        if r.get("localguid") not in hidden
    ]
    # Cursor is the oldest date in the RAW batch (before the acted-on filter), so
    # paging never terminates early just because a page was mostly filtered out.
    dates = [r.get("date") for r in records if r.get("date")]
    oldest = min(dates) if dates else None  # fixed-width strings sort chronologically
    return jsonify(
        {
            "folder": folder,
            "messages": messages,
            "page": {
                "raw_count": raw_count,
                "oldest": oldest,
                "has_more": raw_count >= page_size,
            },
        }
    )


@app.route("/api/message/<path:guid>")
@login_required
def api_message(guid: str):
    try:
        raw = pps.get_raw(guid)
    except PPSError as exc:
        return _pps_error_response(exc, f"get_raw guid={guid!r}")
    return jsonify(_parse_message(raw))


@app.route("/api/actions", methods=["POST"])
@login_required
def api_actions():
    body = request.get_json(silent=True) or {}
    action = body.get("action")
    folder = body.get("folder")
    messages = body.get("messages", [])  # [{localguid, guid, subject, sender, recipient}, ...]
    targetfolder = body.get("targetfolder")

    if action not in VALID_ACTIONS:
        return jsonify({"error": f"unknown action: {action}"}), 400
    if not folder:
        return jsonify({"error": "folder is required"}), 400
    if not messages:
        return jsonify({"error": "no messages selected"}), 400
    if action == "move" and not targetfolder:
        return jsonify({"error": "move requires a targetfolder"}), 400

    items = [
        {
            "localguid": m.get("localguid"),
            "guid": m.get("guid"),
            "subject": m.get("subject"),
            "sender": m.get("sender"),
            "recipient": m.get("recipient"),
        }
        for m in messages
        if m.get("localguid")
    ]
    if not items:
        return jsonify({"error": "no valid messages (missing localguid)"}), 400

    extra = {"targetfolder": targetfolder} if action == "move" else {}
    job_id = queue.enqueue(action, folder, items, extra)
    return jsonify({"job_id": job_id, "queued": len(items)}), 202


@app.route("/api/jobs")
@login_required
def api_jobs():
    return jsonify({"jobs": queue.recent_jobs(limit=25)})


# ------------------------------------------------------------------ helpers

def _pps_error_response(exc: PPSError, context: str):
    """Log a PPS failure with full context and return a clean 502 JSON body."""
    log.error("%s failed: %s", context, exc)
    return (
        jsonify(
            {
                "error": str(exc),
                "detail": exc.body,
                "status": exc.status,
                "reqid": exc.reqid,
            }
        ),
        502,
    )


def _clamp_limit(raw: str | None) -> int:
    default = int(_q_cfg.get("default_limit", 200))
    try:
        value = int(raw) if raw is not None else default
    except (TypeError, ValueError):
        value = default
    return max(1, min(value, 1000))


def _parse_message(raw: bytes) -> dict:
    """Parse raw RFC822 bytes into headers + text/html bodies."""
    msg = email.message_from_bytes(raw, policy=policy.default)
    headers = {
        "from": msg.get("From", ""),
        "to": msg.get("To", ""),
        "cc": msg.get("Cc", ""),
        "date": msg.get("Date", ""),
        "subject": msg.get("Subject", ""),
    }

    text_body = ""
    html_body = ""
    try:
        text_part = msg.get_body(preferencelist=("plain",))
        if text_part is not None:
            text_body = text_part.get_content()
    except Exception:  # noqa: BLE001 - malformed MIME shouldn't 500 the view
        text_body = ""
    try:
        html_part = msg.get_body(preferencelist=("html",))
        if html_part is not None:
            html_body = html_part.get_content()
    except Exception:  # noqa: BLE001
        html_body = ""

    if not text_body and not html_body:
        # Non-multipart or odd structure: fall back to the raw payload.
        try:
            text_body = msg.get_content()
        except Exception:  # noqa: BLE001
            text_body = raw.decode("utf-8", errors="replace")

    return {
        "headers": headers,
        "text": text_body,
        "html": html_body,
        "attachments": _parse_attachments(msg),
        "raw_headers": _raw_headers(raw),
    }


def _parse_attachments(msg) -> list[dict]:
    """List attachment parts as {filename, content_type, size} (size in bytes)."""
    attachments = []
    try:
        for part in msg.iter_attachments():
            payload = part.get_payload(decode=True)
            attachments.append(
                {
                    "filename": part.get_filename() or "(unnamed)",
                    "content_type": part.get_content_type(),
                    "size": len(payload) if payload is not None else None,
                }
            )
    except Exception:  # noqa: BLE001 - malformed MIME shouldn't 500 the view
        pass
    return attachments


def _raw_headers(raw: bytes) -> str:
    """The literal RFC822 header block (everything before the first blank line)."""
    head = re.split(rb"\r?\n\r?\n", raw, maxsplit=1)[0]
    return head.decode("utf-8", errors="replace")


if __name__ == "__main__":
    from waitress import serve

    host = _app_cfg.get("listen", "127.0.0.1")
    port = int(_app_cfg.get("port", 8080))
    log.info("PPS Quarantine Manager listening on http://%s:%s", host, port)
    # Single process so the worker thread and Flask share memory + SQLite connection.
    serve(app, host=host, port=port, threads=8)