#!/usr/bin/env python3
"""Shared core for the local-dev intrusion-detection scripts.

`copyparty.py` and `npm.py` are thin adapters. They only provide:

  * a log-file matcher (glob-ish pattern / suffixes)
  * a line parser producing `Request` objects
  * a flagging rule (which IPs look hostile)
  * column definitions for their CSV output

Everything else -- IP aggregation, whitelisting, console + file reports,
MikroTik RouterOS export -- lives in this module.
"""

from __future__ import annotations
import bad_patterns
from whitelists import (
    WHITELISTS,
    WHITELIST_NETS as _WHITELIST_NETS,
    refresh_whitelists,
)

import argparse
import csv
import gzip
import ipaddress
import json
import os
import re
from collections import Counter
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Iterable, Iterator, Sequence

# ---------------------------------------------------------------------------
# constants
# ---------------------------------------------------------------------------

ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")

# Whitelists (imported from dedicated module)

# Malicious pattern detection

#: (header, getter) pair used for CSV + console tables
Column = tuple[str, Callable[[Any], Any]]


# ---------------------------------------------------------------------------
# List download helpers
# ---------------------------------------------------------------------------

# download helpers and whitelist refresh now live in whitelists.py

# ---------------------------------------------------------------------------
# IP helpers
# ---------------------------------------------------------------------------

def parse_ip(value: str):
    """Return an ip_address object, or None if `value` is not an IP."""
    try:
        return ipaddress.ip_address(value)
    except (ValueError, TypeError):
        return None


def ip_version(value: str) -> int | None:
    addr = parse_ip(value)
    return addr.version if addr else None


def is_ipv6(value: str) -> bool:
    return ip_version(value) == 6


def is_ipv4(value: str) -> bool:
    return ip_version(value) == 4


def whitelist_of(value: str) -> str | None:
    """Return the whitelist name an IP belongs to, else None."""
    addr = parse_ip(value)
    if addr is None:
        return None
    for name, nets in _WHITELIST_NETS.items():
        for net in nets:
            if addr.version == net.version and addr in net:
                return name
    return None


def is_whitelisted(value: str) -> bool:
    return whitelist_of(value) is not None


# WHOIS Cache file path
WHOIS_CACHE_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "data", "whois_cache.json"
)

_whois_cache: dict[str, str] = {}


def load_whois_cache() -> None:
    global _whois_cache
    if not _whois_cache and os.path.exists(WHOIS_CACHE_PATH):
        try:
            with open(WHOIS_CACHE_PATH, "r", encoding="utf-8") as f:
                _whois_cache = json.load(f)
        except Exception:
            _whois_cache = {}


def save_whois_cache() -> None:
    if _whois_cache:
        os.makedirs(os.path.dirname(WHOIS_CACHE_PATH), exist_ok=True)
        try:
            with open(WHOIS_CACHE_PATH, "w", encoding="utf-8") as f:
                json.dump(_whois_cache, f, indent=2)
        except Exception:
            pass


def query_whois_subnet(ip: str) -> str | None:
    """Query whois.radb.net for the route/subnet of an IP address."""
    addr = parse_ip(ip)
    if addr is None:
        return None

    load_whois_cache()
    if ip in _whois_cache:
        return _whois_cache[ip] or None

    import socket
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(3.0)
        s.connect(("whois.radb.net", 43))
        s.sendall(f"{ip}\r\n".encode())
        res = ""
        while True:
            chunk = s.recv(4096).decode(errors="ignore")
            if not chunk:
                break
            res += chunk
        s.close()

        # Find route/route6 patterns
        routes = re.findall(r"^(?:route|route6):\s*(\S+)",
                            res, re.IGNORECASE | re.MULTILINE)
        if routes:
            # Sort by subnet length descending to find the most specific network (longest prefix)
            routes_parsed = []
            for r in routes:
                try:
                    routes_parsed.append(ipaddress.ip_network(r))
                except ValueError:
                    continue
            if routes_parsed:
                # Get the most specific one (highest prefixlen)
                best_route = max(routes_parsed, key=lambda n: n.prefixlen)
                subnet_str = str(best_route)
                _whois_cache[ip] = subnet_str
                save_whois_cache()
                return subnet_str

    except Exception:
        pass

    # Fallback to single host subnet
    subnet_str = f"{ip}/32" if addr.version == 4 else f"{ip}/128"
    _whois_cache[ip] = subnet_str
    save_whois_cache()
    return subnet_str


