#!/usr/bin/env python3
"""qihu-platform auto-block — read nginx access.log, push scanner IPs to blocked_ips.conf.

Usage:  python3 /app/backend/nginx/block_ip.py
Cron:   */5 * * * *  python3 /app/backend/nginx/block_ip.py >> /var/log/qihu-block.log 2>&1
"""
from __future__ import annotations
import ipaddress
import os
import re
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path

ACCESS_LOG = "/var/log/nginx/access.log"
ACCESS_LOG_1 = "/var/log/nginx/access.log.1"
BLOCK_CONF = "/etc/nginx/blocked_ips.conf"
STATE_FILE = "/var/lib/qihu/block_state.json"
WINDOW_SEC = 300  # only consider last 5 min
THRESHOLD_444 = 6
THRESHOLD_4XX = 20
BLOCK_TTL_HOURS = 24

LOG_RE = re.compile(
    r"^(\S+)\s+-\s+\S+\s+\[([^\]]+)\]\s+\"(\S+)\s+(\S+)[^\"]*\"\s+(\d+)\s+\S+\s+\"[^\"]*\"\s+\"([^\"]*)\""
)
SCANNER_UA = re.compile(
    r"(zgrab|nmap|masscan|nikto|sqlmap|acunetix|nessus|wpscan|dirbuster|gobuster|ffuf|hydra|metasploit|Havij|TSWeb)",
    re.IGNORECASE,
)
SCAN_PATH = re.compile(
    r"^/(wp-admin|wp-login|xmlrpc\.php|administrator|phpmyadmin|\.env|\.git/|actuator|server-status|server-info|elmah)",
    re.IGNORECASE,
)


def now_ts() -> int:
    return int(time.time())


def parse_log_line(line: str):
    m = LOG_RE.match(line)
    if not m:
        return None
    ip, ts_str, method, path, status, ua = m.groups()
    return ip, ts_str, method, path, int(status), ua


def nginx_time_to_epoch(s: str) -> int:
    """10/Aug/2026:14:48:20 +0800 -> epoch"""
    try:
        from datetime import datetime
        return int(datetime.strptime(s.split()[0], "%d/%b/%Y:%H:%M:%S").timestamp())
    except Exception:
        return 0


def load_state() -> dict:
    p = Path(STATE_FILE)
    if p.exists():
        try:
            import json
            return json.loads(p.read_text())
        except Exception:
            pass
    return {"blocks": {}, "last_run": 0}


def save_state(state: dict) -> None:
    Path(STATE_FILE).parent.mkdir(parents=True, exist_ok=True)
    import json
    Path(STATE_FILE).write_text(json.dumps(state, indent=2, sort_keys=True))


def is_public_ip(ip: str) -> bool:
    """Skip private / loopback / container IPs."""
    try:
        a = ipaddress.ip_address(ip)
        return not (a.is_private or a.is_loopback or a.is_link_local or a.is_reserved or a.is_multicast)
    except ValueError:
        return False


def scan() -> dict:
    """Return {ip: {444: n, 4xx: n, ua_hit: bool, path_hit: bool}}."""
    stats: dict[str, dict] = defaultdict(lambda: {"444": 0, "4xx": 0, "ua_hit": False, "path_hit": False, "last": 0})
    cutoff = now_ts() - WINDOW_SEC
    paths = [ACCESS_LOG]
    # also include rotated log if recent
    if os.path.exists(ACCESS_LOG_1):
        paths.append(ACCESS_LOG_1)
    for p in paths:
        if not os.path.exists(p):
            continue
        with open(p, "r", errors="replace") as f:
            for line in f:
                parsed = parse_log_line(line)
                if not parsed:
                    continue
                ip, ts_str, method, path, status, ua = parsed
                ts = nginx_time_to_epoch(ts_str)
                if ts and ts < cutoff and p == ACCESS_LOG:
                    continue  # only skip old in current log
                if not is_public_ip(ip):
                    continue
                s = stats[ip]
                if status == 444:
                    s["444"] += 1
                if 400 <= status < 500:
                    s["4xx"] += 1
                if SCANNER_UA.search(ua):
                    s["ua_hit"] = True
                if SCAN_PATH.match(path):
                    s["path_hit"] = True
                s["last"] = max(s["last"], ts)
    return stats


def add_to_blocked(ip: str) -> None:
    """Append `deny <ip>;` to blocked_ips.conf (dedup)."""
    with open(BLOCK_CONF, "r") as f:
        existing = f.read()
    rule = f"deny {ip};"
    if rule in existing:
        return
    with open(BLOCK_CONF, "a") as f:
        f.write(f"\n# auto {time.strftime('%Y-%m-%dT%H:%M:%S')}\n{rule}\n")


def reload_nginx() -> bool:
    try:
        subprocess.run(["nginx", "-t"], check=True, capture_output=True, timeout=10)
        subprocess.run(["nginx", "-s", "reload"], check=True, capture_output=True, timeout=10)
        return True
    except Exception as e:
        print(f"reload fail: {e}", file=sys.stderr)
        return False


def main() -> int:
    state = load_state()
    blocks = state.get("blocks", {})
    stats = scan()
    new_blocks = 0
    for ip, s in stats.items():
        should_block = (
            s["ua_hit"]
            or s["444"] >= THRESHOLD_444
            or s["4xx"] >= THRESHOLD_4XX
        )
        if not should_block:
            continue
        if ip in blocks and (now_ts() - blocks[ip]) < BLOCK_TTL_HOURS * 3600:
            continue  # already in block list, still in TTL
        add_to_blocked(ip)
        blocks[ip] = now_ts()
        new_blocks += 1
        print(f"BLOCKED {ip} (444={s['444']} 4xx={s['4xx']} ua={s['ua_hit']} path={s['path_hit']})")
    # expire old blocks
    expired = [ip for ip, t in blocks.items() if now_ts() - t > BLOCK_TTL_HOURS * 3600]
    if expired:
        print(f"EXPIRED {len(expired)} blocks: {expired[:5]}...")
        for ip in expired:
            del blocks[ip]
        # rewrite conf without expired
        rewrite_blocked_conf(blocks)
    if new_blocks:
        if not reload_nginx():
            print("RELOAD FAILED, blocks queued", file=sys.stderr)
            return 1
    state["blocks"] = blocks
    state["last_run"] = now_ts()
    save_state(state)
    print(f"DONE scanned={len(stats)} new_blocks={new_blocks} expired={len(expired)}")
    return 0


def rewrite_blocked_conf(blocks: dict) -> None:
    """Rewrite blocked_ips.conf with current active blocks + deny directives."""
    lines = [
        "# Auto-managed deny list (see /app/backend/nginx/block_ip.py)",
        "# Format: deny 1.2.3.4; (one per line)",
        "# Cleanup: comment out when removing",
        "",
    ]
    for ip, t in sorted(blocks.items()):
        lines.append(f"# auto {time.strftime('%Y-%m-%dT%H:%M:%S', time.localtime(t))}")
        lines.append(f"deny {ip};")
    Path(BLOCK_CONF).write_text("\n".join(lines) + "\n")


if __name__ == "__main__":
    sys.exit(main())