#!/usr/bin/env python3
"""DDoS detector for public web ports (80/443).
Runs every minute via cron. Counts active connections per IP on 80/443,
alerts Telegram if a single IP exceeds CONN_PER_IP or total exceeds CONN_TOTAL.
State file prevents alert spam (only alerts on threshold crossing).
"""
import os
import re
import subprocess
import datetime

# Load telegram creds from stock-monitor .env (or env)
def load_env(path):
    env = {}
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, v = line.split("=", 1)
                    env[k.strip()] = v.strip()
    except FileNotFoundError:
        pass
    return env

ENV = load_env("/root/stock-monitor/.env")
BOT = os.environ.get("TELEGRAM_BOT_TOKEN") or ENV.get("TELEGRAM_BOT_TOKEN", "")
CHAT = os.environ.get("TELEGRAM_CHAT_ID") or ENV.get("TELEGRAM_CHAT_ID", "5722352727")

PORTS = [80, 443]
CONN_PER_IP = int(os.environ.get("DDOS_PER_IP", "100"))   # alert if 1 IP > this
CONN_TOTAL = int(os.environ.get("DDOS_TOTAL", "500"))      # alert if total on a port > this
STATE_FILE = "/root/ddos-detector/ddos_state.json"

# ports currently exposed to public (from ss output)
def get_connections():
    """Return dict: port -> {ip: count}"""
    out = subprocess.run(
        ["ss", "-tn", "src", ":80", "or", "src", ":443"],
        capture_output=True, text=True
    ).stdout
    result = {p: {} for p in PORTS}
    # ss -tn src :80 or src :443  -> lines like:
    # ESTAB 0 0 127.0.0.1:80 1.2.3.4:55555
    for line in out.splitlines()[1:]:
        parts = line.split()
        if len(parts) < 4:
            continue
        local = parts[3]   # e.g. 0.0.0.0:80 or [::]:443
        remote = parts[4]
        # extract port from local
        m = re.search(r":(\d+)$", local)
        if not m:
            continue
        port = int(m.group(1))
        if port not in result:
            continue
        # extract IP from remote (strip port)
        ip = remote.rsplit(":", 1)[0]
        ip = ip.strip("[]")
        result[port][ip] = result[port].get(ip, 0) + 1
    return result


def send_telegram(msg):
    if not BOT:
        print("NO BOT TOKEN, skip send")
        return
    import urllib.request
    import urllib.parse
    url = f"https://api.telegram.org/bot{BOT}/sendMessage"
    data = urllib.parse.urlencode({"chat_id": CHAT, "text": msg}).encode()
    req = urllib.request.Request(url, data=data, method="POST")
    try:
        urllib.request.urlopen(req, timeout=10)
    except Exception as e:
        print(f"telegram send failed: {e}")


def load_state():
    import json
    try:
        with open(STATE_FILE) as f:
            return json.load(f)
    except (FileNotFoundError, ValueError):
        return {}


def save_state(s):
    import json
    with open(STATE_FILE, "w") as f:
        json.dump(s, f)


def main():
    conns = get_connections()
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    alerts = []
    for port in PORTS:
        ip_map = conns[port]
        total = sum(ip_map.values())
        # top offenders
        top = sorted(ip_map.items(), key=lambda x: -x[1])[:5]
        for ip, cnt in top:
            if cnt > CONN_PER_IP:
                alerts.append(
                    f"⚠️ Port {port}: IP {ip} = {cnt} koneksi (batas {CONN_PER_IP})"
                )
        if total > CONN_TOTAL:
            alerts.append(
                f"⚠️ Port {port}: TOTAL {total} koneksi (batas {CONN_TOTAL})"
            )

    state = load_state()
    new_alert = False
    if alerts:
        # only alert if we weren't already in alert state (avoid spam)
        if not state.get("alerting"):
            new_alert = True
        state["alerting"] = True
        state["last_alerts"] = alerts
    else:
        if state.get("alerting"):
            # send "cleared" message once
            send_telegram(
                f"✅ DDoS alert cleared @ {now}\n"
                f"Port 80: {sum(conns[80].values())} koneksi\n"
                f"Port 443: {sum(conns[443].values())} koneksi\n"
                f"(batas per-IP {CONN_PER_IP}, total {CONN_TOTAL})"
            )
        state["alerting"] = False

    if new_alert:
        top80 = sorted(conns[80].items(), key=lambda x: -x[1])[:5]
        top443 = sorted(conns[443].items(), key=lambda x: -x[1])[:5]
        msg = f"🚨 DDoS ALERT @ {now}\n\n" + "\n".join(alerts) + "\n\n"
        msg += f"Top IP port 80: {top80}\n"
        msg += f"Top IP port 443: {top443}\n"
        msg += f"(batas per-IP {CONN_PER_IP}, total {CONN_TOTAL})"
        send_telegram(msg)

    save_state(state)
    # debug log (no secrets)
    print(f"{now} 80={sum(conns[80].values())} 443={sum(conns[443].values())} alert={state.get('alerting')}")


if __name__ == "__main__":
    main()