#!/usr/bin/env python3
"""Agent for detecting C2 beaconing through network traffic frequency analysis."""

import argparse
import json
import math
from collections import defaultdict
from datetime import datetime, timezone


def parse_zeek_conn_log(log_path):
    """Parse Zeek conn.log and extract connection timestamps per src-dst pair."""
    connections = defaultdict(list)
    try:
        with open(log_path, "r") as f:
            for line in f:
                if line.startswith("#"):
                    continue
                fields = line.strip().split("\t")
                if len(fields) < 7:
                    continue
                ts = float(fields[0])
                src, dst = fields[2], fields[4]
                dst_port = fields[5]
                key = f"{src}->{dst}:{dst_port}"
                connections[key].append(ts)
    except (FileNotFoundError, ValueError):
        pass
    return connections


def calculate_jitter(intervals):
    """Calculate jitter (standard deviation of intervals)."""
    if len(intervals) < 2:
        return 0
    mean = sum(intervals) / len(intervals)
    variance = sum((x - mean) ** 2 for x in intervals) / len(intervals)
    return math.sqrt(variance)


def detect_beaconing(connections, min_connections=10, max_jitter_percent=15):
    """Detect beaconing patterns based on interval regularity."""
    beacons = []
    for key, timestamps in connections.items():
        if len(timestamps) < min_connections:
            continue
        timestamps.sort()
        intervals = [timestamps[i+1] - timestamps[i] for i in range(len(timestamps)-1)]
        if not intervals:
            continue
        mean_interval = sum(intervals) / len(intervals)
        if mean_interval == 0:
            continue
        jitter = calculate_jitter(intervals)
        jitter_percent = (jitter / mean_interval) * 100

        if jitter_percent <= max_jitter_percent:
            parts = key.split("->")
            src = parts[0]
            dst_port = parts[1] if len(parts) > 1 else ""
            beacons.append({
                "flow": key,
                "connection_count": len(timestamps),
                "mean_interval_seconds": round(mean_interval, 2),
                "jitter_seconds": round(jitter, 2),
                "jitter_percent": round(jitter_percent, 2),
                "duration_hours": round((timestamps[-1] - timestamps[0]) / 3600, 2),
                "confidence": "HIGH" if jitter_percent < 5 else "MEDIUM",
            })
    return sorted(beacons, key=lambda x: x["jitter_percent"])


def parse_csv_log(csv_path):
    """Parse generic CSV log with timestamp, src, dst, port columns."""
    connections = defaultdict(list)
    try:
        import csv
        with open(csv_path, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                ts = row.get("timestamp") or row.get("ts") or row.get("time")
                src = row.get("src") or row.get("source") or row.get("src_ip")
                dst = row.get("dst") or row.get("destination") or row.get("dst_ip")
                port = row.get("dst_port") or row.get("port") or ""
                if ts and src and dst:
                    try:
                        ts_float = float(ts)
                    except ValueError:
                        from datetime import datetime as dt
                        try:
                            ts_float = dt.fromisoformat(ts.replace("Z", "+00:00")).timestamp()
                        except ValueError:
                            continue
                    connections[f"{src}->{dst}:{port}"].append(ts_float)
    except (FileNotFoundError, KeyError):
        pass
    return connections


def main():
    parser = argparse.ArgumentParser(
        description="Detect C2 beaconing via frequency analysis"
    )
    parser.add_argument("--conn-log", help="Zeek conn.log path")
    parser.add_argument("--csv", help="CSV log with timestamp, src, dst columns")
    parser.add_argument("--min-connections", type=int, default=10)
    parser.add_argument("--max-jitter", type=float, default=15, help="Max jitter percent")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Beaconing Detection via Frequency Analysis")
    connections = {}

    if args.conn_log:
        connections = parse_zeek_conn_log(args.conn_log)
    elif args.csv:
        connections = parse_csv_log(args.csv)

    print(f"[*] Unique flows: {len(connections)}")
    beacons = detect_beaconing(connections, args.min_connections, args.max_jitter)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "flows_analyzed": len(connections),
        "beacons_detected": len(beacons),
        "beacons": beacons[:50],
        "risk_level": "CRITICAL" if beacons else "LOW",
    }
    print(f"[*] Beacons detected: {len(beacons)}")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"[*] Report saved to {args.output}")
    else:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()