#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-

##################################
# NXFILTER - OPENFW UTM
##################################

import os
import sys
import time
import subprocess
import shutil
import pwd
import grp

NXSET = "/var/efw/nxfilter/settings"
NXBIN = "/var/nxfilter/bin"
NXRUN = "/var/run/nxfilter/nxfilter.pid"
DOH_LIST = "/var/efw/nxfilter/servers-doh"
PORT_APACHE = 85
PORT_DNSMASQ = 54

def log(msg):
    """Displays green formatted log messages."""
    print("\033[1;32m[ NXFilter ] {}\033[0m".format(msg))

def fail(msg):
    """Displays error message in red and terminates execution."""
    print("\033[1;31m[ NXFilter ] {}\033[0m".format(msg))
    sys.exit(1)

def load_settings():
    """Loads configuration variables from file (equivalent to . $NXSET)."""
    settings = {}
    if os.path.exists(NXSET):
        with open(NXSET, "r") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    k, v = line.split("=", 1)
                    # Strip quotes if present
                    v = v.strip("\"'")
                    settings[k.strip()] = v.strip()
    
    if "NXFILTER_ENABLE" not in settings:
        fail("NXFILTER_ENABLE not set")
    if "DROP_DOH" not in settings:
        fail("DROP_DOH not set")

    return settings["NXFILTER_ENABLE"], settings["DROP_DOH"]

enabled, drop_doh = load_settings()

