#!/usr/bin/env python3
"""Manage Windows Firewall rule for WikiServ and show local/public IPs.

Usage examples:
  python manage_firewall.py add --port 8000 --name "WikiServ HTTP"
  python manage_firewall.py remove --name "WikiServ HTTP"
  python manage_firewall.py status --name "WikiServ HTTP"
  python manage_firewall.py ips

Note: add/remove require administrator privileges on Windows.
"""
from __future__ import annotations

import argparse
import json
import socket
import subprocess
import sys
import urllib.request
from typing import Tuple
import ctypes
import platform
from pathlib import Path


def _get_local_ip() -> str:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
    except Exception:
        ip = "127.0.0.1"
    finally:
        try:
            s.close()
        except Exception:
            pass
    return ip


def _get_public_ip(timeout: float = 3.0) -> str:
    urls = [
        "https://api.ipify.org?format=json",
        "https://ifconfig.co/json",
        "https://ipinfo.io/json",
    ]
    for u in urls:
        try:
            with urllib.request.urlopen(u, timeout=timeout) as r:
                data = r.read().decode("utf-8", errors="ignore")
                try:
                    j = json.loads(data)
                    if "ip" in j:
                        return j["ip"]
                    if "ip_address" in j:
                        return j["ip_address"]
                    if "address" in j:
                        return j["address"]
                except Exception:
                    # some services return plain text
                    text = data.strip()
                    if text:
                        return text
        except Exception:
            continue
    return "(unknown)"


def _run_powershell(cmd: str) -> Tuple[int, str, str]:
    try:
        proc = subprocess.run([
            "powershell",
            "-NoProfile",
            "-NonInteractive",
            "-Command",
            cmd,
        ], capture_output=True, text=True, check=False)
        return proc.returncode, proc.stdout.strip(), proc.stderr.strip()
    except FileNotFoundError:
        return 127, "", "powershell not found"


def add_rule(name: str, port: int) -> None:
    ps = (
        f"New-NetFirewallRule -DisplayName '{name}' "
        f"-Direction Inbound -Action Allow -Protocol TCP -LocalPort {port}"
    )
    rc, out, err = _run_powershell(ps)
    if rc == 0:
        print(f"Firewall rule '{name}' added for TCP port {port}.")
    else:
        # fallback to netsh
        netsh = (
            f"netsh advfirewall firewall add rule name=\"{name}\" dir=in action=allow protocol=TCP localport={port}"
        )
        try:
            proc = subprocess.run(netsh, shell=True, capture_output=True, text=True)
            if proc.returncode == 0:
                print(f"Firewall rule '{name}' added (netsh) for TCP port {port}.")
                return
        except Exception as e:
            print(f"Failed to add rule: {e}")
        print("Failed to add firewall rule via PowerShell. stderr:", err)


def _is_admin() -> bool:
    """Return True if running with administrative privileges on Windows."""
    if platform.system().lower() != "windows":
        return False
    try:
        return ctypes.windll.shell32.IsUserAnAdmin() != 0
    except Exception:
        return False


def _elevate_and_reexec() -> int:
    """Try to relaunch the current script elevated via PowerShell Start-Process.

    Returns the subprocess return code (0 if launcher succeeded), or 127 if PowerShell not found.
    The elevated process will run independently; this process should exit after calling this.
    """
    if platform.system().lower() != "windows":
        print("Elevation is supported only on Windows.")
        return 1

    python_exe = sys.executable
    script = Path(__file__).resolve()
    args = sys.argv[1:]
    # Build ArgumentList for PowerShell: comma-separated quoted items
    items = [f"'{str(script)}'"] + [f"'{a.replace("'", "''")}'" for a in args]
    arglist = ",".join(items)
    ps_cmd = f"Start-Process -FilePath '{python_exe}' -ArgumentList {arglist} -Verb RunAs"
    try:
        proc = subprocess.run(["powershell", "-NoProfile", "-NonInteractive", "-Command", ps_cmd], check=False)
        return proc.returncode
    except FileNotFoundError:
        return 127


def remove_rule(name: str) -> None:
    ps = f"Get-NetFirewallRule -DisplayName '{name}' | Remove-NetFirewallRule -ErrorAction SilentlyContinue"
    rc, out, err = _run_powershell(ps)
    if rc == 0:
        print(f"Firewall rule '{name}' removed (PowerShell attempt).")
    else:
        netsh = f"netsh advfirewall firewall delete rule name=\"{name}\""
        try:
            proc = subprocess.run(netsh, shell=True, capture_output=True, text=True)
            if proc.returncode == 0:
                print(f"Firewall rule '{name}' removed (netsh).")
                return
        except Exception as e:
            print(f"Failed to remove rule: {e}")
        print("Failed to remove firewall rule via PowerShell. stderr:", err)


def status_rule(name: str) -> None:
    ps = f"Get-NetFirewallRule -DisplayName '{name}' | Select-Object DisplayName,Enabled,Direction,Action,Profile | ConvertTo-Json -Depth 2"
    rc, out, err = _run_powershell(ps)
    if rc == 0 and out:
        try:
            j = json.loads(out)
            # if single object, json module gives dict; if multiple, list
            if isinstance(j, dict):
                items = [j]
            else:
                items = j
            for it in items:
                print(json.dumps(it, ensure_ascii=False, indent=2))
            return
        except Exception:
            print(out)
            return
    # fallback to netsh
    try:
        proc = subprocess.run(f"netsh advfirewall firewall show rule name=\"{name}\"", shell=True, capture_output=True, text=True)
        if proc.returncode == 0:
            print(proc.stdout.strip())
            return
    except Exception as e:
        print(f"Failed to query rule: {e}")
    print("Rule not found or unable to query. stderr:", err)


def print_ips() -> None:
    local = _get_local_ip()
    public = _get_public_ip()
    print(f"Local IP: {local}")
    print(f"Public IP: {public}")


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="Manage Windows Firewall rule and show IPs for WikiServ")
    sub = parser.add_subparsers(dest="cmd")

    p_add = sub.add_parser("add", help="Add firewall rule")
    p_add.add_argument("--port", type=int, default=8000)
    p_add.add_argument("--name", default="WikiServ HTTP")

    p_rm = sub.add_parser("remove", help="Remove firewall rule")
    p_rm.add_argument("--name", default="WikiServ HTTP")

    p_st = sub.add_parser("status", help="Show firewall rule status/details")
    p_st.add_argument("--name", default="WikiServ HTTP")

    p_ips = sub.add_parser("ips", help="Show local and public IP addresses")

    args = parser.parse_args(argv)
    if args.cmd in ("add", "remove", "status"):
        # actions that require admin
        if not _is_admin():
            print("Administrative privileges required. Attempting to elevate via UAC...")
            rc = _elevate_and_reexec()
            if rc == 127:
                print("PowerShell not found — please run this command in an elevated PowerShell/Command Prompt manually.")
            else:
                print("Elevation request sent. If accepted, the elevated process will run the requested command.")
            return rc

    if args.cmd == "add":
        add_rule(args.name, args.port)
    elif args.cmd == "remove":
        remove_rule(args.name)
    elif args.cmd == "status":
        status_rule(args.name)
    elif args.cmd == "ips":
        print_ips()
    else:
        parser.print_help()
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())