proofpoint-quarantine-portal/pps_client.py

"""Thin client for the Proofpoint Protection Server Quarantine Search REST API.

Covers the three things the PoC needs:
  - search(): GET /rest/v1/quarantine   -> list of message records (JSON)
  - get_raw(): GET .../quarantine?guid= -> raw RFC822 bytes of one message
  - act():     POST /rest/v1/quarantine -> release / move / delete actions

All requests use HTTP Basic auth against the admin port (default 10000).
See "Proofpoint Public APIs rev W", "Proofpoint Quarantine Search REST API".
"""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from urllib.parse import urljoin

import requests
import urllib3

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

# PPS accepts "YYYY-MM-DD HH:MM:SS" UTC time strings for startdate/enddate.
_DATE_FMT = "%Y-%m-%d %H:%M:%S"

_HEADERS = {
    "Accept": "application/json",
    "User-Agent": "pps-quarantine-manager/0.1",
}

# How much of a response body to keep in an error message / log line.
_BODY_SNIPPET = 2000


class PPSError(RuntimeError):
    """Raised when a PPS request fails. Carries status, response body and request-id."""

    def __init__(
        self,
        message: str,
        status: int | None = None,
        body: str = "",
        reqid: str | None = None,
    ):
        super().__init__(message)
        self.status = status
        self.body = body
        self.reqid = reqid


class PPSClient:
    def __init__(
        self,
        base_url: str,
        username: str,
        password: str,
        verify_tls: bool | str = False,
        timeout: int = 120,
        client_cert: str | None = None,
        client_key: str | None = None,
    ):
        # Endpoint is a fixed path on the admin service.
        self.endpoint = urljoin(base_url.rstrip("/") + "/", "rest/v1/quarantine")
        self.timeout = timeout
        self.session = requests.Session()
        self.session.auth = (username, password)
        self.session.headers.update(_HEADERS)
        self.session.verify = verify_tls

        # Client-certificate (mutual TLS) auth. Some PPS deployments sit behind nginx
        # configured to require a client cert ("400 No required SSL certificate was sent").
        # `client_cert` may be a combined cert+key PEM, or a cert paired with `client_key`.
        if client_cert:
            self.session.cert = (client_cert, client_key) if client_key else client_cert
            log.info("Using client certificate for mutual TLS: %s", client_cert)

        if verify_tls is False:
            # Intentionally skipping verification (PoC / self-signed PPS certs). Silence
            # the per-request InsecureRequestWarning but say so once, loudly, at startup.
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
            log.warning(
                "TLS verification is DISABLED for %s — do not use in production.",
                self.endpoint,
            )

    # ------------------------------------------------------------------ search

    def search(
        self,
        folder: str,
        query: str,
        limit: int = 200,
        days_back: int = 7,
    ) -> list[dict]:
        """Search a quarantine folder and return up to `limit` message records.

        `query` is a raw filter fragment the API requires, e.g. "from=*" (match all)
        or "rcpt=@example.com". The API needs at least one of from/rcpt/subject.
        Results are sorted newest-first by the server; we truncate client-side.
        """
        params = self._query_to_params(query)
        params["folder"] = folder
        startdate = datetime.now(timezone.utc) - timedelta(days=days_back)
        params["startdate"] = startdate.strftime(_DATE_FMT)

        resp = self._request("GET", params=params)
        data = resp.json()
        records = data.get("records", []) or []
        log.info(
            "search folder=%r query=%r -> %d record(s)", folder, query, len(records)
        )
        if limit and len(records) > limit:
            records = records[:limit]
        return records

    def get_raw(self, guid: str) -> bytes:
        """Fetch the raw RFC822 bytes for a single message by its stable guid."""
        resp = self._request("GET", params={"guid": guid})
        return resp.content

    # ------------------------------------------------------------------ actions

    def act(
        self,
        action: str,
        folder: str,
        localguids: list[str],
        *,
        targetfolder: str | None = None,
        deletedfolder: str | None = None,
        scan: bool = False,
    ) -> dict:
        """POST an action for one or more messages in `folder`.

        `localguids` are the folder-local handles (not the stable guid). One POST can
        carry many, comma-separated. Returns the parsed JSON status body.
        """
        if not localguids:
            return {"status": "no messages"}
        payload: dict[str, str] = {
            "action": action,
            "folder": folder,
            "localguid": ",".join(localguids),
        }
        if targetfolder:
            payload["targetfolder"] = targetfolder
        if deletedfolder:
            payload["deletedfolder"] = deletedfolder
        if scan:
            payload["scan"] = "t"

        log.info(
            "action=%s folder=%r count=%d targetfolder=%r deletedfolder=%r",
            action, folder, len(localguids), targetfolder, deletedfolder,
        )
        resp = self._request("POST", json=payload)
        # Success is 200 or 204 with a small JSON body; 204 has no body.
        if resp.status_code == 204 or not resp.content:
            result = {"status": "ok"}
        else:
            try:
                result = resp.json()
            except ValueError:
                result = {"status": resp.text.strip()}
        log.info("action=%s folder=%r -> %s", action, folder, result)
        return result

    # ------------------------------------------------------------------ transport

    def _request(self, method: str, *, params: dict | None = None,
                 json: dict | None = None) -> requests.Response:
        """Issue one request, logging it and converting any failure into PPSError."""
        try:
            resp = self.session.request(
                method, self.endpoint, params=params, json=json, timeout=self.timeout
            )
        except requests.RequestException as exc:
            log.error("%s %s failed at transport: %s", method, self.endpoint, exc)
            raise PPSError(f"PPS request failed ({method}): {exc}") from exc

        reqid = resp.headers.get("x-pps-reqid")
        log.debug(
            "%s %s -> HTTP %d (x-pps-reqid=%s)",
            method, resp.url, resp.status_code, reqid,
        )
        self._raise_for_status(resp, method, reqid)
        return resp

    @staticmethod
    def _raise_for_status(resp: requests.Response, method: str, reqid: str | None) -> None:
        if resp.status_code < 400:
            return
        body = (resp.text or "").strip()
        snippet = body[:_BODY_SNIPPET] + ("…" if len(body) > _BODY_SNIPPET else "")
        # The full picture in one place: status, PPS request-id (to grep PPS logs),
        # the resolved request URL, and whatever body PPS returned (often the reason).
        detail = (
            f"PPS returned HTTP {resp.status_code} for {method} {resp.url} "
            f"(x-pps-reqid={reqid or 'n/a'}): {snippet or '<empty body>'}"
        )
        log.error(detail)
        raise PPSError(detail, status=resp.status_code, body=body, reqid=reqid)

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

    @staticmethod
    def _query_to_params(query: str) -> dict[str, str]:
        """Parse a raw filter fragment like 'from=*' or 'rcpt=@x.com&subject=*hi*'."""
        params: dict[str, str] = {}
        for part in (query or "").split("&"):
            part = part.strip()
            if not part or "=" not in part:
                continue
            key, value = part.split("=", 1)
            params[key.strip()] = value.strip()
        if not params:
            # API demands at least one of from/rcpt/subject; default to match-all.
            params["from"] = "*"
        return params