#!/usr/bin/env python3
"""PCAP network traffic analyzer for malware analysis.

Parses PCAP/PCAPNG files and extracts security-relevant information including
DNS queries, HTTP requests, TLS handshakes, connection summaries, potential
C2 beaconing, and carved files.

Usage:
    python pcap_analyzer.py --pcap capture.pcap --output report.json
    python pcap_analyzer.py --pcap capture.pcap --extract-files ./carved/
    python pcap_analyzer.py --pcap capture.pcap --dns-only
"""

from __future__ import annotations

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

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

try:
    from scapy.all import (
        rdpcap, PcapReader, IP, IPv6, TCP, UDP, DNS, DNSQR, DNSRR,
        Raw, conf,
    )
    conf.verb = 0
    HAS_SCAPY = True
except ImportError:
    HAS_SCAPY = False

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


def entropy(data: bytes) -> float:
    """Calculate Shannon entropy of byte data."""
    if not data:
        return 0.0
    freq = collections.Counter(data)
    length = len(data)
    return -sum(
        (count / length) * math.log2(count / length)
        for count in freq.values()
    )


def defang_ioc(value: str) -> str:
    """Defang an IOC for safe display."""
    value = value.replace("http://", "hxxp://")
    value = value.replace("https://", "hxxps://")
    value = value.replace(".", "[.]")
    return value


