#!/usr/bin/env python3
"""Detect Domain Generation Algorithm (DGA) domains using statistical analysis.

Analyzes domain names for characteristics typical of algorithmically generated
domains: high entropy, unusual character distributions, non-linguistic patterns.
"""

from __future__ import annotations

import argparse
import json
import math
import re
import sys
from collections import Counter
from datetime import datetime
from pathlib import Path


# English language bigram frequencies (approximate, normalized)
ENGLISH_BIGRAMS = {
    "th": 0.0356, "he": 0.0307, "in": 0.0243, "er": 0.0205, "an": 0.0199,
    "re": 0.0185, "on": 0.0176, "at": 0.0149, "en": 0.0145, "nd": 0.0135,
    "ti": 0.0134, "es": 0.0134, "or": 0.0128, "te": 0.0120, "of": 0.0117,
    "ed": 0.0117, "is": 0.0113, "it": 0.0112, "al": 0.0109, "ar": 0.0107,
    "st": 0.0105, "to": 0.0104, "nt": 0.0104, "ng": 0.0095, "se": 0.0093,
    "ha": 0.0093, "as": 0.0087, "ou": 0.0087, "io": 0.0083, "le": 0.0083,
    "ve": 0.0083, "co": 0.0079, "me": 0.0079, "de": 0.0076, "hi": 0.0076,
    "ri": 0.0073, "ro": 0.0073, "ic": 0.0070, "ne": 0.0069, "ea": 0.0069,
}

VOWELS = set("aeiou")
CONSONANTS = set("bcdfghjklmnpqrstvwxyz")


def calculate_entropy(domain) -> dict:
    """Calculate Shannon entropy of a domain string."""
    if not domain:
        return 0.0
    counter = Counter(domain.lower())
    length = len(domain)
    entropy = 0.0
    for count in counter.values():
        prob = count / length
        if prob > 0:
            entropy -= prob * math.log2(prob)
    return entropy


def calculate_bigram_score(domain) -> dict:
    """Score domain based on English bigram frequency."""
    domain = domain.lower()
    if len(domain) < 2:
        return 0.0

    score = 0.0
    count = 0
    for i in range(len(domain) - 1):
        bigram = domain[i:i + 2]
        if bigram in ENGLISH_BIGRAMS:
            score += ENGLISH_BIGRAMS[bigram]
        count += 1

    return score / count if count > 0 else 0.0


def vowel_consonant_ratio(domain) -> dict:
    """Calculate vowel to consonant ratio."""
    domain = domain.lower()
    vowel_count = sum(1 for c in domain if c in VOWELS)
    consonant_count = sum(1 for c in domain if c in CONSONANTS)

    if consonant_count == 0:
        return float("inf") if vowel_count > 0 else 0.0
    return vowel_count / consonant_count


def consecutive_consonants_max(domain) -> dict:
    """Find maximum consecutive consonant run."""
    domain = domain.lower()
    max_run = 0
    current_run = 0
    for c in domain:
        if c in CONSONANTS:
            current_run += 1
            max_run = max(max_run, current_run)
        else:
            current_run = 0
    return max_run


def digit_ratio(domain) -> dict:
    """Calculate ratio of digits to total characters."""
    if not domain:
        return 0.0
    return sum(1 for c in domain if c.isdigit()) / len(domain)


def unique_char_ratio(domain) -> dict:
    """Calculate ratio of unique characters."""
    if not domain:
        return 0.0
    return len(set(domain.lower())) / len(domain)


