#!/usr/bin/env python3
"""C2 beacon detection through network traffic timing analysis.

Analyzes connection timing patterns in PCAP files to identify regular-interval
communications indicative of command-and-control beaconing.

Usage:
    python beacon_detector.py --pcap capture.pcap --output beacons.json
    python beacon_detector.py --pcap capture.pcap --min-connections 5 --max-jitter 0.3
"""

from __future__ import annotations

import argparse
import collections
import json
import logging
import math
import os
import platform
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

try:
    import dpkt
    HAS_DPKT = True
except ImportError:
    HAS_DPKT = False

try:
    from scapy.all import PcapReader, IP, IPv6, TCP, UDP, conf
    conf.verb = 0
    HAS_SCAPY = True
except ImportError:
    HAS_SCAPY = False


def is_private_ip(ip: str) -> bool:
    """Check if an IP address is in a private/reserved range."""
    parts = ip.split(".")
    if len(parts) != 4:
        return ":" in ip  # Treat IPv6 as potentially internal
    try:
        first = int(parts[0])
        second = int(parts[1])
    except ValueError:
        return False

    if first == 10:
        return True
    if first == 172 and 16 <= second <= 31:
        return True
    if first == 192 and second == 168:
        return True
    if first == 127:
        return True
    if first == 169 and second == 254:
        return True
    return False


