#!/usr/bin/env python3
import os
import psutil
import json
import time
import shutil
import subprocess
from datetime import datetime
from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, ContextTypes

# ==============================
# CONFIGURAÇÃO BÁSICA
# ==============================
CONFIG_FILE = "/srv/sistema/bot/config.json"
LOG_FILE = "/srv/sistema/bot/attacks/attack_log.txt"

os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

with open(CONFIG_FILE, "r") as f:
    cfg = json.load(f)

BOT_TOKEN = cfg["BOT_TOKEN"]
OWNER_ID = int(cfg["OWNER_ID"])

# Portas seguras (acesso normal)
SAFE_PORTS = [80, 443, 51820]

# Portas sensíveis (tentativa repetida = suspeita)
DANGEROUS_PORTS = [22, 3050, 3060, 21, 23, 3306, 3389, 5900, 25565]

# Países
COUNTRY_CACHE = {}
ATTEMPTS = {}       # chave "ip:porta" -> contagem
BLOCKED_IPS = set() # já bloqueados


# ==============================
# HELPERS DE REDE / FIREWALL
# ==============================
def normalize_ip(ip: str) -> str:
    ip = ip.strip()
    if ip.startswith("[") and ip.endswith("]"):
        ip = ip[1:-1]
    if ip.startswith("::ffff:"):
        ip = ip.split("::ffff:")[-1]
    return ip


def is_internal(ip: str) -> bool:
    ip = normalize_ip(ip)
    return (
        ip.startswith("10.") or
        ip.startswith("192.168.") or
        ip.startswith("172.16.") or
        ip.startswith("172.17.") or
        ip.startswith("172.18.") or
        ip.startswith("172.19.") or
        ip.startswith("172.20.") or
        ip.startswith("172.21.") or
        ip.startswith("172.22.") or
        ip.startswith("172.23.") or
        ip.startswith("172.24.") or
        ip.startswith("172.25.") or
        ip.startswith("172.26.") or
        ip.startswith("172.27.") or
        ip.startswith("172.28.") or
        ip.startswith("172.29.") or
        ip.startswith("172.30.") or
        ip.startswith("172.31.")
    )


def get_country(ip: str):
    """Retorna (code, name). Ex: ('BR', 'Brazil')"""
    ip = normalize_ip(ip)

    if ip in COUNTRY_CACHE:
        return COUNTRY_CACHE[ip]

    try:
        out = subprocess.check_output(
            ["geoiplookup", ip],
            stderr=subprocess.STDOUT
        ).decode()
    except Exception:
        COUNTRY_CACHE[ip] = ("UNK", "Desconhecido")
        return COUNTRY_CACHE[ip]

    if "GeoIP Country Edition" in out and ":" in out:
        part = out.split(":", 1)[1].strip()
        # normalmente: "BR, Brazil"
        bits = [b.strip() for b in part.split(",")]
        if len(bits) >= 2:
            code = bits[0]
            name = ", ".join(bits[1:])
        else:
            code = "UNK"
            name = part
        COUNTRY_CACHE[ip] = (code, name)
    else:
        COUNTRY_CACHE[ip] = ("UNK", "Desconhecido")

    return COUNTRY_CACHE[ip]


def log_block(ip: str, reason: str):
    ip = normalize_ip(ip)
    with open(LOG_FILE, "a") as f:
        f.write(f"{datetime.now()} - BLOQUEADO: {ip} ({reason})\n")


def block_ip(ip: str, reason: str):
    ip = normalize_ip(ip)
    if ip in BLOCKED_IPS:
        # já bloqueado, só loga motivo novo
        log_block(ip, reason)
        return

    BLOCKED_IPS.add(ip)
    # regra UFW (somente bloqueia entrada desse IP)
    subprocess.call(["ufw", "deny", "from", ip])
    log_block(ip, reason)


