#!/usr/bin/env python3
"""Summarize a PCAP and detect beaconing.

Prefers ``pyshark`` (a tshark wrapper) for rich parsing; falls back to a minimal pure-Python
parser that reads the classic pcap format and tallies IPv4 conversations. Analyzes captured
traffic only — it never generates traffic.
"""
from __future__ import annotations

import argparse
import json
import statistics
import struct
import sys
from collections import defaultdict


def _beaconing(timestamps_by_dst: dict[str, list[float]]) -> list[dict]:
    results = []
    for dst, times in timestamps_by_dst.items():
        if len(times) < 4:
            continue
        times = sorted(times)
        deltas = [t2 - t1 for t1, t2 in zip(times, times[1:])]
        mean = statistics.mean(deltas)
        if mean <= 0:
            continue
        stdev = statistics.pstdev(deltas)
        jitter = stdev / mean
        if jitter < 0.25:  # regular interval = likely beacon
            results.append({"dst": dst, "count": len(times),
                            "interval_s": round(mean, 2), "jitter": round(jitter, 3)})
    return sorted(results, key=lambda r: r["jitter"])


def summarize_with_pyshark(path: str) -> dict:
    import pyshark  # type: ignore
    cap = pyshark.FileCapture(path, keep_packets=False)
    convos: dict[tuple, int] = defaultdict(int)
    times_by_dst: dict[str, list[float]] = defaultdict(list)
    dns: set[str] = set()
    sni: set[str] = set()
    for pkt in cap:
        try:
            if hasattr(pkt, "ip"):
                src, dst = pkt.ip.src, pkt.ip.dst
                convos[(src, dst)] += 1
                times_by_dst[dst].append(float(pkt.sniff_timestamp))
            if hasattr(pkt, "dns") and hasattr(pkt.dns, "qry_name"):
                dns.add(pkt.dns.qry_name)
            if hasattr(pkt, "tls") and hasattr(pkt.tls, "handshake_extensions_server_name"):
                sni.add(pkt.tls.handshake_extensions_server_name)
        except AttributeError:
            continue
    cap.close()
    return {
        "engine": "pyshark",
        "top_conversations": sorted(
            [{"src": s, "dst": d, "packets": n} for (s, d), n in convos.items()],
            key=lambda c: c["packets"], reverse=True)[:20],
        "dns_queries": sorted(dns),
        "tls_sni": sorted(sni),
        "beaconing": _beaconing(times_by_dst),
    }


def summarize_minimal(path: str) -> dict:
    """Minimal classic-pcap parser: IPv4 conversations and per-dst timestamps."""
    with open(path, "rb") as fh:
        data = fh.read()
    if len(data) < 24:
        raise ValueError("file too small to be a pcap")
    magic = data[:4]
    if magic in (b"\xd4\xc3\xb2\xa1", b"\xa1\xb2\xc3\xd4"):
        endian = "<" if magic == b"\xd4\xc3\xb2\xa1" else ">"
    else:
        raise ValueError("not a classic pcap (try pyshark for pcapng)")
    off = 24
    convos: dict[tuple, int] = defaultdict(int)
    times_by_dst: dict[str, list[float]] = defaultdict(list)
    while off + 16 <= len(data):
        ts_sec, ts_usec, incl, _orig = struct.unpack_from(endian + "IIII", data, off)
        off += 16
        pkt = data[off:off + incl]
        off += incl
        if len(pkt) < 14 + 20:
            continue
        eth_type = struct.unpack_from(">H", pkt, 12)[0]
        if eth_type != 0x0800:  # IPv4 only
            continue
        ip = pkt[14:]
        src = ".".join(str(b) for b in ip[12:16])
        dst = ".".join(str(b) for b in ip[16:20])
        convos[(src, dst)] += 1
        times_by_dst[dst].append(ts_sec + ts_usec / 1e6)
    return {
        "engine": "minimal",
        "top_conversations": sorted(
            [{"src": s, "dst": d, "packets": n} for (s, d), n in convos.items()],
            key=lambda c: c["packets"], reverse=True)[:20],
        "beaconing": _beaconing(times_by_dst),
        "note": "install pyshark for DNS/TLS/HTTP detail and pcapng support",
    }


def summarize(path: str) -> dict:
    try:
        return summarize_with_pyshark(path)
    except ImportError:
        return summarize_minimal(path)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Summarize malware PCAP")
    sub = parser.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("summarize")
    p.add_argument("pcap")
    args = parser.parse_args(argv)
    print(json.dumps(summarize(args.pcap), indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())