#!/usr/bin/env python3
"""Agent for detecting DNS tunneling using Zeek log analysis."""

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


ENTROPY_THRESHOLD = 3.5
MIN_QUERIES_PER_DOMAIN = 20
MAX_NORMAL_SUBDOMAIN_LEN = 30
TUNNEL_QUERY_TYPES = {"TXT", "NULL", "CNAME", "MX"}


def shannon_entropy(data):
    """Calculate Shannon entropy of a string."""
    if not data:
        return 0.0
    freq = defaultdict(int)
    for c in data:
        freq[c] += 1
    n = len(data)
    return -sum((cnt/n) * math.log2(cnt/n) for cnt in freq.values())


def load_dns_log(filepath):
    """Load Zeek dns.log (TSV format)."""
    entries = []
    try:
        with open(filepath, "r") as f:
            for line in f:
                if line.startswith("#"):
                    continue
                parts = line.strip().split("\t")
                if len(parts) >= 10:
                    entries.append({
                        "ts": parts[0],
                        "uid": parts[1],
                        "src": parts[2],
                        "src_port": parts[3],
                        "dst": parts[4],
                        "dst_port": parts[5],
                        "query": parts[9] if len(parts) > 9 else "",
                        "qtype": parts[13] if len(parts) > 13 else "",
                        "answers": parts[21] if len(parts) > 21 else "",
                    })
    except (OSError, IndexError) as e:
        print(f"[!] Error loading DNS log: {e}")
    return entries


def analyze_domain_statistics(entries):
    """Compute per-domain statistics for tunneling detection."""
    domain_data = defaultdict(lambda: {
        "queries": [], "subdomains": [], "qtypes": defaultdict(int),
        "sources": set(), "total_subdomain_len": 0,
    })

    for entry in entries:
        query = entry.get("query", "")
        if not query or query == "-":
            continue
        parts = query.rstrip(".").split(".")
        if len(parts) < 2:
            continue
        domain = ".".join(parts[-2:])
        subdomain = ".".join(parts[:-2])
        d = domain_data[domain]
        d["queries"].append(query)
        d["subdomains"].append(subdomain)
        d["qtypes"][entry.get("qtype", "")] += 1
        d["sources"].add(entry.get("src", ""))
        d["total_subdomain_len"] += len(subdomain)

    return domain_data


def detect_tunneling(domain_data):
    """Apply tunneling detection heuristics."""
    findings = []
    for domain, data in domain_data.items():
        query_count = len(data["queries"])
        if query_count < MIN_QUERIES_PER_DOMAIN:
            continue

        avg_subdomain_len = data["total_subdomain_len"] / query_count
        all_subdomain_text = "".join(data["subdomains"])
        entropy = shannon_entropy(all_subdomain_text)

        tunnel_qtype_count = sum(
            data["qtypes"].get(qt, 0) for qt in TUNNEL_QUERY_TYPES
        )
        tunnel_qtype_ratio = tunnel_qtype_count / query_count if query_count else 0

        score = 0
        if entropy > ENTROPY_THRESHOLD:
            score += 40
        if avg_subdomain_len > MAX_NORMAL_SUBDOMAIN_LEN:
            score += 30
        if tunnel_qtype_ratio > 0.5:
            score += 20
        if query_count > 500:
            score += 10

        if score >= 40:
            findings.append({
                "domain": domain,
                "query_count": query_count,
                "avg_subdomain_length": round(avg_subdomain_len, 1),
                "entropy": round(entropy, 3),
                "tunnel_qtype_ratio": round(tunnel_qtype_ratio, 3),
                "unique_sources": len(data["sources"]),
                "tunnel_score": score,
                "severity": "CRITICAL" if score >= 70 else "HIGH" if score >= 50 else "MEDIUM",
            })

    findings.sort(key=lambda f: f["tunnel_score"], reverse=True)
    return findings


def main():
    parser = argparse.ArgumentParser(
        description="DNS tunneling detection agent using Zeek logs"
    )
    parser.add_argument("dns_log", help="Path to Zeek dns.log")
    parser.add_argument("--min-queries", type=int, default=20)
    parser.add_argument("--entropy-threshold", type=float, default=3.5)
    parser.add_argument("--output", "-o", help="Output JSON report")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    global MIN_QUERIES_PER_DOMAIN, ENTROPY_THRESHOLD
    MIN_QUERIES_PER_DOMAIN = args.min_queries
    ENTROPY_THRESHOLD = args.entropy_threshold

    print("[*] DNS Tunneling Detection Agent (Zeek)")
    entries = load_dns_log(args.dns_log)
    if not entries:
        print("[!] No DNS entries loaded")
        sys.exit(1)

    print(f"[*] Loaded {len(entries)} DNS queries")
    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "source_file": args.dns_log,
        "total_queries": len(entries),
        "findings": [],
    }

    domain_data = analyze_domain_statistics(entries)
    findings = detect_tunneling(domain_data)
    report["findings"] = findings
    report["risk_level"] = (
        "CRITICAL" if any(f["severity"] == "CRITICAL" for f in findings)
        else "HIGH" if findings else "LOW"
    )

    print(f"[*] Detected {len(findings)} suspected DNS tunnels")
    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()

#!/usr/bin/env python3
"""
DNS Tunneling Detection Script for Zeek Logs
Analyzes dns.log for tunneling indicators: query length, entropy,
unique subdomains, record types, and query volume.
"""

import json
import csv
import argparse
import datetime
import math
from collections import defaultdict, Counter
from pathlib import Path