def consolidate_by_subnet(ips: Iterable[str], threshold: int = 3) -> set[str]:
    """Group IPs by WHOIS subnet; if >= threshold IPs in a subnet, block the subnet.

    Otherwise, block individual IPs.
    """
    ip_list = list(ips)
    subnet_to_ips: dict[str, set[str]] = {}
    ip_to_subnet: dict[str, str] = {}

    for ip in ip_list:
        subnet = query_whois_subnet(ip)
        if subnet:
            # Verify the IP is actually in the subnet
            try:
                net = ipaddress.ip_network(subnet, strict=False)
                addr = ipaddress.ip_address(ip)
                if addr in net:
                    subnet_to_ips.setdefault(subnet, set()).add(ip)
                    ip_to_subnet[ip] = subnet
                    continue
            except ValueError:
                pass
        # Fallback to self
        fallback_net = f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128"
        subnet_to_ips.setdefault(fallback_net, set()).add(ip)
        ip_to_subnet[ip] = fallback_net

    blocked_targets: set[str] = set()
    for subnet, mapped_ips in subnet_to_ips.items():
        if len(mapped_ips) >= threshold:
            blocked_targets.add(subnet)
        else:
            for ip in mapped_ips:
                blocked_targets.add(
                    f"{ip}/32" if ip_version(ip) == 4 else f"{ip}/128")

    return blocked_targets


def sort_ips(ips: Iterable[str]) -> list[str]:
    """Numeric sort, IPv4 before IPv6, unparseable last."""
    def key(ip: str):
        addr = parse_ip(ip)
        if addr is None:
            return (9, b"", ip)
        return (addr.version, addr.packed, "")
    return sorted(ips, key=key)


def split_versions(ips: Iterable[str]) -> tuple[list[str], list[str]]:
    """Split an IP iterable into (sorted IPv4, sorted IPv6)."""
    ips = list(ips)
    return (sort_ips(ip for ip in ips if is_ipv4(ip)),
            sort_ips(ip for ip in ips if is_ipv6(ip)))


# ---------------------------------------------------------------------------
# data model
# ---------------------------------------------------------------------------

@dataclass(slots=True)
class Request:
    """One normalised request line, produced by any log parser."""

    ts: str = ""
    ip: str = ""
    method: str | None = None
    path: str | None = None
    status: int | None = None
    upstream_status: int | None = None
    host: str | None = None
    user: str | None = None
    ua: str | None = None
    referer: str | None = None
    denied: bool = False
    banned: bool = False
    extra: dict = field(default_factory=dict)

    @property
    def is_client_error(self) -> bool:
        return self.status is not None and 400 <= self.status < 500

    @property
    def is_server_error(self) -> bool:
        return self.status is not None and 500 <= self.status < 600

    @property
    def is_error(self) -> bool:
        return self.denied or self.is_client_error or self.is_server_error


