#! /usr/bin/env python3
import argparse
import sys
import logging
import subprocess

logger = logging.getLogger(__name__)

PROFILES = {
    "minecraft": {
        "comment": "Minecraft Server",
        "ports": [(25565, "tcp"), (25565, "udp")],
    },
    "unturned": {
        "comment": "Unturned Server",
        "ports": [(27015, "tcp"), (27015, "udp"), (27016, "tcp"), (27016, "udp")],
    },
    "zola": {
        "comment": "Zola",
        "ports": [(1111, "tcp")],
    },
    "kdeconnect": {
        "comment": "KDE Connect",
        "ports": [("1714:1764", "udp"), ("1714:1764", "tcp")],
    },
    "asterisk": {
        "comment": "Asterisk",
        "ports": [
            (5060, "udp"),
            (4569, "udp"),
            (4569, "tcp"),
            (5060, "tcp"),
            (5061, "tcp"),
            ("10000:20000", "udp"),
        ],
    },
    "elproxy": {
        "comment": "ElProxy",
        "ports": [(8100, "tcp"), (5198, "udp"), (5199, "udp"), (5200, "tcp")],
    },
}


def main() -> int:
    # Handle program arguments
    ap = argparse.ArgumentParser(
        prog="ufw-gen", description="Generate UFW allow commands"
    )
    ap.add_argument(
        "profile",
        help="Profile to generate UFW allow commands for",
        choices=PROFILES.keys(),
    )
    ap.add_argument("--from", help="Source IP address", default="any", dest="source")
    ap.add_argument("--to", help="Destination IP address", default="any", dest="dest")
    ap.add_argument(
        "--no-sudo", help="Don't prefix commands with sudo", action="store_true"
    )
    ap.add_argument(
        "--dry-run", help="Generate UFW commands in dry-run mode", action="store_true"
    )
    ap.add_argument(
        "--execute", "-x", help="Execute generated UFW commands", action="store_true"
    )
    ap.add_argument(
        "-v", "--verbose", help="Enable verbose logging", action="store_true"
    )
    args = ap.parse_args()

    # Configure logging
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s:	%(message)s",
    )

    # Generate UFW allow commands
    profile = PROFILES[args.profile]
    for port, proto in profile["ports"]:
        sudo_str = "" if args.no_sudo else "sudo "
        dry_run_str = "--dry-run" if args.dry_run else ""

        # Generate the command
        command = ["sudo", "ufw"] if not args.no_sudo else ["ufw"]
        if args.dry_run:
            command.append("--dry-run")
        command.extend(
            [
                "allow",
                "from",
                args.source,
                "to",
                args.dest,
                "port",
                str(port),
                "proto",
                proto,
                "comment",
                profile["comment"],
            ]
        )

        # Run
        print(" ".join(command))
        if args.execute:
            result = subprocess.run(command).returncode
            if result != 0:
                logger.error("Failed to run command: %s", command)
                return result

    return 0


if __name__ == "__main__":
    sys.exit(main())