class BeaconDetector:
    """Detects C2 beaconing patterns in network traffic."""

    def __init__(
        self,
        min_connections: int = 10,
        max_jitter: float = 0.3,
        min_interval: float = 5.0,
        max_interval: float = 86400.0,
    ):
        self.min_connections = min_connections
        self.max_jitter = max_jitter
        self.min_interval = min_interval  # Minimum beacon interval in seconds
        self.max_interval = max_interval  # Maximum beacon interval in seconds
        self.connection_times: dict[str, list[float]] = collections.defaultdict(list)
        self.connection_sizes: dict[str, list[int]] = collections.defaultdict(list)
        self.packet_count = 0

    def parse_pcap_dpkt(self, pcap_path: str) -> None:
        """Extract connection timing from PCAP using dpkt."""
        import socket as _socket
        logger.info("Parsing PCAP with dpkt...")

        with open(pcap_path, "rb") as f:
            try:
                pcap = dpkt.pcap.Reader(f)
            except ValueError:
                f.seek(0)
                pcap = dpkt.pcapng.Reader(f)

            for timestamp, buf in pcap:
                self.packet_count += 1
                try:
                    eth = dpkt.ethernet.Ethernet(buf)
                    if not isinstance(eth.data, dpkt.ip.IP):
                        continue

                    ip = eth.data
                    src_ip = _socket.inet_ntoa(ip.src)
                    dst_ip = _socket.inet_ntoa(ip.dst)
                    pkt_len = len(buf)

                    if isinstance(ip.data, dpkt.tcp.TCP):
                        tcp = ip.data
                        # Track SYN packets (new connections) and data packets
                        if tcp.flags & dpkt.tcp.TH_SYN or len(tcp.data) > 0:
                            key = f"{src_ip}->{dst_ip}:{tcp.dport}"
                            self.connection_times[key].append(timestamp)
                            self.connection_sizes[key].append(pkt_len)

                    elif isinstance(ip.data, dpkt.udp.UDP):
                        udp = ip.data
                        if udp.dport != 53 and udp.sport != 53:  # Skip DNS
                            key = f"{src_ip}->{dst_ip}:{udp.dport}/udp"
                            self.connection_times[key].append(timestamp)
                            self.connection_sizes[key].append(pkt_len)

                except Exception:
                    continue

                if self.packet_count % 50000 == 0:
                    logger.info(f"Processed {self.packet_count} packets...")

        logger.info(f"Total packets processed: {self.packet_count}")
        logger.info(f"Unique connection pairs: {len(self.connection_times)}")

    def parse_pcap_scapy(self, pcap_path: str) -> None:
        """Extract connection timing from PCAP using Scapy."""
        logger.info("Parsing PCAP with Scapy...")

        with PcapReader(pcap_path) as reader:
            for pkt in reader:
                self.packet_count += 1
                ts = float(pkt.time) if hasattr(pkt, 'time') else 0

                if pkt.haslayer(IP):
                    src_ip = pkt[IP].src
                    dst_ip = pkt[IP].dst
                elif pkt.haslayer(IPv6):
                    src_ip = pkt[IPv6].src
                    dst_ip = pkt[IPv6].dst
                else:
                    continue

                pkt_len = len(pkt)

                if pkt.haslayer(TCP):
                    tcp = pkt[TCP]
                    if tcp.flags & 0x02 or len(bytes(tcp.payload)) > 0:
                        key = f"{src_ip}->{dst_ip}:{tcp.dport}"
                        self.connection_times[key].append(ts)
                        self.connection_sizes[key].append(pkt_len)

                elif pkt.haslayer(UDP):
                    udp = pkt[UDP]
                    if udp.dport != 53 and udp.sport != 53:
                        key = f"{src_ip}->{dst_ip}:{udp.dport}/udp"
                        self.connection_times[key].append(ts)
                        self.connection_sizes[key].append(pkt_len)

                if self.packet_count % 50000 == 0:
                    logger.info(f"Processed {self.packet_count} packets...")

        logger.info(f"Total packets processed: {self.packet_count}")
        logger.info(f"Unique connection pairs: {len(self.connection_times)}")

    def analyze_timing(self, times: list[float]) -> dict[str, Any]:
        """Analyze timing pattern of a connection series."""
        sorted_times = sorted(times)
        deltas = [sorted_times[i + 1] - sorted_times[i]
                  for i in range(len(sorted_times) - 1)]

        if not deltas:
            return {}

        mean = sum(deltas) / len(deltas)
        if mean <= 0:
            return {}

        variance = sum((d - mean) ** 2 for d in deltas) / len(deltas)
        std_dev = math.sqrt(variance)
        jitter = std_dev / mean

        # Median
        sorted_deltas = sorted(deltas)
        mid = len(sorted_deltas) // 2
        median = (
            sorted_deltas[mid]
            if len(sorted_deltas) % 2 == 1
            else (sorted_deltas[mid - 1] + sorted_deltas[mid]) / 2
        )

        # Mode (approximate - bucket into 1-second bins)
        bins = collections.Counter(int(d) for d in deltas)
        mode_bin = bins.most_common(1)[0][0] if bins else 0

        # Percentiles
        p5 = sorted_deltas[max(0, int(len(sorted_deltas) * 0.05))]
        p95 = sorted_deltas[min(len(sorted_deltas) - 1, int(len(sorted_deltas) * 0.95))]

        # Skewness
        skewness = 0.0
        if std_dev > 0:
            skewness = sum((d - mean) ** 3 for d in deltas) / (len(deltas) * std_dev ** 3)

        return {
            "count": len(times),
            "delta_count": len(deltas),
            "mean_interval": round(mean, 2),
            "median_interval": round(median, 2),
            "mode_interval": mode_bin,
            "std_dev": round(std_dev, 2),
            "jitter_ratio": round(jitter, 4),
            "min_delta": round(min(deltas), 2),
            "max_delta": round(max(deltas), 2),
            "p5_delta": round(p5, 2),
            "p95_delta": round(p95, 2),
            "skewness": round(skewness, 4),
            "duration_seconds": round(sorted_times[-1] - sorted_times[0], 2),
            "first_seen": sorted_times[0],
            "last_seen": sorted_times[-1],
        }

    def analyze_sizes(self, sizes: list[int]) -> dict[str, Any]:
        """Analyze packet sizes for a connection."""
        if not sizes:
            return {}

        mean = sum(sizes) / len(sizes)
        variance = sum((s - mean) ** 2 for s in sizes) / len(sizes)
        std_dev = math.sqrt(variance)

        return {
            "mean_size": round(mean, 1),
            "std_dev_size": round(std_dev, 1),
            "min_size": min(sizes),
            "max_size": max(sizes),
            "size_jitter": round(std_dev / mean, 4) if mean > 0 else 0,
        }

    def classify_beacon(self, timing: dict, sizes: dict) -> dict[str, Any]:
        """Classify the confidence level and characteristics of a potential beacon."""
        jitter = timing.get("jitter_ratio", float("inf"))
        interval = timing.get("mean_interval", 0)
        count = timing.get("count", 0)
        duration = timing.get("duration_seconds", 0)

        # Confidence scoring
        confidence_score = 0
        indicators = []

        # Low jitter = strong beacon indicator
        if jitter < 0.05:
            confidence_score += 3
            indicators.append("very low jitter (<5%)")
        elif jitter < 0.1:
            confidence_score += 2
            indicators.append("low jitter (<10%)")
        elif jitter < 0.2:
            confidence_score += 1
            indicators.append("moderate jitter (<20%)")

        # Regular interval within expected C2 range
        if self.min_interval <= interval <= self.max_interval:
            confidence_score += 1
            indicators.append(f"interval in typical C2 range ({interval:.0f}s)")

        # Long duration = persistent beacon
        if duration > 3600:
            confidence_score += 1
            indicators.append(f"long duration ({duration / 3600:.1f} hours)")

        # Many connections
        if count > 50:
            confidence_score += 1
            indicators.append(f"high connection count ({count})")

        # Consistent packet sizes (another beacon indicator)
        if sizes and sizes.get("size_jitter", 1) < 0.15:
            confidence_score += 1
            indicators.append("consistent packet sizes")

        # Common beacon intervals (exact minutes/hours)
        for common in [30, 60, 120, 300, 600, 900, 1800, 3600]:
            if abs(interval - common) / common < 0.1:
                confidence_score += 1
                indicators.append(f"near common interval ({common}s)")
                break

        if confidence_score >= 5:
            confidence = "high"
        elif confidence_score >= 3:
            confidence = "medium"
        else:
            confidence = "low"

        return {
            "confidence": confidence,
            "confidence_score": confidence_score,
            "indicators": indicators,
        }

    def detect(self) -> list[dict[str, Any]]:
        """Run beacon detection on all tracked connections."""
        beacons = []

        for key, times in self.connection_times.items():
            if len(times) < self.min_connections:
                continue

            timing = self.analyze_timing(times)
            if not timing:
                continue

            # Filter by jitter threshold
            if timing["jitter_ratio"] > self.max_jitter:
                continue

            # Filter by interval range
            if not (self.min_interval <= timing["mean_interval"] <= self.max_interval):
                continue

            sizes = self.analyze_sizes(self.connection_sizes.get(key, []))
            classification = self.classify_beacon(timing, sizes)

            # Parse connection key
            parts = key.split("->")
            src_ip = parts[0]
            dst_parts = parts[1] if len(parts) > 1 else ""

            beacon = {
                "connection": key,
                "src_ip": src_ip,
                "dst": dst_parts,
                "is_external": not is_private_ip(dst_parts.split(":")[0]) if dst_parts else False,
                "timing": timing,
                "packet_sizes": sizes,
                **classification,
            }

            beacons.append(beacon)

        # Sort by confidence score (descending), then jitter (ascending)
        beacons.sort(
            key=lambda x: (-x["confidence_score"], x["timing"]["jitter_ratio"])
        )

        return beacons

    def generate_report(self) -> dict[str, Any]:
        """Generate full beacon detection report."""
        beacons = self.detect()

        # Categorize
        high_confidence = [b for b in beacons if b["confidence"] == "high"]
        medium_confidence = [b for b in beacons if b["confidence"] == "medium"]
        low_confidence = [b for b in beacons if b["confidence"] == "low"]
        external_beacons = [b for b in beacons if b.get("is_external")]

        return {
            "metadata": {
                "tool": "beacon_detector.py",
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "platform": platform.system(),
                "parameters": {
                    "min_connections": self.min_connections,
                    "max_jitter": self.max_jitter,
                    "min_interval": self.min_interval,
                    "max_interval": self.max_interval,
                },
                "total_packets": self.packet_count,
                "unique_connections": len(self.connection_times),
            },
            "summary": {
                "total_beacons_detected": len(beacons),
                "high_confidence": len(high_confidence),
                "medium_confidence": len(medium_confidence),
                "low_confidence": len(low_confidence),
                "external_beacons": len(external_beacons),
            },
            "beacons": beacons,
        }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Detect C2 beaconing in PCAP files via timing analysis",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --pcap capture.pcap --output beacons.json
  %(prog)s --pcap capture.pcap --min-connections 5 --max-jitter 0.3
  %(prog)s --pcap capture.pcap --min-interval 10 --max-interval 7200
        """,
    )
    parser.add_argument("--input", "--pcap", "-f", dest="pcap", required=True, help="Path to PCAP/PCAPNG file")
    parser.add_argument("--output", "-o", help="Output JSON report path (default: stdout)")
    parser.add_argument(
        "--min-connections", type=int, default=10,
        help="Minimum connections for beacon candidate (default: 10)",
    )
    parser.add_argument(
        "--max-jitter", type=float, default=0.3,
        help="Maximum jitter ratio (std_dev/mean) (default: 0.3)",
    )
    parser.add_argument(
        "--min-interval", type=float, default=5.0,
        help="Minimum beacon interval in seconds (default: 5)",
    )
    parser.add_argument(
        "--max-interval", type=float, default=86400.0,
        help="Maximum beacon interval in seconds (default: 86400)",
    )
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    if not os.path.isfile(args.pcap):
        logger.error(f"PCAP file not found: {args.pcap}")
        sys.exit(1)

    detector = BeaconDetector(
        min_connections=args.min_connections,
        max_jitter=args.max_jitter,
        min_interval=args.min_interval,
        max_interval=args.max_interval,
    )

    # Parse with best available library
    if HAS_DPKT:
        try:
            detector.parse_pcap_dpkt(args.pcap)
        except Exception as e:
            logger.warning(f"dpkt failed: {e}")
            if HAS_SCAPY:
                detector.parse_pcap_scapy(args.pcap)
            else:
                sys.exit(1)
    elif HAS_SCAPY:
        detector.parse_pcap_scapy(args.pcap)
    else:
        logger.error("Install scapy or dpkt: pip install scapy dpkt")
        sys.exit(1)

    report = detector.generate_report()
    report_json = json.dumps(report, indent=2, default=str)

    if args.output:
        with open(args.output, "w") as f:
            f.write(report_json)
        logger.info(f"Report written to: {args.output}")
    else:
        print(report_json)

    # Summary
    s = report["summary"]
    print(f"\n=== Beacon Detection Summary ===", file=sys.stderr)
    print(f"Connections analyzed: {report['metadata']['unique_connections']}", file=sys.stderr)
    print(f"Beacons detected: {s['total_beacons_detected']}", file=sys.stderr)
    print(f"  High confidence: {s['high_confidence']}", file=sys.stderr)
    print(f"  Medium confidence: {s['medium_confidence']}", file=sys.stderr)
    print(f"  Low confidence: {s['low_confidence']}", file=sys.stderr)
    print(f"  External destinations: {s['external_beacons']}", file=sys.stderr)

    if report["beacons"]:
        print("\nTop beacon candidates:", file=sys.stderr)
        for b in report["beacons"][:10]:
            t = b["timing"]
            print(
                f"  [{b['confidence'].upper()}] {b['connection']} "
                f"interval={t['mean_interval']}s jitter={t['jitter_ratio']} "
                f"count={t['count']}",
                file=sys.stderr,
            )


if __name__ == "__main__":
    main()