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 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)
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():
    folder = request.args.get("folder") or _q_cfg.get("default_folder", "Quarantine")
    limit = _clamp_limit(request.args.get("limit"))
    try:
        records = pps.search(
            folder,
            _q_cfg.get("list_query", "from=*"),
            limit=limit,
            days_back=int(_q_cfg.get("default_days_back", 7)),
        )
    except PPSError as exc:
        return _pps_error_response(exc, f"search folder={folder!r}")

    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
    ]
    return jsonify({"folder": folder, "count": len(messages), "messages": messages})


@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}, ...]
    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")}
        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}


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)