#!/usr/bin/env python3
import ipaddress
import json
import os
import pathlib
import re
import subprocess
import sys

STATE_DIR = pathlib.Path("/run/lemoa-network-scanner")
STATE_FILE = STATE_DIR / "state.json"


def fail(message):
    print(message, file=sys.stderr)
    raise SystemExit(2)


def command(args):
    result = subprocess.run(args, text=True, capture_output=True)
    if result.returncode:
        fail(result.stderr.strip() or f"Command failed: {' '.join(args)}")
    return result


def valid_interface(name):
    return (
        name != "lo"
        and name.replace("-", "").replace("_", "").isalnum()
        and (pathlib.Path("/sys/class/net") / name).exists()
    )


def valid_network(range_text):
    try:
        network = ipaddress.ip_network(range_text, strict=False)
    except ValueError as exc:
        fail(str(exc))
    if network.version != 4 or network.prefixlen < 16 or network.prefixlen > 30:
        fail("IPv4 range must be between /16 and /30")
    return network


def valid_scan_target(target_text):
    """Return validated Nmap target arguments for a CIDR or bounded IPv4 range."""
    short_match = re.fullmatch(r"(\d{1,3}\.\d{1,3}\.\d{1,3}\.)(\d{1,3})-(\d{1,3})", target_text)
    full_match = re.fullmatch(r"(\d{1,3}(?:\.\d{1,3}){3})-(\d{1,3}(?:\.\d{1,3}){3})", target_text)
    if not short_match and not full_match:
        return [str(valid_network(target_text))]
    try:
        if short_match:
            start = ipaddress.ip_address(f"{short_match.group(1)}{short_match.group(2)}")
            end = ipaddress.ip_address(f"{short_match.group(1)}{short_match.group(3)}")
        else:
            start = ipaddress.ip_address(full_match.group(1))
            end = ipaddress.ip_address(full_match.group(2))
    except ValueError as exc:
        fail(str(exc))
    count = int(end) - int(start) + 1
    if end < start or count > 65536:
        fail("Scan range must increase and contain no more than 65,536 addresses")
    return [str(block) for block in ipaddress.summarize_address_range(start, end)]


def restore():
    if os.geteuid() != 0:
        fail("Administrator privileges are required")
    if not STATE_FILE.exists():
        print("No temporary address is active")
        return
    state = json.loads(STATE_FILE.read_text())
    device = state["device"]
    cidr = state["cidr"]
    if valid_interface(device):
        subprocess.run(["ip", "address", "del", cidr, "dev", device], text=True, capture_output=True)
    STATE_FILE.unlink(missing_ok=True)
    print(f"Removed {cidr} from {device}")


def apply(device, range_text, address_text):
    if os.geteuid() != 0:
        fail("Administrator privileges are required")
    if not valid_interface(device):
        fail("Invalid network interface")
    network = valid_network(range_text)
    address = ipaddress.ip_address(address_text)
    if address not in network or address in {network.network_address, network.broadcast_address}:
        fail("Temporary address is not usable in the selected range")
    restore()
    cidr = f"{address}/{network.prefixlen}"
    command(["ip", "link", "set", "dev", device, "up"])
    # IPv4 labels share Linux's 15-character interface-name limit. Appending
    # ':lemoa-scan' breaks predictable USB names such as enx3cab72d591a0 and
    # produces the misleading "Attribute failed policy validation" error.
    # The state file already retains the exact device and CIDR for cleanup.
    command(["ip", "address", "add", cidr, "dev", device])
    STATE_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
    STATE_FILE.write_text(json.dumps({"device": device, "cidr": cidr}))
    os.chmod(STATE_FILE, 0o600)
    print(f"Added temporary address {cidr} to {device}")


def scan(device, range_text):
    if os.geteuid() != 0:
        fail("Administrator privileges are required")
    if not valid_interface(device):
        fail("Invalid network interface")
    targets = valid_scan_target(range_text)
    # Match `sudo nmap -sn TARGET`: privileged Nmap chooses ARP on a directly
    # connected network and retains its normal retry/timeout behavior.
    os.execv(
        "/usr/bin/nmap",
        ["nmap", "-sn", "-n", "-e", device, "-oX", "-", *targets],
    )


def main():
    if len(sys.argv) < 2:
        fail("Usage: helper apply INTERFACE CIDR ADDRESS | scan INTERFACE CIDR | restore")
    if sys.argv[1] == "restore" and len(sys.argv) == 2:
        restore()
    elif sys.argv[1] == "apply" and len(sys.argv) == 5:
        apply(sys.argv[2], sys.argv[3], sys.argv[4])
    elif sys.argv[1] == "scan" and len(sys.argv) == 4:
        scan(sys.argv[2], sys.argv[3])
    else:
        fail("Invalid arguments")


if __name__ == "__main__":
    main()
