flask_api.py

#!/usr/bin/env python3
from flask import Flask, request, jsonify
import subprocess, time, os, json, socket
from scapy.utils import RawPcapReader, PcapWriter
from cachetools import TTLCache
import logging

app = Flask(__name__)

# defaults (tune paths as needed)
PCAP = "/home/ubuntu/git/getting_netflow/mitm_output.pcap"
TMP_JSON = "/tmp/ndpi_flows.json"
NDPI = "/usr/bin/ndpiReader"

# cache: keep a generous maxsize; ttl 1h
cache = TTLCache(maxsize=10000, ttl=3600)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)


def send_pcap_to_arkime(filepath: str, host: str, port: int):
    try:
        with socket.create_connection((host, port)) as sock:
            f = sock.makefile('wb')
            writer = PcapWriter(f, linktype=1, sync=True)
            for pkt, meta in RawPcapReader(filepath):
                writer.write(pkt)
                writer.flush()
            f.flush()
            sock.shutdown(socket.SHUT_WR)
    except Exception as e:
        logging.error(f"Failed to send PCAP to Arkime: {e}")
        return False
    return True


@app.route("/update-cache")
def update_cache(force=False):
    """ Run ndpiReader on the given pcap and populate the in-memory cache.
        Key format used: src_ip,src_port,dest_ip,dst_port
    """
    global cache, TMP_JSON, PCAP
    PCAP = request.args.get("file", PCAP)
    TMP_JSON = PCAP.split('.', 1)[0] + ".json"

    try:
        mtime = os.path.getmtime(PCAP)
    except FileNotFoundError:
        logging.error(f"PCAP file not found: {PCAP}")
        cache.clear()
        return jsonify({"error": f"PCAP file not found: {PCAP}"}), 404

    if force or mtime != cache.get("pcap_mtime", 0):
        try:
            cmd = [NDPI, "-i", PCAP, "-F", "-k", TMP_JSON, "-K", "json"]
            subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL)
        except subprocess.CalledProcessError as e:
            logging.error(f"Error running ndpiReader: {e}")
            return jsonify({"error": "Error processing PCAP file with ndpiReader"}), 500

    flows = []
    if os.path.exists(TMP_JSON):
        try:
            with open(TMP_JSON) as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        flows.append(json.loads(line))
                    except json.JSONDecodeError:
                        logging.warning(f"Skipping invalid JSON line")
                        continue
        except Exception as e:
            logging.error(f"Error reading TMP_JSON: {e}")
            return jsonify({"error": "Error reading TMP_JSON"}), 500

    # populate cache (overwrite existing entries for the same key)
    inserted = 0
    for flow in flows:
        si = flow.get("src_ip")
        sp = flow.get("src_port")
        di = flow.get("dest_ip") or flow.get("dst_ip")
        dp = flow.get("dst_port") or flow.get("dest_port")
        key = f"{si},{sp},{di},{dp}"
        if si and sp and di and dp:
            logging.info(f"[CACHE-INSERT] {key}")
            cache[key] = flow
#            print(f"key format {key}")
#            print(f"value format {flow}")
            inserted += 1

    cache["last_updated"] = time.time()
    cache["pcap_mtime"] = mtime
    logging.info(f"[CACHE] inserted {inserted} flows from {TMP_JSON}")
    return jsonify({"message": "Cache updated"}), 200

def extract_app(flow):
    ndpi = flow.get("ndpi", {}) or {}
    category = ndpi.get("category") or "Unknown"  # فقط از category بگیر
    proto = ndpi.get("proto")
    hostname = ndpi.get("hostname")
    domain = ndpi.get("domainame")
    url = (ndpi.get("http") or {}).get("url")

    parts = []
    if proto: parts.append(proto)
    if hostname: parts.append(hostname)
    if domain: parts.append(domain)
    if url: parts.append(url)

    app_protocol = " | ".join(parts) if parts else "Unknown"
    return category, app_protocol


def extract_risks(flow):
    risks_out = []
    risks = (flow.get("ndpi") or {}).get("flow_risk", {}) or {}
    for rid, info in risks.items():
        risk_name = info.get("risk")
        severity = info.get("severity")
        if risk_name and severity:
            risks_out.append(f"{risk_name} ({severity})")
    return "\n".join(risks_out)