class PcapAnalyzer:
    """Analyzes PCAP files for malware-related network indicators."""

    def __init__(self, pcap_path: str):
        self.pcap_path = pcap_path
        self.dns_queries: list[dict] = []
        self.dns_responses: list[dict] = []
        self.http_requests: list[dict] = []
        self.http_responses: list[dict] = []
        self.tls_handshakes: list[dict] = []
        self.connections: list[dict] = []
        self.connection_times: dict[str, list[float]] = collections.defaultdict(list)
        self.file_objects: list[dict] = []
        self.packet_count = 0
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None

    def analyze_with_scapy(self) -> None:
        """Parse PCAP using Scapy (slower but more flexible)."""
        logger.info("Analyzing with Scapy...")
        try:
            # Use PcapReader for memory efficiency on large files
            with PcapReader(self.pcap_path) as reader:
                for pkt in reader:
                    self.packet_count += 1
                    ts = float(pkt.time) if hasattr(pkt, 'time') else 0

                    if self.start_time is None or ts < self.start_time:
                        self.start_time = ts
                    if self.end_time is None or ts > self.end_time:
                        self.end_time = ts

                    self._process_scapy_packet(pkt, ts)

                    if self.packet_count % 10000 == 0:
                        logger.info(f"Processed {self.packet_count} packets...")
        except Exception as e:
            logger.error(f"Scapy analysis failed: {e}")
            raise

    def _process_scapy_packet(self, pkt, timestamp: float):
        """Process a single Scapy packet."""
        # Extract IP layer
        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:
            return

        # DNS
        if pkt.haslayer(DNS):
            self._process_dns_scapy(pkt, src_ip, dst_ip, timestamp)

        # TCP connections
        if pkt.haslayer(TCP):
            src_port = pkt[TCP].sport
            dst_port = pkt[TCP].dport
            flags = pkt[TCP].flags

            # Track connection timing for beacon detection
            if flags & 0x02:  # SYN flag
                key = f"{src_ip}:{src_port}->{dst_ip}:{dst_port}"
                self.connection_times[f"{src_ip}->{dst_ip}:{dst_port}"].append(timestamp)

            # HTTP detection (basic)
            if pkt.haslayer(Raw):
                payload = bytes(pkt[Raw].load)
                if dst_port == 80 or payload.startswith(b"GET ") or payload.startswith(b"POST "):
                    self._process_http_request(payload, src_ip, dst_ip, dst_port, timestamp)
                elif src_port == 80 or payload.startswith(b"HTTP/"):
                    self._process_http_response(payload, src_ip, dst_ip, src_port, timestamp)

            # TLS detection
            if pkt.haslayer(Raw):
                payload = bytes(pkt[Raw].load)
                if len(payload) > 5 and payload[0] == 0x16:  # TLS Handshake
                    self._process_tls(payload, src_ip, dst_ip, dst_port, timestamp)

        # UDP connections
        if pkt.haslayer(UDP):
            src_port = pkt[UDP].sport
            dst_port = pkt[UDP].dport
            key = f"{src_ip}->{dst_ip}:{dst_port}/udp"
            self.connection_times[key].append(timestamp)

    def _process_dns_scapy(self, pkt, src_ip, dst_ip, timestamp):
        """Extract DNS query/response information."""
        dns = pkt[DNS]
        if dns.qr == 0 and dns.haslayer(DNSQR):  # Query
            qname = dns[DNSQR].qname.decode("utf-8", errors="replace").rstrip(".")
            qtype_map = {1: "A", 2: "NS", 5: "CNAME", 15: "MX", 16: "TXT", 28: "AAAA", 33: "SRV"}
            qtype = qtype_map.get(dns[DNSQR].qtype, str(dns[DNSQR].qtype))
            self.dns_queries.append({
                "timestamp": timestamp,
                "src_ip": src_ip,
                "dst_ip": dst_ip,
                "query": qname,
                "type": qtype,
            })
        elif dns.qr == 1:  # Response
            qname = dns[DNSQR].qname.decode("utf-8", errors="replace").rstrip(".") if dns.haslayer(DNSQR) else ""
            answers = []
            for i in range(dns.ancount):
                try:
                    rr = dns.an[i]
                    if hasattr(rr, "rdata"):
                        rdata = rr.rdata
                        if isinstance(rdata, bytes):
                            rdata = rdata.decode("utf-8", errors="replace")
                        answers.append(str(rdata))
                except Exception:
                    pass
            rcode_map = {0: "NOERROR", 1: "FORMERR", 2: "SERVFAIL", 3: "NXDOMAIN", 5: "REFUSED"}
            self.dns_responses.append({
                "timestamp": timestamp,
                "src_ip": src_ip,
                "query": qname,
                "answers": answers,
                "rcode": rcode_map.get(dns.rcode, str(dns.rcode)),
            })

    def _process_http_request(self, payload, src_ip, dst_ip, dst_port, timestamp):
        """Parse HTTP request."""
        try:
            lines = payload.split(b"\r\n")
            if not lines:
                return
            request_line = lines[0].decode("utf-8", errors="replace")
            parts = request_line.split(" ", 2)
            if len(parts) < 2:
                return

            method = parts[0]
            uri = parts[1]
            headers = {}
            for line in lines[1:]:
                if b":" in line:
                    key, _, value = line.partition(b":")
                    headers[key.decode("utf-8", errors="replace").strip().lower()] = \
                        value.decode("utf-8", errors="replace").strip()

            self.http_requests.append({
                "timestamp": timestamp,
                "src_ip": src_ip,
                "dst_ip": dst_ip,
                "dst_port": dst_port,
                "method": method,
                "uri": uri,
                "host": headers.get("host", dst_ip),
                "user_agent": headers.get("user-agent", ""),
                "content_type": headers.get("content-type", ""),
                "content_length": headers.get("content-length", ""),
            })
        except Exception:
            pass

    def _process_http_response(self, payload, src_ip, dst_ip, src_port, timestamp):
        """Parse HTTP response."""
        try:
            lines = payload.split(b"\r\n")
            if not lines:
                return
            status_line = lines[0].decode("utf-8", errors="replace")
            parts = status_line.split(" ", 2)
            if len(parts) < 2:
                return

            status_code = int(parts[1]) if parts[1].isdigit() else 0
            headers = {}
            for line in lines[1:]:
                if b":" in line:
                    key, _, value = line.partition(b":")
                    headers[key.decode("utf-8", errors="replace").strip().lower()] = \
                        value.decode("utf-8", errors="replace").strip()

            self.http_responses.append({
                "timestamp": timestamp,
                "src_ip": src_ip,
                "dst_ip": dst_ip,
                "status_code": status_code,
                "content_type": headers.get("content-type", ""),
                "content_length": headers.get("content-length", ""),
                "server": headers.get("server", ""),
            })
        except Exception:
            pass

    def _process_tls(self, payload, src_ip, dst_ip, dst_port, timestamp):
        """Extract TLS handshake information."""
        try:
            if len(payload) < 6:
                return

            content_type = payload[0]
            tls_version_major = payload[1]
            tls_version_minor = payload[2]

            if content_type != 0x16:  # Not a handshake
                return

            # Parse handshake type
            if len(payload) > 5:
                handshake_type = payload[5]
                handshake_types = {
                    1: "ClientHello",
                    2: "ServerHello",
                    11: "Certificate",
                    12: "ServerKeyExchange",
                    16: "ClientKeyExchange",
                }

                version_map = {
                    (3, 1): "TLS 1.0",
                    (3, 2): "TLS 1.1",
                    (3, 3): "TLS 1.2",
                    (3, 4): "TLS 1.3",
                }
                version = version_map.get(
                    (tls_version_major, tls_version_minor),
                    f"Unknown ({tls_version_major}.{tls_version_minor})",
                )

                # Extract SNI from ClientHello
                sni = ""
                if handshake_type == 1 and len(payload) > 43:
                    sni = self._extract_sni(payload)

                self.tls_handshakes.append({
                    "timestamp": timestamp,
                    "src_ip": src_ip,
                    "dst_ip": dst_ip,
                    "dst_port": dst_port,
                    "version": version,
                    "handshake_type": handshake_types.get(handshake_type, str(handshake_type)),
                    "sni": sni,
                })
        except Exception:
            pass

    def _extract_sni(self, payload: bytes) -> str:
        """Extract Server Name Indication from TLS ClientHello."""
        try:
            # Skip TLS record header (5) + handshake header (4) + version (2) + random (32)
            offset = 5 + 4 + 2 + 32

            if offset >= len(payload):
                return ""

            # Session ID length
            session_id_len = payload[offset]
            offset += 1 + session_id_len

            if offset + 2 >= len(payload):
                return ""

            # Cipher suites length
            cipher_suites_len = struct.unpack("!H", payload[offset:offset + 2])[0]
            offset += 2 + cipher_suites_len

            if offset >= len(payload):
                return ""

            # Compression methods length
            comp_methods_len = payload[offset]
            offset += 1 + comp_methods_len

            if offset + 2 >= len(payload):
                return ""

            # Extensions length
            extensions_len = struct.unpack("!H", payload[offset:offset + 2])[0]
            offset += 2

            # Parse extensions
            end = min(offset + extensions_len, len(payload))
            while offset + 4 < end:
                ext_type = struct.unpack("!H", payload[offset:offset + 2])[0]
                ext_len = struct.unpack("!H", payload[offset + 2:offset + 4])[0]
                offset += 4

                if ext_type == 0 and offset + ext_len <= end:  # SNI extension
                    # Skip SNI list length (2) + type (1) + name length (2)
                    if offset + 5 < end:
                        name_len = struct.unpack("!H", payload[offset + 3:offset + 5])[0]
                        if offset + 5 + name_len <= end:
                            return payload[offset + 5:offset + 5 + name_len].decode("ascii", errors="replace")

                offset += ext_len

        except Exception:
            pass
        return ""

    def analyze_with_dpkt(self) -> None:
        """Parse PCAP using dpkt (faster for large files)."""
        logger.info("Analyzing with dpkt...")
        try:
            with open(self.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
                    if self.start_time is None or timestamp < self.start_time:
                        self.start_time = timestamp
                    if self.end_time is None or timestamp > self.end_time:
                        self.end_time = timestamp

                    try:
                        eth = dpkt.ethernet.Ethernet(buf)
                        if isinstance(eth.data, dpkt.ip.IP):
                            ip = eth.data
                            src_ip = self._ip_to_str(ip.src)
                            dst_ip = self._ip_to_str(ip.dst)

                            if isinstance(ip.data, dpkt.tcp.TCP):
                                tcp = ip.data
                                if tcp.flags & dpkt.tcp.TH_SYN:
                                    key = f"{src_ip}->{dst_ip}:{tcp.dport}"
                                    self.connection_times[key].append(timestamp)

                                if tcp.data:
                                    self._process_dpkt_tcp(tcp.data, src_ip, dst_ip,
                                                          tcp.sport, tcp.dport, timestamp)

                            elif isinstance(ip.data, dpkt.udp.UDP):
                                udp = ip.data
                                if udp.data:
                                    self._process_dpkt_udp(udp.data, src_ip, dst_ip,
                                                          udp.sport, udp.dport, timestamp)

                    except Exception:
                        continue

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

        except Exception as e:
            logger.error(f"dpkt analysis failed: {e}")
            raise

    @staticmethod
    def _ip_to_str(ip_bytes: bytes) -> str:
        """Convert IP bytes to string."""
        import socket
        try:
            if len(ip_bytes) == 4:
                return socket.inet_ntoa(ip_bytes)
            elif len(ip_bytes) == 16:
                return socket.inet_ntop(socket.AF_INET6, ip_bytes)
        except Exception:
            pass
        return ".".join(str(b) for b in ip_bytes)

    def _process_dpkt_tcp(self, data, src_ip, dst_ip, sport, dport, timestamp):
        """Process TCP payload with dpkt."""
        # HTTP
        if dport == 80 or data[:4] in (b"GET ", b"POST", b"HEAD", b"PUT "):
            try:
                req = dpkt.http.Request(data)
                self.http_requests.append({
                    "timestamp": timestamp,
                    "src_ip": src_ip,
                    "dst_ip": dst_ip,
                    "dst_port": dport,
                    "method": req.method,
                    "uri": req.uri,
                    "host": req.headers.get("host", dst_ip),
                    "user_agent": req.headers.get("user-agent", ""),
                    "content_type": req.headers.get("content-type", ""),
                })
            except Exception:
                pass
        elif sport == 80 or data[:5] == b"HTTP/":
            try:
                resp = dpkt.http.Response(data)
                self.http_responses.append({
                    "timestamp": timestamp,
                    "src_ip": src_ip,
                    "dst_ip": dst_ip,
                    "status_code": int(resp.status) if resp.status.isdigit() else 0,
                    "content_type": resp.headers.get("content-type", ""),
                    "server": resp.headers.get("server", ""),
                })
            except Exception:
                pass

        # TLS
        if len(data) > 5 and data[0] == 0x16:
            self._process_tls(data, src_ip, dst_ip, dport, timestamp)

    def _process_dpkt_udp(self, data, src_ip, dst_ip, sport, dport, timestamp):
        """Process UDP payload with dpkt."""
        if sport == 53 or dport == 53:
            try:
                dns = dpkt.dns.DNS(data)
                if dns.qr == 0:  # Query
                    for q in dns.qd:
                        qtype_map = {1: "A", 2: "NS", 5: "CNAME", 15: "MX",
                                     16: "TXT", 28: "AAAA", 33: "SRV"}
                        self.dns_queries.append({
                            "timestamp": timestamp,
                            "src_ip": src_ip,
                            "dst_ip": dst_ip,
                            "query": q.name,
                            "type": qtype_map.get(q.type, str(q.type)),
                        })
                else:  # Response
                    qname = dns.qd[0].name if dns.qd else ""
                    answers = []
                    for rr in dns.an:
                        if rr.type == 1:  # A record
                            import socket
                            try:
                                answers.append(socket.inet_ntoa(rr.rdata))
                            except Exception:
                                answers.append(str(rr.rdata))
                        elif rr.type == 28:  # AAAA
                            import socket
                            try:
                                answers.append(socket.inet_ntop(socket.AF_INET6, rr.rdata))
                            except Exception:
                                answers.append(str(rr.rdata))
                        else:
                            answers.append(str(rr.rdata))

                    rcode_map = {0: "NOERROR", 1: "FORMERR", 2: "SERVFAIL",
                                 3: "NXDOMAIN", 5: "REFUSED"}
                    self.dns_responses.append({
                        "timestamp": timestamp,
                        "src_ip": src_ip,
                        "query": qname,
                        "answers": answers,
                        "rcode": rcode_map.get(dns.rcode, str(dns.rcode)),
                    })
            except Exception:
                pass

    def detect_beacons(self, min_connections: int = 10, max_jitter: float = 0.3) -> list[dict]:
        """Detect beaconing behavior from connection timing data."""
        beacons = []

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

            sorted_times = sorted(times)
            deltas = [sorted_times[i + 1] - sorted_times[i]
                      for i in range(len(sorted_times) - 1)]

            if not deltas:
                continue

            mean_delta = sum(deltas) / len(deltas)
            if mean_delta <= 0:
                continue

            variance = sum((d - mean_delta) ** 2 for d in deltas) / len(deltas)
            std_dev = variance ** 0.5
            jitter = std_dev / mean_delta if mean_delta > 0 else float("inf")

            if jitter <= max_jitter:
                beacons.append({
                    "connection": key,
                    "count": len(times),
                    "mean_interval_seconds": round(mean_delta, 2),
                    "std_dev_seconds": round(std_dev, 2),
                    "jitter_ratio": round(jitter, 4),
                    "first_seen": sorted_times[0],
                    "last_seen": sorted_times[-1],
                    "confidence": "high" if jitter < 0.1 else "medium",
                })

        return sorted(beacons, key=lambda x: x["jitter_ratio"])

    def get_connection_summary(self) -> list[dict]:
        """Summarize all tracked connections."""
        summary = []
        for key, times in self.connection_times.items():
            sorted_times = sorted(times)
            summary.append({
                "connection": key,
                "count": len(times),
                "first_seen": sorted_times[0],
                "last_seen": sorted_times[-1],
                "duration_seconds": round(sorted_times[-1] - sorted_times[0], 2),
            })
        return sorted(summary, key=lambda x: x["count"], reverse=True)

    def generate_report(self) -> dict[str, Any]:
        """Generate the full analysis report."""
        beacons = self.detect_beacons()
        conn_summary = self.get_connection_summary()

        # Unique DNS domains
        unique_domains = list(set(q["query"] for q in self.dns_queries))

        # NXDOMAIN responses (potential DGA)
        nxdomains = [r for r in self.dns_responses if r.get("rcode") == "NXDOMAIN"]

        # Unique external IPs
        all_dst_ips = set()
        for req in self.http_requests:
            all_dst_ips.add(req["dst_ip"])
        for hs in self.tls_handshakes:
            all_dst_ips.add(hs["dst_ip"])
        for conn in conn_summary:
            parts = conn["connection"].split("->")
            if len(parts) == 2:
                all_dst_ips.add(parts[1].split(":")[0])

        report = {
            "metadata": {
                "tool": "pcap_analyzer.py",
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "pcap_file": os.path.abspath(self.pcap_path),
                "pcap_size_bytes": os.path.getsize(self.pcap_path),
                "platform": platform.system(),
            },
            "overview": {
                "total_packets": self.packet_count,
                "capture_start": self.start_time,
                "capture_end": self.end_time,
                "duration_seconds": round(
                    (self.end_time - self.start_time), 2
                ) if self.start_time and self.end_time else 0,
                "unique_dns_domains": len(unique_domains),
                "total_dns_queries": len(self.dns_queries),
                "total_http_requests": len(self.http_requests),
                "total_tls_handshakes": len(self.tls_handshakes),
                "unique_destination_ips": len(all_dst_ips),
                "nxdomain_count": len(nxdomains),
                "detected_beacons": len(beacons),
            },
            "dns": {
                "queries": self.dns_queries[:500],  # Limit output size
                "responses": self.dns_responses[:500],
                "unique_domains": sorted(unique_domains)[:200],
                "nxdomain_queries": [r["query"] for r in nxdomains][:100],
            },
            "http": {
                "requests": self.http_requests[:500],
                "responses": self.http_responses[:500],
                "unique_user_agents": list(set(
                    r.get("user_agent", "") for r in self.http_requests if r.get("user_agent")
                )),
            },
            "tls": {
                "handshakes": self.tls_handshakes[:500],
                "unique_sni": list(set(
                    h["sni"] for h in self.tls_handshakes if h.get("sni")
                )),
            },
            "beacons": beacons,
            "connections": conn_summary[:200],
            "destination_ips": sorted(all_dst_ips),
        }

        return report

    def analyze(self) -> None:
        """Run the analysis using the best available library."""
        if HAS_DPKT:
            try:
                self.analyze_with_dpkt()
                return
            except Exception as e:
                logger.warning(f"dpkt failed, falling back to scapy: {e}")

        if HAS_SCAPY:
            self.analyze_with_scapy()
            return

        logger.error(
            "No packet parsing library available. "
            "Install scapy (pip install scapy) or dpkt (pip install dpkt)."
        )
        sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="PCAP network traffic analyzer for malware analysis",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --pcap capture.pcap --output report.json
  %(prog)s --pcap capture.pcap --dns-only
  %(prog)s --pcap capture.pcap --extract-files ./carved/
        """,
    )
    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(
        "--dns-only",
        action="store_true",
        help="Only analyze DNS traffic",
    )
    parser.add_argument(
        "--extract-files",
        help="Directory to extract carved files into",
    )
    parser.add_argument(
        "--min-beacon-connections",
        type=int, default=10,
        help="Minimum connections to consider for beacon detection (default: 10)",
    )
    parser.add_argument(
        "--max-jitter",
        type=float, default=0.3,
        help="Maximum jitter ratio for beacon detection (default: 0.3)",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Enable 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)

    if args.extract_files:
        os.makedirs(args.extract_files, exist_ok=True)

    analyzer = PcapAnalyzer(args.pcap)
    analyzer.analyze()

    report = analyzer.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 to stderr
    overview = report["overview"]
    print(f"\n=== PCAP Analysis Summary ===", file=sys.stderr)
    print(f"Packets: {overview['total_packets']}", file=sys.stderr)
    print(f"Duration: {overview['duration_seconds']}s", file=sys.stderr)
    print(f"DNS queries: {overview['total_dns_queries']} ({overview['unique_dns_domains']} unique domains)", file=sys.stderr)
    print(f"HTTP requests: {overview['total_http_requests']}", file=sys.stderr)
    print(f"TLS handshakes: {overview['total_tls_handshakes']}", file=sys.stderr)
    print(f"NXDOMAIN responses: {overview['nxdomain_count']}", file=sys.stderr)
    print(f"Detected beacons: {overview['detected_beacons']}", file=sys.stderr)
    if report["beacons"]:
        print("\nBeacon candidates:", file=sys.stderr)
        for b in report["beacons"][:5]:
            print(f"  {b['connection']} - interval={b['mean_interval_seconds']}s "
                  f"jitter={b['jitter_ratio']} ({b['confidence']})", file=sys.stderr)


if __name__ == "__main__":
    main()