import json
import subprocess
from pathlib import Path
from datetime import datetime

# --- CONFIG ---
MWAC_DIR = Path(r"C:\ProgramData\Malwarebytes\MBAMService\MwacDetections")
RULE_NAME_PREFIX = "MBAM_AutoBlock_"
FIREWALL_DIRECTION = "Inbound"
FIREWALL_PROFILE = "Any"
REPORT_FILE = Path("mbam_block_report.txt")
# ---------------


def is_hex_header(line: str) -> bool:
    line = line.strip()
    return len(line) == 64 and all(c in "0123456789ABCDEFabcdef" for c in line)


def load_ips_from_folder(folder: Path) -> set[str]:
    ips = set()

    files = list(folder.glob("*.json"))
    print(f"[*] Found {len(files)} JSON files")

    for file in files:
        try:
            raw = file.read_text(encoding="utf-8", errors="ignore")
        except Exception:
            continue

        lines = raw.splitlines()

        # Remove the hash header if present
        if lines and is_hex_header(lines[0]):
            json_text = "\n".join(lines[1:])
        else:
            json_text = raw

        try:
            data = json.loads(json_text)
        except Exception:
            continue

        threats = data.get("threats", [])
        for t in threats:
            main = t.get("mainTrace", {})
            website = main.get("websiteData", {})
            ip = website.get("ip")
            if ip:
                ips.add(ip)

    return ips


def load_existing_firewall_ips() -> set[str]:
    """
    Load ALL firewall rules once and extract all RemoteAddress IPs.
    """
    print("[*] Loading existing firewall rules...")

    cmd = [
        "powershell",
        "-Command",
        (
            "Get-NetFirewallRule | "
            "Get-NetFirewallAddressFilter | "
            "Select-Object -ExpandProperty RemoteAddress"
        ),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print("[!] Failed to load firewall rules")
        return set()

    raw_addresses = result.stdout.splitlines()

    # Normalize and filter valid IPv4 addresses
    ips = {addr.strip() for addr in raw_addresses if "." in addr}

    print(f"[*] Found {len(ips)} existing firewall IP entries")
    return ips


def add_firewall_rule(ip: str):
    rule_name = f"{RULE_NAME_PREFIX}{ip.replace('.', '_')}"
    cmd = [
        "powershell",
        "-Command",
        (
            "New-NetFirewallRule "
            f"-DisplayName '{rule_name}' "
            f"-Direction {FIREWALL_DIRECTION} "
            "-Action Block "
            f"-RemoteAddress {ip} "
            f"-Profile {FIREWALL_PROFILE}"
        ),
    ]

    print(f"[+] Adding firewall block rule for {ip}")
    subprocess.run(cmd, capture_output=True, text=True)


def export_report(all_ips: set[str], new_rules: list[str]):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    with REPORT_FILE.open("w", encoding="utf-8") as f:
        f.write("=== Malwarebytes → Windows Firewall Block Report ===\n")
        f.write(f"Generated: {timestamp}\n\n")
        f.write(f"Total unique IPs detected: {len(all_ips)}\n")
        f.write(f"New firewall rules added: {len(new_rules)}\n\n")

        if new_rules:
            f.write("Newly added rules:\n")
            for ip in new_rules:
                f.write(f"  - {ip}\n")
            f.write("\n")

        f.write("All detected IPs:\n")
        for ip in sorted(all_ips):
            f.write(f"  - {ip}\n")

    print(f"\n[*] Report saved to: {REPORT_FILE.resolve()}")


def main():
    print("[*] Scanning Malwarebytes MwacDetections folder...")
    ips = load_ips_from_folder(MWAC_DIR)
    print(f"[*] Extracted {len(ips)} IPs from logs")

    existing_ips = load_existing_firewall_ips()

    new_rules = []

    for ip in sorted(ips):
        if ip in existing_ips:
            print(f"[-] Already blocked: {ip}")
        else:
            add_firewall_rule(ip)
            new_rules.append(ip)

    export_report(ips, new_rules)

    print("[*] Done.")
    input("\nPress Enter to exit...")


if __name__ == "__main__":
    main()