#!/usr/bin/env python3
"""
harvest-detector.py — flag passive TLS collection patterns on a perimeter node.

Reads a NetFlow / Zeek conn.log style stream (one JSON object per line) and emits
flows that look like "harvest now, decrypt later" collection rather than
legitimate use.

Heuristics (all configurable):
  - full TLS handshake observed
  - no session resumption ticket honored
  - small application-data volume after handshake (< 4 KB)
  - flow duration > 2 s
  - peer is not in known-CDN allowlist
  - repeated handshakes from same peer to same SNI without prior resumption

Usage:
    zeek-cut -F '\t' < conn.log | jq -c . | python harvest-detector.py
    cat flows.ndjson | python harvest-detector.py --threshold 3
"""

from __future__ import annotations
import argparse
import json
import sys
from collections import defaultdict, deque
from dataclasses import dataclass
from typing import Iterable


@dataclass
class Flow:
    src: str
    dst: str
    sni: str
    full_handshake: bool
    session_resumed: bool
    app_data_bytes: int
    duration_ms: int
    cipher: str | None = None

    @classmethod
    def from_json(cls, obj: dict) -> "Flow":
        return cls(
            src=obj.get("id.orig_h") or obj["src"],
            dst=obj.get("id.resp_h") or obj["dst"],
            sni=obj.get("server_name") or obj.get("sni") or "",
            full_handshake=bool(obj.get("full_handshake", obj.get("established", True))),
            session_resumed=bool(obj.get("resumed", False)),
            app_data_bytes=int(obj.get("orig_bytes", 0)) + int(obj.get("resp_bytes", 0)),
            duration_ms=int(float(obj.get("duration", 0)) * 1000),
            cipher=obj.get("cipher"),
        )


KNOWN_CDN = {
    "cloudflare", "fastly", "akamai", "amazonaws", "azureedge",
    "cloudfront", "edgecastcdn", "stackpathdns", "vercel",
}


def looks_like_cdn(sni: str) -> bool:
    s = sni.lower()
    return any(k in s for k in KNOWN_CDN)


def looks_like_collection(f: Flow) -> bool:
    if not f.full_handshake:
        return False
    if f.session_resumed:
        return False
    if f.app_data_bytes > 4096:
        return False
    if f.duration_ms < 2000:
        return False
    if looks_like_cdn(f.sni):
        return False
    return True


class Aggregator:
    """
    Track repeated handshakes per (src, sni) pair. After `threshold`
    suspicious handshakes within `window` flows, escalate.
    """

    def __init__(self, threshold: int = 3, window: int = 1024) -> None:
        self.threshold = threshold
        self.window = window
        self.hits: dict[tuple[str, str], deque[int]] = defaultdict(lambda: deque(maxlen=window))
        self.seen = 0

    def update(self, f: Flow) -> str | None:
        self.seen += 1
        if not looks_like_collection(f):
            return None
        key = (f.src, f.sni)
        self.hits[key].append(self.seen)
        if len(self.hits[key]) >= self.threshold:
            return f"{f.src} -> {f.sni}: {len(self.hits[key])} suspicious handshakes (last in window of {self.window})"
        return None


def run(stream: Iterable[str], threshold: int) -> None:
    agg = Aggregator(threshold=threshold)
    for line in stream:
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        try:
            f = Flow.from_json(obj)
        except KeyError:
            continue
        single = looks_like_collection(f)
        agg_hit = agg.update(f)
        if single:
            print(json.dumps({
                "kind": "single",
                "src": f.src, "dst": f.dst, "sni": f.sni,
                "app_data": f.app_data_bytes, "duration_ms": f.duration_ms,
            }))
        if agg_hit:
            print(json.dumps({"kind": "aggregate", "note": agg_hit}))


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--threshold", type=int, default=3,
                   help="suspicious handshakes per (src, sni) before aggregate alert")
    args = p.parse_args()
    run(sys.stdin, args.threshold)


if __name__ == "__main__":
    main()