app.py

"""PCAPGraph — Flask server.

Drag-drop a .pcap/.pcapng, stream-parse it into flows, build a vis.js graph,
and serve it. Click a public IP to enrich it (offline-first).

All limits are env-configurable (see DESIGN.md §6):
    MAX_UPLOAD_MB   default 500
    MAX_PACKETS     default 5,000,000
    PARSE_TIMEOUT   default 120 (seconds)
    GEOIP_DIR       default ./geoip
    ABUSEIPDB_KEY   optional, enables reputation in online mode
"""

from __future__ import annotations

import os
import tempfile

from flask import Flask, jsonify, render_template, request
from werkzeug.utils import secure_filename

import pcap_parser
import graph_builder
import enrichment
import virustotal

MAX_UPLOAD_MB = int(os.environ.get("MAX_UPLOAD_MB", "500"))
MAX_PACKETS = int(os.environ.get("MAX_PACKETS", "5000000"))
PARSE_TIMEOUT = float(os.environ.get("PARSE_TIMEOUT", "120"))

ALLOWED_EXT = (".pcap", ".pcapng", ".cap")

app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = MAX_UPLOAD_MB * 1024 * 1024

UPLOAD_DIR = os.path.join(os.path.dirname(__file__), "uploads")
os.makedirs(UPLOAD_DIR, exist_ok=True)

# The capture FILE is always deleted after parsing. The parsed conversation
# table for the most-recent upload is kept in memory so reverse-lookup can scan
# every flow without re-uploading. Single-user local tool; one slot is enough.
LAST = {"result": None}


@app.route("/")
def index():
    return render_template(
        "index.html",
        max_upload_mb=MAX_UPLOAD_MB,
        geo_available=enrichment.geo_available(),
        vt_available=virustotal.available(),
    )


@app.route("/upload", methods=["POST"])
def upload():
    if "file" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400
    f = request.files["file"]
    if not f.filename:
        return jsonify({"error": "Empty filename"}), 400
    if not f.filename.lower().endswith(ALLOWED_EXT):
        return jsonify({"error": f"Unsupported file type. Use {', '.join(ALLOWED_EXT)}"}), 400

    name = secure_filename(f.filename)
    fd, path = tempfile.mkstemp(suffix="_" + name, dir=UPLOAD_DIR)
    os.close(fd)
    try:
        f.save(path)
        result = pcap_parser.parse(path, max_packets=MAX_PACKETS, timeout=PARSE_TIMEOUT)
        if result.parsed_packets == 0:
            return jsonify({"error": "No IPv4/IPv6 packets found in capture."}), 422
        graph = graph_builder.build(result)
        graph["filename"] = name
        graph["vt_available"] = virustotal.available()
        LAST["result"] = result
        return jsonify(graph)
    except Exception as exc:  # noqa: BLE001
        app.logger.exception("parse failed")
        return jsonify({"error": f"Failed to parse capture: {exc}"}), 500
    finally:
        try:
            os.remove(path)
        except OSError:
            pass


@app.route("/enrich/<ip>")
def enrich(ip):
    online = request.args.get("online", "0") in ("1", "true", "yes")
    return jsonify(enrichment.enrich(ip, online=online))


@app.route("/geo_batch", methods=["POST"])
def geo_batch():
    """Resolve a list of IPs to countries (ip-api.com batch). Powers the
    per-host Countries drill-down box."""
    data = request.get_json(silent=True) or {}
    ips = data.get("ips")
    if not isinstance(ips, list):
        return jsonify({"error": "expected {\"ips\": [...]}"}), 400
    return jsonify(enrichment.batch_country(ips[:1000]))


@app.route("/reverse")
def reverse():
    """Find which internal hosts contacted a given malicious IP / domain."""
    query = request.args.get("q", "")
    result = LAST.get("result")
    if result is None:
        return jsonify({"error": "Upload a capture first."}), 409
    return jsonify(graph_builder.reverse_lookup(result, query))


@app.route("/vt/<kind>/<path:value>")
def vt(kind, value):
    """VirusTotal detection count for an IP or domain (needs an API token)."""
    if kind not in ("ip", "domain"):
        return jsonify({"available": False, "note": "kind must be 'ip' or 'domain'."}), 400
    return jsonify(virustotal.lookup(kind, value))


@app.route("/vt_scan_info")
def vt_scan_info():
    """How many indicators a full scan would query — drives the warning dialog."""
    result = LAST.get("result")
    if result is None:
        return jsonify({"error": "Upload a capture first."}), 409
    cap = virustotal.scan_max()
    inds, n_ips, n_domains = graph_builder.scan_indicators(result, cap)
    return jsonify({"available": virustotal.available(), "cap": cap,
                    "ips": n_ips, "domains": n_domains, "planned": len(inds)})


@app.route("/vt_scan")
def vt_scan():
    """Bulk-check the capture's public IPs/domains on VirusTotal → top-5 malicious.
    Heavy on the VT API; the UI confirms before calling this."""
    result = LAST.get("result")
    if result is None:
        return jsonify({"error": "Upload a capture first."}), 409
    return jsonify(graph_builder.malicious_scan(result))


@app.errorhandler(413)
def too_large(_e):
    return jsonify({"error": f"File exceeds the {MAX_UPLOAD_MB} MB limit. "
                             f"Pre-slice it with editcap/tshark, or raise "
                             f"MAX_UPLOAD_MB."}), 413


if __name__ == "__main__":
    print(f"PCAPGraph → http://127.0.0.1:5000  "
          f"(limits: {MAX_UPLOAD_MB}MB / {MAX_PACKETS} pkts / {PARSE_TIMEOUT}s, "
          f"geo {'on' if enrichment.geo_available() else 'off'}, "
          f"virustotal {'on' if virustotal.available() else 'off'})")
    app.run(host="127.0.0.1", port=5000, debug=False)