app.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291
  1. """PPS Quarantine Manager — Flask frontend + API.
  2. Serves the single-page UI and proxies to the Proofpoint Quarantine Search REST API.
  3. Slow actions (release/move/delete) are handed to the background JobQueue and return
  4. immediately; the UI polls /api/jobs for progress.
  5. """
  6. from __future__ import annotations
  7. import email
  8. import logging
  9. import tomllib
  10. from email import policy
  11. from functools import wraps
  12. from pathlib import Path
  13. from flask import (
  14. Flask,
  15. abort,
  16. jsonify,
  17. redirect,
  18. render_template,
  19. request,
  20. session,
  21. url_for,
  22. )
  23. from pps_client import PPSClient, PPSError
  24. from worker import VALID_ACTIONS, JobQueue
  25. CONFIG_PATH = Path(__file__).with_name("config.toml")
  26. def load_config() -> dict:
  27. if not CONFIG_PATH.exists():
  28. raise SystemExit(
  29. f"Missing {CONFIG_PATH.name}. Copy config.example.toml to config.toml and edit it."
  30. )
  31. with CONFIG_PATH.open("rb") as fh:
  32. return tomllib.load(fh)
  33. config = load_config()
  34. _pps_cfg = config["pps"]
  35. _q_cfg = config["quarantine"]
  36. _app_cfg = config["app"]
  37. _auth_cfg = config["auth"]
  38. logging.basicConfig(
  39. level=getattr(logging, str(_app_cfg.get("log_level", "INFO")).upper(), logging.INFO),
  40. format="%(asctime)s %(levelname)-7s %(name)s: %(message)s",
  41. )
  42. log = logging.getLogger("pps.app")
  43. pps = PPSClient(
  44. base_url=_pps_cfg["base_url"],
  45. username=_pps_cfg["username"],
  46. password=_pps_cfg["password"],
  47. verify_tls=_pps_cfg.get("verify_tls", False),
  48. timeout=int(_pps_cfg.get("timeout", 120)),
  49. client_cert=_pps_cfg.get("client_cert") or None,
  50. client_key=_pps_cfg.get("client_key") or None,
  51. )
  52. queue = JobQueue(
  53. _app_cfg.get("db_path", "jobs.db"),
  54. pps,
  55. _q_cfg,
  56. ops_log_path=_app_cfg.get("worker_log", "worker.log"),
  57. )
  58. queue.start()
  59. app = Flask(__name__)
  60. app.secret_key = _app_cfg["secret_key"]
  61. # ------------------------------------------------------------------ auth
  62. def login_required(view):
  63. @wraps(view)
  64. def wrapped(*args, **kwargs):
  65. if not session.get("user"):
  66. if request.path.startswith("/api/"):
  67. abort(401)
  68. return redirect(url_for("login", next=request.path))
  69. return view(*args, **kwargs)
  70. return wrapped
  71. @app.route("/login", methods=["GET", "POST"])
  72. def login():
  73. error = None
  74. if request.method == "POST":
  75. username = request.form.get("username", "")
  76. password = request.form.get("password", "")
  77. if username == _auth_cfg["username"] and password == _auth_cfg["password"]:
  78. session["user"] = username
  79. return redirect(request.args.get("next") or url_for("index"))
  80. error = "Invalid credentials"
  81. return render_template("login.html", error=error)
  82. @app.route("/logout")
  83. def logout():
  84. session.clear()
  85. return redirect(url_for("login"))
  86. # ------------------------------------------------------------------ pages
  87. @app.route("/")
  88. @login_required
  89. def index():
  90. return render_template("index.html")
  91. # ------------------------------------------------------------------ JSON API
  92. @app.route("/api/config")
  93. @login_required
  94. def api_config():
  95. return jsonify(
  96. {
  97. "folders": _q_cfg.get("folders", []),
  98. "default_folder": _q_cfg.get("default_folder", "Quarantine"),
  99. "default_limit": int(_q_cfg.get("default_limit", 200)),
  100. "report_release_folder": _q_cfg.get("report_release_folder"),
  101. "deleted_folder": _q_cfg.get("deleted_folder"),
  102. }
  103. )
  104. @app.route("/api/messages")
  105. @login_required
  106. def api_messages():
  107. folder = request.args.get("folder") or _q_cfg.get("default_folder", "Quarantine")
  108. limit = _clamp_limit(request.args.get("limit"))
  109. try:
  110. records = pps.search(
  111. folder,
  112. _q_cfg.get("list_query", "from=*"),
  113. limit=limit,
  114. days_back=int(_q_cfg.get("default_days_back", 7)),
  115. )
  116. except PPSError as exc:
  117. return _pps_error_response(exc, f"search folder={folder!r}")
  118. hidden = queue.acted_localguids(folder)
  119. messages = [
  120. {
  121. "date": r.get("date"),
  122. "from": r.get("from"),
  123. "rcpts": r.get("rcpts", []),
  124. "subject": r.get("subject"),
  125. "guid": r.get("guid"),
  126. "localguid": r.get("localguid"),
  127. "size": r.get("size"),
  128. }
  129. for r in records
  130. if r.get("localguid") not in hidden
  131. ]
  132. # Ordering is handled client-side (the toolbar Sort control); the API returns
  133. # newest-first by date.
  134. return jsonify({"folder": folder, "count": len(messages), "messages": messages})
  135. @app.route("/api/message/<path:guid>")
  136. @login_required
  137. def api_message(guid: str):
  138. try:
  139. raw = pps.get_raw(guid)
  140. except PPSError as exc:
  141. return _pps_error_response(exc, f"get_raw guid={guid!r}")
  142. return jsonify(_parse_message(raw))
  143. @app.route("/api/actions", methods=["POST"])
  144. @login_required
  145. def api_actions():
  146. body = request.get_json(silent=True) or {}
  147. action = body.get("action")
  148. folder = body.get("folder")
  149. messages = body.get("messages", []) # [{localguid, guid, subject, sender, recipient}, ...]
  150. targetfolder = body.get("targetfolder")
  151. if action not in VALID_ACTIONS:
  152. return jsonify({"error": f"unknown action: {action}"}), 400
  153. if not folder:
  154. return jsonify({"error": "folder is required"}), 400
  155. if not messages:
  156. return jsonify({"error": "no messages selected"}), 400
  157. if action == "move" and not targetfolder:
  158. return jsonify({"error": "move requires a targetfolder"}), 400
  159. items = [
  160. {
  161. "localguid": m.get("localguid"),
  162. "guid": m.get("guid"),
  163. "subject": m.get("subject"),
  164. "sender": m.get("sender"),
  165. "recipient": m.get("recipient"),
  166. }
  167. for m in messages
  168. if m.get("localguid")
  169. ]
  170. if not items:
  171. return jsonify({"error": "no valid messages (missing localguid)"}), 400
  172. extra = {"targetfolder": targetfolder} if action == "move" else {}
  173. job_id = queue.enqueue(action, folder, items, extra)
  174. return jsonify({"job_id": job_id, "queued": len(items)}), 202
  175. @app.route("/api/jobs")
  176. @login_required
  177. def api_jobs():
  178. return jsonify({"jobs": queue.recent_jobs(limit=25)})
  179. # ------------------------------------------------------------------ helpers
  180. def _pps_error_response(exc: PPSError, context: str):
  181. """Log a PPS failure with full context and return a clean 502 JSON body."""
  182. log.error("%s failed: %s", context, exc)
  183. return (
  184. jsonify(
  185. {
  186. "error": str(exc),
  187. "detail": exc.body,
  188. "status": exc.status,
  189. "reqid": exc.reqid,
  190. }
  191. ),
  192. 502,
  193. )
  194. def _clamp_limit(raw: str | None) -> int:
  195. default = int(_q_cfg.get("default_limit", 200))
  196. try:
  197. value = int(raw) if raw is not None else default
  198. except (TypeError, ValueError):
  199. value = default
  200. return max(1, min(value, 1000))
  201. def _parse_message(raw: bytes) -> dict:
  202. """Parse raw RFC822 bytes into headers + text/html bodies."""
  203. msg = email.message_from_bytes(raw, policy=policy.default)
  204. headers = {
  205. "from": msg.get("From", ""),
  206. "to": msg.get("To", ""),
  207. "cc": msg.get("Cc", ""),
  208. "date": msg.get("Date", ""),
  209. "subject": msg.get("Subject", ""),
  210. }
  211. text_body = ""
  212. html_body = ""
  213. try:
  214. text_part = msg.get_body(preferencelist=("plain",))
  215. if text_part is not None:
  216. text_body = text_part.get_content()
  217. except Exception: # noqa: BLE001 - malformed MIME shouldn't 500 the view
  218. text_body = ""
  219. try:
  220. html_part = msg.get_body(preferencelist=("html",))
  221. if html_part is not None:
  222. html_body = html_part.get_content()
  223. except Exception: # noqa: BLE001
  224. html_body = ""
  225. if not text_body and not html_body:
  226. # Non-multipart or odd structure: fall back to the raw payload.
  227. try:
  228. text_body = msg.get_content()
  229. except Exception: # noqa: BLE001
  230. text_body = raw.decode("utf-8", errors="replace")
  231. return {"headers": headers, "text": text_body, "html": html_body}
  232. if __name__ == "__main__":
  233. from waitress import serve
  234. host = _app_cfg.get("listen", "127.0.0.1")
  235. port = int(_app_cfg.get("port", 8080))
  236. log.info("PPS Quarantine Manager listening on http://%s:%s", host, port)
  237. # Single process so the worker thread and Flask share memory + SQLite connection.
  238. serve(app, host=host, port=port, threads=8)