ENTROPY_THRESHOLD = 3.5
QUERY_LENGTH_THRESHOLD = 50
SUBDOMAIN_LENGTH_THRESHOLD = 25
UNIQUE_SUBDOMAIN_THRESHOLD = 50
QUERY_COUNT_THRESHOLD = 100

SUSPICIOUS_RECORD_TYPES = {"TXT", "NULL", "CNAME", "MX", "KEY", "SRV"}


def calculate_entropy(text: str) -> float:
    if not text:
        return 0.0
    counts = Counter(text)
    length = len(text)
    return -sum((c / length) * math.log2(c / length) for c in counts.values())


def parse_zeek_dns(input_path: str) -> list[dict]:
    path = Path(input_path)
    events = []
    if path.suffix == ".json":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            events = data if isinstance(data, list) else data.get("events", [])
    elif path.suffix == ".csv":
        with open(path, "r", encoding="utf-8-sig") as f:
            events = [dict(row) for row in csv.DictReader(f)]
    elif path.suffix == ".log":
        with open(path, "r", encoding="utf-8") as f:
            headers = None
            for line in f:
                if line.startswith("#fields"):
                    headers = line.strip().split("\t")[1:]
                elif line.startswith("#"):
                    continue
                elif headers:
                    values = line.strip().split("\t")
                    if len(values) == len(headers):
                        events.append(dict(zip(headers, values)))
    return events


def detect_dns_tunneling(queries: list[dict]) -> list[dict]:
    domain_stats = defaultdict(lambda: {
        "queries": 0, "unique_subdomains": set(), "total_length": 0,
        "max_length": 0, "record_types": Counter(), "sources": set(),
        "subdomain_entropy_samples": [],
    })

    for q in queries:
        query = q.get("query", "")
        src = q.get("id.orig_h", q.get("src_ip", q.get("source_ip", "")))
        qtype = q.get("qtype_name", q.get("query_type", "A"))
        if not query or len(query.split(".")) < 3:
            continue

        parts = query.rstrip(".").split(".")
        base_domain = ".".join(parts[-2:])
        subdomain = ".".join(parts[:-2])

        stats = domain_stats[(src, base_domain)]
        stats["queries"] += 1
        stats["unique_subdomains"].add(subdomain)
        stats["total_length"] += len(query)
        stats["max_length"] = max(stats["max_length"], len(query))
        stats["record_types"][qtype] += 1
        stats["sources"].add(src)
        if len(stats["subdomain_entropy_samples"]) < 100:
            stats["subdomain_entropy_samples"].append(subdomain)

    findings = []
    for (src, base_domain), stats in domain_stats.items():
        if stats["queries"] < 20:
            continue

        avg_len = stats["total_length"] / stats["queries"]
        unique_subs = len(stats["unique_subdomains"])

        avg_entropy = 0.0
        if stats["subdomain_entropy_samples"]:
            entropies = [calculate_entropy(s) for s in stats["subdomain_entropy_samples"]]
            avg_entropy = sum(entropies) / len(entropies)

        risk = 0
        indicators = []

        if avg_len > QUERY_LENGTH_THRESHOLD:
            risk += 25
            indicators.append(f"High avg query length: {avg_len:.1f}")
        if unique_subs > UNIQUE_SUBDOMAIN_THRESHOLD:
            risk += 25
            indicators.append(f"High unique subdomains: {unique_subs}")
        if avg_entropy > ENTROPY_THRESHOLD:
            risk += 25
            indicators.append(f"High subdomain entropy: {avg_entropy:.2f}")
        if stats["queries"] > QUERY_COUNT_THRESHOLD:
            risk += 10
            indicators.append(f"High query volume: {stats['queries']}")

        suspicious_types = sum(stats["record_types"].get(rt, 0) for rt in SUSPICIOUS_RECORD_TYPES)
        if suspicious_types > stats["queries"] * 0.5:
            risk += 15
            indicators.append(f"Unusual record types: {dict(stats['record_types'])}")

        if risk < 30:
            continue

        risk_level = "CRITICAL" if risk >= 70 else "HIGH" if risk >= 50 else "MEDIUM"
        findings.append({
            "source_ip": src,
            "base_domain": base_domain,
            "query_count": stats["queries"],
            "unique_subdomains": unique_subs,
            "avg_query_length": round(avg_len, 1),
            "max_query_length": stats["max_length"],
            "avg_subdomain_entropy": round(avg_entropy, 3),
            "record_types": dict(stats["record_types"]),
            "risk_score": risk,
            "risk_level": risk_level,
            "indicators": indicators,
        })

    return sorted(findings, key=lambda x: x["risk_score"], reverse=True)


def run_hunt(input_path: str, output_dir: str) -> None:
    print(f"[*] DNS Tunneling Hunt - {datetime.datetime.now().isoformat()}")
    queries = parse_zeek_dns(input_path)
    print(f"[*] Loaded {len(queries)} DNS queries")
    findings = detect_dns_tunneling(queries)
    print(f"[!] DNS tunneling detections: {len(findings)}")

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    with open(output_path / "dns_tunnel_findings.json", "w", encoding="utf-8") as f:
        json.dump({"hunt_id": f"TH-DNSTUNNEL-{datetime.date.today().isoformat()}",
                    "total_queries": len(queries), "findings": findings}, f, indent=2)
    print(f"[+] Results written to {output_dir}")


def main():
    parser = argparse.ArgumentParser(description="DNS Tunneling Detection")
    parser.add_argument("--input", "-i", required=True)
    parser.add_argument("--output", "-o", default="./dns_tunnel_hunt_output")
    args = parser.parse_args()
    run_hunt(args.input, args.output)


if __name__ == "__main__":
    main()