def scan_connections():
    """
    Analisa conexões com ss -tun e retorna lista de 'eventos' para o monitor.
    Cada evento = dict com:
    type: 'blocked_international' | 'blocked_sensitive' | 'warn_sensitive'
    ip, port, country_code, country_name, attempts
    """
    events = []

    try:
        out = subprocess.check_output(
            ["ss", "-tun"],
            stderr=subprocess.STDOUT
        ).decode()
    except Exception:
        return events

    lines = out.splitlines()
    if len(lines) <= 1:
        return events

    for line in lines[1:]:
        if not line.strip():
            continue

        parts = line.split()
        if len(parts) < 6:
            continue

        local = parts[4]  # Local Address:Port (onde o servidor está escutando)
        peer = parts[5]   # Peer Address:Port (quem conectou)

        if ":" not in local or ":" not in peer:
            continue

        try:
            local_ip, local_port_str = local.rsplit(":", 1)
            peer_ip, peer_port_str = peer.rsplit(":", 1)
        except ValueError:
            continue

        peer_ip = normalize_ip(peer_ip)

        # porta do SERVIDOR (é essa que importa, não a do cliente)
        try:
            local_port = int(local_port_str)
        except ValueError:
            continue

        # ignora LAN / loopback
        if is_internal(peer_ip) or peer_ip.startswith("127.") or peer_ip.startswith("::1"):
            continue

        country_code, country_name = get_country(peer_ip)
        key = f"{peer_ip}:{local_port}"

        # 1) QUALQUER IP que não seja BR é bloqueado (independente da porta)
        if country_code != "BR":
            reason = f"Acesso internacional na porta {local_port}"
            block_ip(peer_ip, reason)
            events.append({
                "type": "blocked_international",
                "ip": peer_ip,
                "port": local_port,
                "country_code": country_code,
                "country_name": country_name,
                "attempts": 1
            })
            continue

        # 2) Se for Brasil: ignora portas normais (HTTP/HTTPS/WireGuard)
        if local_port in SAFE_PORTS:
            continue

        # 3) Portas sensíveis (SSH, Firebird, etc.) — só Brasil
        if local_port in DANGEROUS_PORTS:
            ATTEMPTS[key] = ATTEMPTS.get(key, 0) + 1
            attempts = ATTEMPTS[key]

            if attempts >= 3:
                reason = f"Muitas tentativas na porta sensível {local_port}"
                block_ip(peer_ip, reason)
                events.append({
                    "type": "blocked_sensitive",
                    "ip": peer_ip,
                    "port": local_port,
                    "country_code": country_code,
                    "country_name": country_name,
                    "attempts": attempts
                })
                ATTEMPTS[key] = 0
            else:
                events.append({
                    "type": "warn_sensitive",
                    "ip": peer_ip,
                    "port": local_port,
                    "country_code": country_code,
                    "country_name": country_name,
                    "attempts": attempts
                })

    return events