@dataclass
class IPStats:
    """Aggregated per-IP counters."""

    ip: str
    version: int = 0
    total: int = 0
    errors: int = 0
    c4xx: int = 0
    c5xx: int = 0
    denied: int = 0
    banned: int = 0
    first: str = ""
    last: str = ""
    paths: Counter = field(default_factory=Counter)
    methods: Counter = field(default_factory=Counter)
    hosts: Counter = field(default_factory=Counter)
    uas: Counter = field(default_factory=Counter)
    users: set[str] = field(default_factory=set)
    malicious_hits: int = 0
    malicious_patterns: dict[str, int] = field(default_factory=dict)

    @property
    def error_ratio(self) -> float:
        return round(self.errors / self.total, 3) if self.total else 0.0

    @property
    def authenticated(self) -> bool:
        """True if the IP ever presented a real (non-anonymous) user."""
        return any(u and u != "*" for u in self.users)

    @property
    def whitelist(self) -> str:
        return whitelist_of(self.ip) or ""

    @property
    def users_str(self) -> str:
        return ",".join(sorted(self.users)) or "-"

    @property
    def methods_str(self) -> str:
        return ",".join(sorted(self.methods)) or "-"

    def to_dict(self) -> dict:
        return {
            "ip": self.ip,
            "version": self.version,
            "total": self.total,
            "errors": self.errors,
            "error_ratio": self.error_ratio,
            "status_4xx": self.c4xx,
            "status_5xx": self.c5xx,
            "denied": self.denied,
            "banned": self.banned,
            "first": self.first,
            "last": self.last,
            "whitelist": self.whitelist or None,
            "users": sorted(self.users),
            "methods": dict(self.methods),
            "hosts": dict(self.hosts),
            "paths": dict(self.paths.most_common()),
            "uas": dict(self.uas.most_common(10)),
            "malicious_hits": self.malicious_hits,
            "malicious_patterns": dict(self.malicious_patterns),
        }


def aggregate(requests: Iterable[Request]) -> dict[str, IPStats]:
    """Fold request records into per-IP statistics."""
    stats: dict[str, IPStats] = {}
    for r in requests:
        s = stats.get(r.ip)
        if s is None:
            s = stats[r.ip] = IPStats(
                ip=r.ip, version=ip_version(r.ip) or 0, first=r.ts)
        s.total += 1
        s.last = r.ts
        if r.denied:
            s.denied += 1
        if r.banned:
            s.banned += 1
        if r.is_client_error:
            s.c4xx += 1
        if r.is_server_error:
            s.c5xx += 1
        if r.is_error:
            s.errors += 1
        if r.path:
            s.paths[r.path] += 1
        if r.method:
            s.methods[r.method] += 1
        if r.host:
            s.hosts[r.host] += 1
        if r.ua and r.ua != "-":
            s.uas[r.ua] += 1
        if r.user:
            s.users.add(r.user)
        if r.path and bad_patterns.is_malicious(r.path):
            s.malicious_hits += 1
            for p in bad_patterns.classify_path(r.path):
                s.malicious_patterns[p.name] = s.malicious_patterns.get(
                    p.name, 0) + 1
    return stats


# --- flagging rules --------------------------------------------------------

def flag_scanners(stats: dict[str, IPStats], min_errors: int = 5,
                  min_ratio: float = 0.5) -> dict[str, IPStats]:
    """IPs with many errors AND a high error ratio."""
    return {
        ip: s for ip, s in stats.items()
        if s.errors >= min_errors and s.error_ratio >= min_ratio
    }


def flag_denied_only(stats: dict[str, IPStats],
                     min_denied: int = 1) -> dict[str, IPStats]:
    """IPs that only ever hit denied paths and never authenticated."""
    return {
        ip: s for ip, s in stats.items()
        if s.denied >= min_denied and not s.authenticated
    }


def flag_malicious_patterns(stats: dict[str, IPStats],
                            min_severity: int = 2,
                            min_hits: int = 1) -> dict[str, IPStats]:
    """IPs with hits against malicious path patterns at or above severity threshold."""
    return {
        ip: s for ip, s in stats.items()
        if s.malicious_hits >= min_hits and any(
            bad_patterns.PATTERNS_BY_NAME[name].severity >= min_severity
            for name in s.malicious_patterns
        )
    }


def by_total(stats: dict[str, IPStats]) -> list[IPStats]:
    """Stats sorted by request count, descending."""
    return sorted(stats.values(), key=lambda s: (-s.total, s.ip))


# ---------------------------------------------------------------------------
# filesystem helpers
# ---------------------------------------------------------------------------

def validate_dir(path: str, need_write: bool = False) -> str:
    if not os.path.exists(path):
        raise FileNotFoundError(f"path does not exist: {path}")
    if not os.path.isdir(path):
        raise NotADirectoryError(f"path is not a directory: {path}")
    if not os.access(path, os.R_OK):
        raise PermissionError(f"path not readable: {path}")
    if need_write and not os.access(path, os.W_OK):
        raise PermissionError(f"path not writable: {path}")
    return path