def nxfilter_is_up():
    """Checks if the NXFilter service responds to ping.sh."""
    try:
        p = subprocess.Popen([os.path.join(NXBIN, "ping.sh")], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, _ = p.communicate()
        return "Still working" in out
    except Exception:
        return False

def get_apache_listening():
    """Checks if Apache is listening on the configured port."""
    try:
        p = subprocess.Popen(["netstat", "-lnt"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, _ = p.communicate()
        target = ":{}".format(PORT_APACHE)
        for line in out.splitlines():
            parts = line.split()
            if len(parts) >= 6 and parts[5] == "LISTEN":
                if parts[3].endswith(target):
                    return True
        return False
    except Exception:
        return False

def get_dnsmasq_listening():
    """Checks if dnsmasq is listening on the configured port."""
    try:
        p = subprocess.Popen(["netstat", "-lnt"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, _ = p.communicate()
        target = ":{}".format(PORT_DNSMASQ)
        for line in out.splitlines():
            parts = line.split()
            if len(parts) >= 6 and parts[5] == "LISTEN":
                if parts[3].endswith(target):
                    return True
        return False
    except Exception:
        return False

def get_wpad_port():
    """Checks WPAD port configuration inside dhcpd.conf.tmpl."""
    target_str = 'option wpad "http://.*:{}/proxy\\.pac"'.format(PORT_APACHE)
    if os.path.exists("/etc/dhcpd.conf.tmpl"):
        try:
            import re
            with open("/etc/dhcpd.conf.tmpl", "r") as f:
                content = f.read()
                return bool(re.search(target_str, content))
        except Exception:
            return False
    return False

#########################################
# Apache Port Handling
#########################################
def set_apache_port():
    if enabled == "on":
        if get_apache_listening():
            log("Apache already listening on port {} — no restart needed".format(PORT_APACHE))
            return

        log("Configuring Apache for NXFilter (port 85)")
        shutil.copyfile("/etc/nxfilter/tmpl/nx-httpd.conf.tmpl", "/etc/httpd/sites-available/000-default.conf.tmpl")
        shutil.copyfile("/etc/nxfilter/tmpl/nx-httpd.conf.tmpl", "/etc/httpd/sites-available/000-default.conf")
        shutil.copyfile("/etc/nxfilter/tmpl/nx-ports.conf.tmpl", "/etc/httpd/ports.conf.tmpl")
        subprocess.call(["/usr/local/bin/restarthttpd", "--force"])
    else:
        log("Restoring Apache default configuration")
        shutil.copyfile("/etc/nxfilter/tmpl/default-httpd.conf.tmpl", "/etc/httpd/sites-available/000-default.conf.tmpl")
        shutil.copyfile("/etc/nxfilter/tmpl/default-ports.conf.tmpl", "/etc/httpd/ports.conf.tmpl")
        subprocess.call(["/usr/local/bin/restarthttpd", "--force"])

#########################################
# WPAD
#########################################
def set_wpad():
    if enabled == "on":
        if get_wpad_port():
            log("WPAD already configured on port 85 — no DHCP restart needed")
            return

        log("Configuring WPAD for NXFilter (port 85)")
        shutil.copyfile("/etc/nxfilter/tmpl/nx-dhcpd.conf.tmpl", "/etc/dhcpd.conf.tmpl")
        subprocess.call(["/usr/local/bin/restartdhcp", "--force"])
    else:
        log("Restoring default WPAD configuration")
        shutil.copyfile("/etc/nxfilter/tmpl/default-dhcpd.conf.tmpl", "/etc/dhcpd.conf.tmpl")
        subprocess.call(["/usr/local/bin/restartdhcp", "--force"])

#########################################
# DNS & dnsmasq
#########################################
def set_dns_redirect():
    dnsmasq_dir = "/etc/dnsmasq/dnsmasq.d"
    conf_file = os.path.join(dnsmasq_dir, "nxfilter.conf")

    if enabled == "on":
        if get_dnsmasq_listening():
            log("DNSMasq already listening on port 54 — no restart needed")
            return

        if not os.path.exists(dnsmasq_dir):
            os.makedirs(dnsmasq_dir)

        with open(conf_file, "w") as f:
            f.write("port=54\n")

        subprocess.call(["/usr/local/bin/restartdnsmasq", "--force"])
        log("Updating dnsmasq port 54!")
    else:
        if not os.path.exists(dnsmasq_dir):
            os.makedirs(dnsmasq_dir)

        if os.path.exists(conf_file):
            os.remove(conf_file)

        subprocess.call(["/usr/local/bin/restartdnsmasq", "--force"])
        log("Updating dnsmasq port default 53!")

#########################################
# INPUT Firewall
#########################################
def set_input_firewall():
    target = "/etc/firewall/inputfw/nxfilter.conf"
    if enabled == "on":
        shutil.copyfile("/etc/nxfilter/tmpl/nx-firewall.conf", target)
        subprocess.call(["/usr/local/bin/setxtaccess"])
        log("Updating Inputfw Rules!")
    else:
        if os.path.exists(target):
            os.remove(target)
        subprocess.call(["/usr/local/bin/setxtaccess"])
        log("Updating Inputfw Rules!")

#########################################
# DoH Firewall
#########################################
def run_iptables(args):
    cmd = ["iptables", "-w", "5"] + args
    with open(os.devnull, 'wb') as devnull:
        return subprocess.call(cmd, stdout=devnull, stderr=devnull)

def enable_nxfilter_doh():
    # Create chain if it does not exist
    if run_iptables(["-S", "DOH_DROP"]) != 0:
        run_iptables(["-N", "DOH_DROP"])

    # Flush rules
    run_iptables(["-F", "DOH_DROP"])

    # Ensure hook in CUSTOMFORWARD
    if run_iptables(["-C", "CUSTOMFORWARD", "-j", "DOH_DROP"]) != 0:
        run_iptables(["-I", "CUSTOMFORWARD", "-j", "DOH_DROP"])

    # Allow local NXFilter
    run_iptables(["-A", "DOH_DROP", "-d", "127.0.0.1", "-p", "udp", "--dport", "53", "-j", "RETURN"])
    run_iptables(["-A", "DOH_DROP", "-d", "127.0.0.1", "-p", "tcp", "--dport", "53", "-j", "RETURN"])

    # DoT (853)
    run_iptables(["-A", "DOH_DROP", "-p", "tcp", "--dport", "853", "-j", "NFLOG", "--nflog-prefix", "NXFILTER_DOH:REJECT "])
    run_iptables(["-A", "DOH_DROP", "-p", "tcp", "--dport", "853", "-j", "REJECT"])
    run_iptables(["-A", "DOH_DROP", "-p", "udp", "--dport", "853", "-j", "NFLOG", "--nflog-prefix", "NXFILTER_DOH:REJECT "])
    run_iptables(["-A", "DOH_DROP", "-p", "udp", "--dport", "853", "-j", "REJECT"])

    # DoH 443 by list
    if os.path.isfile(DOH_LIST):
        with open(DOH_LIST, "r") as f:
            for ip in f:
                ip = ip.strip()
                if not ip:
                    continue
                run_iptables(["-A", "DOH_DROP", "-d", ip, "-p", "tcp", "--dport", "443", "-j", "NFLOG", "--nflog-prefix", "NXFILTER_DOH:REJECT "])
                run_iptables(["-A", "DOH_DROP", "-d", ip, "-p", "tcp", "--dport", "443", "-j", "REJECT"])

    log("DoH / DoT blocking with NFLOG enabled")

def disable_nxfilter_doh():
    if run_iptables(["-C", "CUSTOMFORWARD", "-j", "DOH_DROP"]) == 0:
        run_iptables(["-D", "CUSTOMFORWARD", "-j", "DOH_DROP"])

    run_iptables(["-F", "DOH_DROP"])
    run_iptables(["-X", "DOH_DROP"])

    log("DoH / DoT blocking disabled")

#########################################
# NXFilter runtime
#########################################
def get_nxfilter_pid():
    try:
        p = subprocess.Popen(["pgrep", "-f", "nxd.Main"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, _ = p.communicate()
        pids = out.strip().splitlines()
        return pids[0] if pids else None
    except Exception:
        return None

def stop():
    if not get_nxfilter_pid():
        log("NXFilter already stopped")
        if os.path.exists(NXRUN):
            try:
                os.remove(NXRUN)
            except OSError:
                pass
        return

    # Stop NXFilter suppressing JAVA_TOOL_OPTIONS messages
    env = os.environ.copy()
    env["JAVA_TOOL_OPTIONS"] = ""
    shutdown_cmd = os.path.join(NXBIN, "shutdown.sh")
    
    p = subprocess.Popen([shutdown_cmd, "-d"], env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    out, _ = p.communicate()
    for line in out.splitlines():
        if "Picked up JAVA_TOOL_OPTIONS" not in line:
            print(line)

    time.sleep(5)

    if os.path.exists(NXRUN):
        try:
            os.remove(NXRUN)
        except OSError:
            pass

    log("NXFilter stopped")

def start():
    log("Starting NXFilter...")

    env = os.environ.copy()
    if "JAVA_TOOL_OPTIONS" in env:
        del env["JAVA_TOOL_OPTIONS"]

    startup_cmd = os.path.join(NXBIN, "startup.sh")
    p = subprocess.Popen([startup_cmd, "-d"], env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    out, _ = p.communicate()
    for line in out.splitlines():
        if "Picked up JAVA_TOOL_OPTIONS" not in line:
            print(line)

    # Wait for NXFilter to respond to ping.sh
    for _ in range(30):
        time.sleep(1)
        if nxfilter_is_up():
            break

    if not nxfilter_is_up():
        fail("NXFilter did not respond. (service not UP)")

    pid = get_nxfilter_pid()
    if not pid:
        fail("NXFilter process not found")

    with open(NXRUN, "w") as f:
        f.write("{}\n".format(pid))

    try:
        uid = pwd.getpwnam("nobody").pw_uid
        gid = grp.getgrnam("nogroup").gr_gid
        os.chown(NXRUN, uid, gid)
    except Exception:
        pass

    log("NXFilter running (healthcheck OK, PID {})".format(pid))

#########################################
# MAIN
#########################################
def main():
    if len(sys.argv) < 2:
        print("Usage: nxfilter --restart")
        sys.exit(0)

    action = sys.argv[1]

    if action == "--restart":
        if enabled == "on":
            log("NXFilter enabled")
            stop()
            time.sleep(3)
            set_apache_port()
            time.sleep(3)
            set_dns_redirect()
            time.sleep(3)
            set_wpad()
            time.sleep(3)
            set_input_firewall()

            if drop_doh == "on":
                enable_nxfilter_doh()
            else:
                disable_nxfilter_doh()

            start()

            log("NXFilter fully operational")
        else:
            log("NXFilter disabled — restoring system")
            stop()
            time.sleep(3)
            set_apache_port()
            time.sleep(3)
            set_dns_redirect()
            time.sleep(3)
            set_wpad()
            time.sleep(3)
            set_input_firewall()
            disable_nxfilter_doh()

            log("System restored to default")

    elif action == "--reload-doh":
        log("Applying DoH firewall rules only (no full restart)")

        if enabled != "on":
            disable_nxfilter_doh()
            sys.exit(0)

        if drop_doh == "on":
            enable_nxfilter_doh()
        else:
            disable_nxfilter_doh()

        log("DoH rules reload finished")

    else:
        print("Usage: nxfilter --restart")

if __name__ == "__main__":
    main()