# pcapx/core/auth.py

from collections import defaultdict, Counter
from base64 import b64decode
import re

from scapy.all import TCP, Raw, IP, IPv6


def _get_ips(pkt):
    if pkt.haslayer(IP):
        return pkt[IP].src, pkt[IP].dst
    if pkt.haslayer(IPv6):
        return pkt[IPv6].src, pkt[IPv6].dst
    return None, None


def extract_auth_evidence(packets):
    credentials = []
    auth_attempts = Counter()

    # TCP stream buffers
    streams = defaultdict(bytes)

    # FTP session state
    ftp_sessions = {}

    for pkt in packets:
        if not pkt.haslayer(TCP) or not pkt.haslayer(Raw):
            continue

        src, dst = _get_ips(pkt)
        if not src or not dst:
            continue

        tcp = pkt[TCP]
        payload = pkt[Raw].load

        stream_key = (src, tcp.sport, dst, tcp.dport)
        streams[stream_key] += payload

        try:
            text = streams[stream_key].decode(errors="ignore")
        except Exception:
            continue

        text_upper = text.upper()

        # =====================================================
        # FTP AUTH (CORRECT, STREAM-BASED)
        # =====================================================
        if tcp.dport == 21 or tcp.sport == 21:
            for line in text.split("\r\n"):
                if line.startswith("USER "):
                    ftp_sessions[stream_key] = {
                        "protocol": "FTP",
                        "src": src,
                        "dst": dst,
                        "username": line[5:].strip(),
                        "password": None,
                        "evidence": [line.strip()],
                    }
                    auth_attempts[(src, dst, "FTP")] += 1

                elif line.startswith("PASS ") and stream_key in ftp_sessions:
                    ftp_sessions[stream_key]["password"] = line[5:].strip()
                    ftp_sessions[stream_key]["evidence"].append(line.strip())

                    credentials.append(ftp_sessions[stream_key])
                    del ftp_sessions[stream_key]

        # =====================================================
        # HTTP BASIC AUTH
        # =====================================================
        if b"Authorization: Basic" in payload:
            try:
                b64 = payload.split(b"Basic", 1)[1].strip()
                userpass = b64decode(b64).decode(errors="ignore")
                if ":" in userpass:
                    u, p = userpass.split(":", 1)
                    credentials.append({
                        "protocol": "HTTP",
                        "type": "Basic",
                        "username": u,
                        "password": p,
                        "src": src,
                        "dst": dst,
                        "evidence": [userpass],
                    })
            except Exception:
                pass

        # =====================================================
        # HTTP FORM / JSON (BEST-EFFORT, FACT ONLY)
        # =====================================================
        lower = text.lower()
        if "password=" in lower and ("username=" in lower or "user=" in lower):
            credentials.append({
                "protocol": "HTTP",
                "type": "Form",
                "src": src,
                "dst": dst,
                "evidence": [line.strip() for line in text.splitlines() if "password=" in line.lower()],
            })

    # Bruteforce = many auth attempts, no guessing
    bruteforce = [
        {"protocol": proto, "src": src, "dst": dst, "attempts": cnt}
        for (src, dst, proto), cnt in auth_attempts.items()
        if cnt >= 10
    ]

    return {
        "credentials": credentials,
        "auth_attempts": [
            {"protocol": proto, "src": src, "dst": dst, "count": cnt}
            for (src, dst, proto), cnt in auth_attempts.items()
        ],
        "bruteforce_indicators": bruteforce,
    }