def ensure_dirs(*paths: str) -> None:
    for p in paths:
        os.makedirs(p, exist_ok=True)


def find_logs(path: str, pattern: re.Pattern | None = None,
              suffixes: Sequence[str] = (),
              exclude_prefixes: Sequence[str] = ()) -> list[str]:
    """List matching log file names (not full paths), sorted."""
    out = []
    for name in os.listdir(path):
        if not os.path.isfile(os.path.join(path, name)):
            continue
        if any(name.startswith(p) for p in exclude_prefixes):
            continue
        if pattern is not None and not pattern.match(name):
            continue
        if suffixes and not name.endswith(tuple(suffixes)):
            continue
        out.append(name)
    return sorted(out)


def read_lines(path: str) -> Iterator[str]:
    """Yield ANSI-stripped, newline-stripped lines. Transparently gunzips."""
    opener = gzip.open if path.endswith(".gz") else open
    with opener(path, "rt", encoding="utf-8", errors="replace") as f:
        for line in f:
            yield ANSI_RE.sub("", line.rstrip("\r\n"))


# ---------------------------------------------------------------------------
# output helpers
# ---------------------------------------------------------------------------

def write_csv(path: str, header: Sequence[str],
              rows: Iterable[Sequence[Any]]) -> str:
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(header)
        w.writerows(rows)
    return path


def write_json(path: str, obj: Any) -> str:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2, default=str)
    return path


def write_lines(path: str, lines: Iterable[str]) -> str:
    with open(path, "w", encoding="utf-8") as f:
        for line in lines:
            f.write(line + "\n")
    return path


def columns_header(columns: Sequence[Column]) -> list[str]:
    return [c[0] for c in columns]


def columns_row(columns: Sequence[Column], obj: Any) -> list[Any]:
    return [fn(obj) for _, fn in columns]


def render_table(headers: Sequence[str], rows: Sequence[Sequence[Any]],
                 max_width: int = 48) -> str:
    """Aligned fixed-width text table. Numeric columns right-aligned."""
    def cell(v: Any) -> str:
        s = "" if v is None else str(v)
        return s if len(s) <= max_width else s[:max_width - 1] + "~"

    body = [[cell(v) for v in row] for row in rows]
    aligns = []
    for i in range(len(headers)):
        numeric = bool(rows) and all(
            isinstance(row[i], (int, float)) and not isinstance(row[i], bool)
            for row in rows)
        aligns.append(">" if numeric else "<")
    widths = [len(h) for h in headers]
    for row in body:
        for i, v in enumerate(row):
            widths[i] = max(widths[i], len(v))

    def line(values: Sequence[str]) -> str:
        return "  ".join(
            f"{v:{a}{w}}" for v, a, w in zip(values, aligns, widths)).rstrip()

    out = [line(list(headers)), "-" * (sum(widths) + 2 * (len(widths) - 1))]
    out += [line(row) for row in body]
    return "\n".join(out)


def render_kv(rows: Sequence[tuple[str, Any]]) -> str:
    width = max((len(k) for k, _ in rows), default=0) + 1
    return "\n".join(f"{k + ':':<{width + 1}} {v}" for k, v in rows)


def timestamp_name() -> str:
    return datetime.now().strftime("%Y-%m-%d-%H-%M-%S")


# ---------------------------------------------------------------------------
# MikroTik RouterOS export
# ---------------------------------------------------------------------------

def split_cidr_versions(targets: Iterable[str]) -> tuple[list[str], list[str]]:
    """Split target CIDRs/IPs into (sorted IPv4 list, sorted IPv6 list)."""
    v4 = []
    v6 = []
    for t in targets:
        # t might have CIDR suffix
        ip = t.split("/")[0]
        if is_ipv4(ip):
            v4.append(t)
        elif is_ipv6(ip):
            v6.append(t)

    # Sort CIDR/IP strings
    def key_cidr(cidr: str):
        ip = cidr.split("/")[0]
        addr = parse_ip(ip)
        if addr is None:
            return (9, b"", cidr)
        return (addr.version, addr.packed, cidr)

    return sorted(v4, key=key_cidr), sorted(v6, key=key_cidr)