def analyze_domain(domain) -> dict:
    """Analyze a single domain for DGA characteristics."""
    # Extract second-level domain (remove TLD)
    parts = domain.lower().strip().rstrip(".").split(".")
    if len(parts) >= 2:
        sld = parts[-2]  # second-level domain
        tld = parts[-1]
    else:
        sld = parts[0]
        tld = ""

    # Calculate features
    entropy = calculate_entropy(sld)
    bigram_score = calculate_bigram_score(sld)
    vc_ratio = vowel_consonant_ratio(sld)
    max_consonants = consecutive_consonants_max(sld)
    digits = digit_ratio(sld)
    unique_ratio = unique_char_ratio(sld)
    length = len(sld)

    # DGA scoring heuristics
    dga_score = 0.0

    # High entropy suggests randomness
    if entropy > 3.5:
        dga_score += 0.25
    if entropy > 4.0:
        dga_score += 0.15

    # Low bigram score suggests non-English
    if bigram_score < 0.002:
        dga_score += 0.2
    if bigram_score < 0.001:
        dga_score += 0.1

    # Abnormal vowel/consonant ratio
    if vc_ratio < 0.2 or vc_ratio > 1.5:
        dga_score += 0.15

    # Long consonant runs
    if max_consonants >= 4:
        dga_score += 0.15
    if max_consonants >= 6:
        dga_score += 0.1

    # Contains digits mixed with letters
    if 0 < digits < 0.5:
        dga_score += 0.1

    # Unusual length
    if length > 20:
        dga_score += 0.1
    if length > 30:
        dga_score += 0.1

    # High unique character ratio (more random)
    if unique_ratio > 0.8 and length > 8:
        dga_score += 0.1

    # Normalize to [0, 1]
    dga_score = min(dga_score, 1.0)

    # Classification
    if dga_score >= 0.6:
        classification = "likely_dga"
    elif dga_score >= 0.4:
        classification = "suspicious"
    else:
        classification = "likely_legitimate"

    return {
        "domain": domain,
        "sld": sld,
        "tld": tld,
        "classification": classification,
        "dga_score": round(dga_score, 3),
        "features": {
            "entropy": round(entropy, 3),
            "bigram_score": round(bigram_score, 5),
            "vowel_consonant_ratio": round(vc_ratio, 3) if vc_ratio != float("inf") else "inf",
            "max_consecutive_consonants": max_consonants,
            "digit_ratio": round(digits, 3),
            "unique_char_ratio": round(unique_ratio, 3),
            "length": length,
        }
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Detect DGA (Domain Generation Algorithm) domains"
    )
    parser.add_argument(
        "--input", "-i", required=True,
        help="Input file with one domain per line, or single domain"
    )
    parser.add_argument(
        "--output", "-o",
        help="Output file path (JSON)"
    )
    parser.add_argument(
        "--threshold", "-t", type=float, default=0.4,
        help="DGA score threshold for flagging (default: 0.4)"
    )
    parser.add_argument(
        "--format", choices=["json", "text", "csv"], default="text",
        help="Output format"
    )
    parser.add_argument(
        "--only-dga", action="store_true",
        help="Only output domains classified as DGA/suspicious"
    )
    args = parser.parse_args()

    # Read domains
    input_path = Path(args.input)
    if input_path.exists():
        domains = [
            line.strip() for line in input_path.read_text().split("\n")
            if line.strip() and not line.startswith("#")
        ]
    else:
        # Treat as single domain
        domains = [args.input]

    # Analyze
    results = [analyze_domain(d) for d in domains]

    if args.only_dga:
        results = [r for r in results if r["dga_score"] >= args.threshold]

    report = {
        "tool": "dga_detector",
        "timestamp": datetime.now().isoformat(),
        "total_domains": len(domains),
        "dga_detected": sum(1 for r in results if r["classification"] == "likely_dga"),
        "suspicious": sum(1 for r in results if r["classification"] == "suspicious"),
        "legitimate": sum(1 for r in results if r["classification"] == "likely_legitimate"),
        "threshold": args.threshold,
        "results": results,
    }

    if args.format == "json":
        output = json.dumps(report, indent=2)
    elif args.format == "csv":
        output = "domain,classification,dga_score,entropy,length\n"
        for r in results:
            output += f"{r['domain']},{r['classification']},{r['dga_score']},{r['features']['entropy']},{r['features']['length']}\n"
    else:
        output = f"=== DGA Detection Report ===\n"
        output += f"Total: {report['total_domains']} | DGA: {report['dga_detected']} | Suspicious: {report['suspicious']} | Legitimate: {report['legitimate']}\n\n"
        for r in results:
            marker = "!!!" if r["classification"] == "likely_dga" else (" ? " if r["classification"] == "suspicious" else "   ")
            output += f"[{marker}] {r['domain']} — score: {r['dga_score']} ({r['classification']})\n"

    if args.output:
        Path(args.output).write_text(output)
        print(f"[+] Report saved to {args.output}")
    else:
        print(output)


if __name__ == "__main__":
    main()