@app.route("/enrich")
def enrich():
    src_ip = request.args.get("src_ip")
    src_port = request.args.get("src_port")
    dest_ip = request.args.get("dest_ip")
    dst_port = request.args.get("dst_port")
    #print(f"request.args format : {request.args}") 
    key = f"{src_ip},{src_port},{dest_ip},{dst_port}"
    logging.info(f"[CACHE-LOOKUP] {key}")
    flow = cache.get(key)

    if flow and isinstance(flow, dict):
        logging.info(f"[CACHE-HIT] {key}")
        app_category, app_protocol = extract_app(flow)
        risks = extract_risks(flow)

        # --- Additional fields for dataset enrichment ---
        ndpi_data = flow.get("ndpi", {})
        xfer = flow.get("xfer", {})
        iat = flow.get("iat", {})
        pktlen = flow.get("pktlen", {})
        tcp_flags = flow.get("tcp_flags", {})

        enriched_data = {
            "app_category": app_category,
            "app_protocol": app_protocol,
            "app_risk": risks,

            # --- Flow-level features ---
            "duration": flow.get("duration"),
            "protocol": flow.get("proto"),
            "src_ip": flow.get("src_ip"),
            "src_port": flow.get("src_port"),
            "dst_ip": flow.get("dest_ip"),
            "dst_port": flow.get("dst_port"),

            # --- Transfer statistics ---
            "src2dst_bytes": xfer.get("src2dst_bytes"),
            "dst2src_bytes": xfer.get("dst2src_bytes"),
            "src2dst_packets": xfer.get("src2dst_packets"),
            "dst2src_packets": xfer.get("dst2src_packets"),
            "data_ratio": xfer.get("data_ratio"),

            # --- Inter-arrival time (IAT) stats ---
            "iat_flow_avg": iat.get("flow_avg"),
            "iat_flow_stddev": iat.get("flow_stddev"),
            "iat_c_to_s_avg": iat.get("c_to_s_avg"),
            "iat_s_to_c_avg": iat.get("s_to_c_avg"),

            # --- Packet length statistics ---
            "pktlen_c_to_s_avg": pktlen.get("c_to_s_avg"),
            "pktlen_c_to_s_stddev": pktlen.get("c_to_s_stddev"),
            "pktlen_s_to_c_avg": pktlen.get("s_to_c_avg"),
            "pktlen_s_to_c_stddev": pktlen.get("s_to_c_stddev"),

            # --- TCP flags summary ---
            "tcp_ack_count": tcp_flags.get("ack_count"),
            "tcp_psh_count": tcp_flags.get("psh_count"),
            "tcp_syn_count": tcp_flags.get("syn_count"),
            "tcp_fin_count": tcp_flags.get("fin_count"),

            # --- nDPI metadata ---
            "encrypted": ndpi_data.get("encrypted"),
            "confidence": list(ndpi_data.get("confidence", {}).values())[0] if ndpi_data.get("confidence") else None,
            "breed": ndpi_data.get("breed"),
            "category": ndpi_data.get("category"),
            "proto_by_ip": ndpi_data.get("proto_by_ip"),
            "hostname": ndpi_data.get("hostname"),
        }

        # --- Optionally compute combined risk score ---
        flow_risk = ndpi_data.get("flow_risk", {})
        total_risk_score = sum(r.get("risk_score", {}).get("total", 0) for r in flow_risk.values())
        enriched_data["risk_score_total"] = total_risk_score
        return jsonify(enriched_data)

    logging.info(f"[CACHE-MISS] {key}")
    return jsonify({"error": "no match"}), 200


@app.route("/debug_cache")
def debug_cache():
    """Return lightweight cache stats."""
    return jsonify({
        "flows_count": len([k for k in cache.keys() if k not in ("last_updated","pcap_mtime")]),
        "last_updated": cache.get("last_updated", "N/A")
    })


@app.route("/cache/inspect")
def cache_inspect():
    """ Inspect cache.
        action=keys (default) | flow | search
        - flow requires key parameter; raw=1 returns full flow dict
    """
    def is_meta(k): return k in ("last_updated", "pcap_mtime")

    action = request.args.get("action", "keys")
    limit = min(int(request.args.get("limit", 100)), 2000)
    offset = int(request.args.get("offset", 0))

    if action == "keys":
        prefix = request.args.get("prefix", "")
        keys = [k for k in cache.keys() if isinstance(k, str) and not is_meta(k)]
        if prefix:
            keys = [k for k in keys if k.startswith(prefix)]
        total = len(keys)
        return jsonify({
            "action": "keys",
            "total": total,
            "count": min(len(keys)-offset, limit),
            "keys": keys[offset: offset+limit]
        })

    if action == "flow":
        key = request.args.get("key")
        if not key:
            return jsonify({"error": "missing key parameter for action=flow"}), 400
        flow = cache.get(key)
        if not flow:
            return jsonify({"error": "not found"}), 404
        raw = request.args.get("raw", "0")
        if raw == "1":
            return jsonify({"action": "flow", "key": key, "flow": flow})
        out = {
            "key": key,
            "src_ip": flow.get("src_ip"),
            "src_port": flow.get("src_port"),
            "dest_ip": flow.get("dest_ip") or flow.get("dst_ip"),
            "dst_port": flow.get("dst_port") or flow.get("dest_port"),
            "proto": (flow.get("ndpi") or {}).get("proto") or flow.get("proto"),
        }
        return jsonify({"action": "flow", "key": key, "flow": out})

    if action == "search":
        ip_q = request.args.get("ip", "").strip()
        port_q = request.args.get("port", "").strip()
        proto_q = request.args.get("proto", "").strip().lower()
        matches = []
        for k, flow in cache.items():
            if is_meta(k): continue
            if not isinstance(flow, dict): continue
            s_ip = str(flow.get("src_ip") or "")
            d_ip = str(flow.get("dest_ip") or flow.get("dst_ip") or "")
            s_port = str(flow.get("src_port") or "")
            d_port = str(flow.get("dst_port") or flow.get("dest_port") or "")
            proto = str((flow.get("ndpi") or {}).get("proto") or flow.get("proto") or "")
            if ip_q and not (ip_q in s_ip or ip_q in d_ip): continue
            if port_q and not (port_q in s_port or port_q in d_port): continue
            if proto_q and proto_q not in proto.lower(): continue
            matches.append({
                "key": k,
                "src": f"{s_ip}:{s_port}",
                "dst": f"{d_ip}:{d_port}",
                "proto": proto
            })
        total = len(matches)
        results = matches[offset: offset + limit]
        return jsonify({
            "action": "search",
            "total": total,
            "count": len(results),
            "results": results
        })

    return jsonify({"error": "unknown action"}), 400


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000)