# ==============================
# INFO DE SISTEMA
# ==============================
def get_uptime():
    uptime_seconds = time.time() - psutil.boot_time()
    hours = int(uptime_seconds // 3600)
    minutes = int((uptime_seconds % 3600) // 60)
    return f"{hours}h {minutes}min"


# ==============================
# HANDLERS DO BOT
# ==============================
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    await update.message.reply_text(
        "🔥 *Ultimate Defender PRO Online!*\n\n"
        "COMANDOS DISPONÍVEIS:\n"
        "• /status  – Status completo do servidor\n"
        "• /attacks – Últimos ataques registrados\n"
        "• /live    – Monitor inteligente em tempo real\n"
        "• /who     – Conexões ativas\n"
        "• /ports   – Portas em escuta\n"
        "• /ufw     – Status do firewall UFW\n"
        "• /ping    – Teste rápido\n"
        "• /scan    – Scan ClamAV (se instalado)\n",
        parse_mode="Markdown"
    )


async def status(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    cpu = psutil.cpu_percent()
    ram = psutil.virtual_memory().percent
    disk = psutil.disk_usage("/").percent
    uptime = get_uptime()

    msg = (
        "🖥 *STATUS DO SERVIDOR*\n\n"
        f"• CPU: `{cpu}%`\n"
        f"• RAM: `{ram}%`\n"
        f"• DISCO: `{disk}%`\n"
        f"• UPTIME: `{uptime}`\n"
    )

    await update.message.reply_text(msg, parse_mode="Markdown")


async def attacks(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    if not os.path.exists(LOG_FILE):
        await update.message.reply_text("Nenhum ataque registrado ainda.")
        return

    data = open(LOG_FILE).read().strip()
    if data == "":
        await update.message.reply_text("Nenhum ataque registrado ainda.")
    else:
        await update.message.reply_text(
            f"📌 *Últimos ataques:*\n\n```{data[-2000:]}```",
            parse_mode="Markdown"
        )


async def who(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    out = subprocess.getoutput("ss -tun | head -n 20")
    await update.message.reply_text(
        f"🌐 *Conexões ativas (top 20):*\n\n```{out}```",
        parse_mode="Markdown"
    )


async def ports(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    out = subprocess.getoutput("ss -tuln | grep LISTEN || ss -tuln")
    await update.message.reply_text(
        f"📡 *Portas em escuta:*\n\n```{out}```",
        parse_mode="Markdown"
    )


async def ufw(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    out = subprocess.getoutput("ufw status verbose")
    await update.message.reply_text(
        f"🧱 *Firewall UFW:*\n\n```{out}```",
        parse_mode="Markdown"
    )


async def ping(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    now = datetime.now().strftime("%H:%M:%S")
    await update.message.reply_text(f"🏓 Pong! {now}")


async def scan(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """
    Scan antivírus básico usando ClamAV (se tiver instalado).
    """
    if update.effective_user.id != OWNER_ID:
        return

    if shutil.which("clamscan") is None:
        await update.message.reply_text(
            "⚠️ ClamAV não está instalado.\n"
            "Instale com: `sudo apt install clamav`",
            parse_mode="Markdown"
        )
        return

    await update.message.reply_text(
        "🧬 Iniciando scan ClamAV em /srv/sistema...\n"
        "Isso pode demorar, vou te mandar o final do relatório.",
    )

    # scan recursivo em /srv/sistema (ajuste o caminho se quiser)
    out = subprocess.getoutput("clamscan -r /srv/sistema | tail -n 25")

    await update.message.reply_text(
        f"🧪 *Resultado ClamAV (últimas linhas):*\n\n```{out}```",
        parse_mode="Markdown"
    )


# ==============================
# MONITOR INTELIGENTE (/live)
# ==============================
async def live_job(context: ContextTypes.DEFAULT_TYPE):
    chat_id = context.job.data["chat_id"]
    events = scan_connections()

    for e in events:
        if e["type"] == "blocked_international":
            text = (
                "🚫 *ATAQUE BLOQUEADO (INTERNACIONAL)*\n"
                f"IP: `{e['ip']}`\n"
                f"País: {e['country_name']} ({e['country_code']})\n"
                f"Porta alvo: {e['port']}\n"
                f"Motivo: Acesso internacional bloqueado."
            )
        elif e["type"] == "blocked_sensitive":
            text = (
                "🚫 *BLOQUEADO (PORTA SENSÍVEL)*\n"
                f"IP: `{e['ip']}`\n"
                f"País: {e['country_name']} ({e['country_code']})\n"
                f"Porta alvo: {e['port']}\n"
                f"Tentativas: {e['attempts']}\n"
                f"Motivo: Múltiplas tentativas em porta sensível."
            )
        else:  # warn_sensitive
            text = (
                "⚠️ *Acesso suspeito (Brasil)*\n"
                f"IP: `{e['ip']}`\n"
                f"País: {e['country_name']} ({e['country_code']})\n"
                f"Porta alvo: {e['port']}\n"
                f"Tentativas: {e['attempts']}/3\n"
                f"Motivo: Tentativa em porta sensível."
            )

        await context.bot.send_message(
            chat_id=chat_id,
            text=text,
            parse_mode="Markdown"
        )


async def live(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if update.effective_user.id != OWNER_ID:
        return

    chat_id = update.effective_chat.id

    # Remove monitor antigo desse chat, se existir
    for job in context.job_queue.get_jobs_by_name(f"live_{chat_id}"):
        job.schedule_removal()

    context.job_queue.run_repeating(
        live_job,
        interval=3,
        first=0,
        name=f"live_{chat_id}",
        data={"chat_id": chat_id},
    )

    await update.message.reply_text(
        "📡 Monitor inteligente ativado!\n"
        "Vou avisar só quando tiver coisa estranha de verdade "
        "(internacional ou porta sensível)."
    )


# ==============================
# LOOP PRINCIPAL
# ==============================
def run_bot():
    app = ApplicationBuilder().token(BOT_TOKEN).build()

    app.add_handler(CommandHandler("start", start))
    app.add_handler(CommandHandler("help", start))
    app.add_handler(CommandHandler("status", status))
    app.add_handler(CommandHandler("attacks", attacks))
    app.add_handler(CommandHandler("live", live))
    app.add_handler(CommandHandler("who", who))
    app.add_handler(CommandHandler("ports", ports))
    app.add_handler(CommandHandler("ufw", ufw))
    app.add_handler(CommandHandler("ping", ping))
    app.add_handler(CommandHandler("scan", scan))

    print("🔥 Ultimate Defender PRO rodando…")
    app.run_polling()


if __name__ == "__main__":
    run_bot()