def mikrotik_script(tag: str, blocked_ips_or_subnets: Iterable[str],
                    whitelists: dict[str, list[str]] | None = None,
                    generated: str | None = None) -> tuple[str, int, int]:
    """Build a RouterOS .rsc script.

    Whitelisted ranges become `accept` rules placed above the `drop` rules.
    Blocked targets can be individual IPs (e.g. 1.2.3.4/32) or CIDR blocks (e.g. 1.2.3.0/24).

    Returns (script_text, blocked_count, whitelisted_cidr_count).
    """
    whitelists = WHITELISTS if whitelists is None else whitelists
    generated = generated or datetime.now().isoformat(timespec="seconds")

    # Filter out whitelisted entries (check individual IP representation or base subnet)
    def clean_blocked(targets: Iterable[str]) -> list[str]:
        cleaned = []
        for t in targets:
            ip = t.split("/")[0]
            if not is_whitelisted(ip):
                cleaned.append(t)
        return cleaned

    blocked = clean_blocked(set(blocked_ips_or_subnets))
    block_v4, block_v6 = split_cidr_versions(blocked)

    lines: list[str] = [f"# {tag} firewall export - generated {generated}", ""]
    wl_count = 0

    for name, cidrs in whitelists.items():
        nets = [ipaddress.ip_network(c) for c in cidrs]
        wl_count += len(nets)
        for version, cmd in ((4, "/ip firewall address-list"),
                             (6, "/ipv6 firewall address-list")):
            chunk = [n for n in nets if n.version == version]
            if not chunk:
                continue
            lines.append(f"# {name} IPv{version} ranges (whitelisted)")
            lines.append(cmd)
            lines += [f"add list={name} address={n}" for n in chunk]
            lines.append("")

    if block_v4:
        lines.append(f"# Blocked IPv4 ({len(block_v4)}, non-whitelisted)")
        lines.append("/ip firewall address-list")
        lines += [f"add list=Blocked address={t}" for t in block_v4]
        lines.append("")
    if block_v6:
        lines.append(f"# Blocked IPv6 ({len(block_v6)}, non-whitelisted)")
        lines.append("/ipv6 firewall address-list")
        lines += [f"add list=Blocked address={t}" for t in block_v6]
        lines.append("")

    for version, cmd, blocks in ((4, "/ip firewall filter", block_v4),
                                 (6, "/ipv6 firewall filter", block_v6)):
        lines.append(cmd)
        for name in whitelists:
            lines.append(
                f"add chain=input action=accept src-address-list={name} "
                f'comment="{tag} {name.lower()} whitelist (v{version})"')
        if blocks:
            for chain in ("input", "forward"):
                lines.append(
                    f"add chain={chain} action=drop src-address-list=Blocked "
                    f'comment="{tag} blocked (v{version})"')
        lines.append("")

    return "\n".join(lines), len(blocked), wl_count


def write_mikrotik(path: str, tag: str,
                   blocked_ips_or_subnets: Iterable[str]) -> tuple[str, int, int]:
    text, n_blocked, n_whitelisted = mikrotik_script(
        tag, blocked_ips_or_subnets)
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
    return path, n_blocked, n_whitelisted


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def arg_parser(description: str, default_log_path: str,
               default_report_dir: str) -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=description)
    p.add_argument("--log-path", default=default_log_path,
                   help=f"log directory (default: {default_log_path})")
    p.add_argument("--report-dir", default=default_report_dir,
                   help=f"report output dir (default: {default_report_dir})")
    p.add_argument("--top", type=int, default=30,
                   help="rows in the top-IP table (default: 30)")
    p.add_argument("--quiet", action="store_true",
                   help="write files, print nothing but the paths")
    p.add_argument("--no-rsc", action="store_true",
                   help="skip MikroTik .rsc generation")
    return p