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

###############################################################################
# Script Name : wireguard
# Description : WireGuard Management (Server / Client / Peers)
###############################################################################

import os
import sys
import re
import time
import glob
import fcntl
import socket
import subprocess

# Define umask 077 equivalent to the Bash script
os.umask(077)

SERVER_FILE = "/var/efw/wireguard/server.settings"
CLIENTS_FILE = "/var/efw/wireguard/clients.settings"
PEERS_FILE = "/var/efw/wireguard/peers.settings"
LOG = "/var/log/wireguard/wireguard.log"

RUN_DIR = "/var/run/wireguard"
LOCK_SERVER = "/run/wireguard-server.lock"
LOCK_CLIENT = "/run/wireguard-client.lock"
STATUS_DIR = "/home/httpd/cgi-bin"

def log(msg):
    """Writes log messages with timestamp."""
    ensure_dir(os.path.dirname(LOG))
    timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
    with open(LOG, "a") as f:
        f.write("[{}] {}\n".format(timestamp, msg))

def ensure_dir(directory):
    """Ensures that a directory exists."""
    if not os.path.exists(directory):
        try:
            os.makedirs(directory)
        except OSError:
            pass

def valid_iface(iface):
    """Validates the WireGuard interface name."""
    return bool(re.match(r'^wg[0-9A-Za-z_-]+$', iface))

def is_ip(host):
    """Checks if the string is a valid IPv4 address."""
    return bool(re.match(r'^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$', host))

def wait_dns(host, timeout=10):
    """Waits for DNS resolution using dig. Silent on success; fails fast on failure."""
    if is_ip(host):
        return True

    interval = 2
    waited = 0

    while waited < timeout:
        cmd = ["dig", "+time=2", "+tries=1", "+short", host]
        with open(os.devnull, 'wb') as devnull:
            res = subprocess.call(cmd, stdout=devnull, stderr=devnull)
            if res == 0:
                return True
        time.sleep(interval)
        waited += interval

    return False

def wg_stop(iface):
    """Stops/removes the network interface and cleans up PID and Perl status files."""
    if not valid_iface(iface):
        return
    
    log("Stopping WireGuard interface {}".format(iface))
    with open(os.devnull, 'wb') as devnull:
        res = subprocess.call(["wg-quick", "down", iface], stdout=devnull, stderr=devnull)
        if res != 0:
            subprocess.call(["ip", "link", "del", iface], stdout=devnull, stderr=devnull)

    remove_status(iface)

def create_status(iface, role):
    """Creates the PID file and the Perl status script."""
    pidfile = os.path.join(RUN_DIR, "{}.pid".format(iface))
    pl = os.path.join(STATUS_DIR, "status-{}.pl".format(iface))
    pid = None

    target_comm = "wg-crypt-{}".format(iface)

    # Retry up to 10 times (2 seconds total) during system startup
    for _ in range(10):
        try:
            p1 = subprocess.Popen(["ps", "-eo", "pid,args"], stdout=subprocess.PIPE)
            output, _ = p1.communicate()
            for line in output.splitlines():
                parts = line.strip().split(None, 1)
                if len(parts) == 2:
                    # Strips brackets to match kernel threads like [wg-crypt-wg0]
                    comm_clean = parts[1].strip("[]")
                    if comm_clean == target_comm or parts[1] == target_comm:
                        pid = parts[0]
                        break
        except Exception:
            pass
        
        if pid:
            break
        time.sleep(0.2)

    # Writes the PID file if found
    if pid:
        try:
            with open(pidfile, "w") as f:
                f.write("{}\n".format(pid))

            import pwd, grp
            uid = pwd.getpwnam("nobody").pw_uid
            gid = grp.getgrnam("nogroup").gr_gid
            os.chown(pidfile, uid, gid)
        except Exception as e:
            log("Error writing PID file {}: {}".format(pidfile, e))
    else:
        log("Warning: Could not find PID for process {}".format(target_comm))

    if role == "server":
        label = "WireGuard Server Interface {}".format(iface)
    else:
        label = "WireGuard Client Interface {}".format(iface)

    perl_content = """#!/usr/bin/perl
require 'header.pl';
require '/home/httpd/cgi-bin/endianinc.pl';

my $wireguard = ['wg-crypt-{}', '{}', ''];
register_status(_('{}'),$wireguard);
""".format(iface, pidfile, label)

    with open(pl, "w") as f:
        f.write(perl_content)

    os.chmod(pl, 0755)

def remove_status(iface):
    """Removes PID and Perl status files."""
    pidfile = os.path.join(RUN_DIR, "{}.pid".format(iface))
    pl = os.path.join(STATUS_DIR, "status-{}.pl".format(iface))

    for path in (pidfile, pl):
        if os.path.exists(path):
            try:
                os.remove(path)
            except OSError:
                pass

def cleanup_confs(tag, *keep):
    """Removes obsolete interface configurations."""
    for c in glob.glob("/etc/wireguard/wg*.conf"):
        if not os.path.isfile(c):
            continue

        managed = False
        with open(c, "r") as f:
            for line in f:
                if line.startswith("# Managed-by: {}".format(tag)):
                    managed = True
                    break

        if not managed:
            continue

        base = os.path.basename(c)
        iface = os.path.splitext(base)[0]

        if iface in keep:
            continue

        log("Removing {}".format(iface))
        wg_stop(iface)
        try:
            os.remove(c)
        except OSError:
            pass

def fix_key(key):
    """Ensures proper key formatting keeping '=' at the end if necessary."""
    if key.endswith("="):
        key = key[:-1]
    return key + "="

def stop_all_peers():
    """Stops all registered client interfaces."""
    if not os.path.exists(CLIENTS_FILE):
        return

    with open(CLIENTS_FILE, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue

            NAME = parts[1] if len(parts) > 1 else ""
            IFACE = parts[2] if len(parts) > 2 else ""

            if not NAME or NAME.startswith("#"):
                continue

            if IFACE:
                wg_stop(IFACE)

def stop_server():
    """Stops the configured server interface."""
    if not os.path.exists(SERVER_FILE):
        return

    iface = None
    with open(SERVER_FILE, 'r') as f:
        for line in f:
            line = line.strip()
            if line.startswith("interface="):
                iface = line.split("=", 1)[1].strip()
                break

    if iface:
        wg_stop(iface)

def wg_client():
    """Handles WireGuard clients initialization."""
    if not os.path.exists(CLIENTS_FILE):
        sys.exit(0)

    lock_file = open(LOCK_CLIENT, "w")
    try:
        fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        sys.exit(0)

    active = {}

    with open(CLIENTS_FILE, "r") as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue

            EN = parts[0]
            NAME = parts[1] if len(parts) > 1 else ""
            IFACE = parts[2] if len(parts) > 2 else ""
            PRIV = parts[3] if len(parts) > 3 else ""
            PUB = parts[4] if len(parts) > 4 else ""
            END = parts[5] if len(parts) > 5 else ""
            END_PUB = parts[6] if len(parts) > 6 else ""
            ALLOWED = parts[7] if len(parts) > 7 else ""
            ADDR = parts[8] if len(parts) > 8 else ""
            KEEP = parts[9] if len(parts) > 9 else ""

            if not NAME or NAME.startswith("#"):
                continue

            conf = "/etc/wireguard/{}.conf".format(IFACE)

            if EN == "0":
                wg_stop(IFACE)
                if os.path.exists(conf):
                    try:
                        os.remove(conf)
                    except OSError:
                        pass
                continue

            active[IFACE] = True

            host = END.split(":")[0]
            if not wait_dns(host):
                msg = "DNS resolution failure: could not resolve host '{}' (endpoint field of interface {} in clients.settings)".format(host, IFACE)
                log(msg)
                print(msg)
                continue

            conf_content = []
            conf_content.append("# Managed-by: wireguard-client")
            conf_content.append("[Interface]")
            conf_content.append("PrivateKey = {}".format(fix_key(PRIV)))
            if ADDR:
                conf_content.append("Address = {}".format(ADDR))
            conf_content.append("")
            conf_content.append("# Peer: {}".format(NAME))
            conf_content.append("[Peer]")
            conf_content.append("PublicKey = {}".format(fix_key(END_PUB)))
            if ALLOWED:
                conf_content.append("AllowedIPs = {}".format(ALLOWED))
            if KEEP:
                conf_content.append("PersistentKeepalive = {}".format(KEEP))
            conf_content.append("Endpoint = {}".format(END))

            with open(conf, "w") as cf:
                cf.write("\n".join(conf_content) + "\n")

            os.chmod(conf, 0600)
            wg_stop(IFACE)

            with open(LOG, "a") as log_file:
                subprocess.call(["wg-quick", "up", IFACE], stdout=log_file, stderr=log_file)

            create_status(IFACE, "client")

    cleanup_confs("wireguard-client", *active.keys())

    for pl in glob.glob(os.path.join(STATUS_DIR, "status-wg*.pl")):
        if not os.path.isfile(pl):
            continue

        is_server = False
        with open(pl, "r") as f:
            if any("WireGuard Server Interface" in line for line in f):
                is_server = True

        if is_server:
            continue

        base = os.path.basename(pl)
        iface = base.replace("status-", "").replace(".pl", "")

        if iface not in active:
            log("Removing status for client no longer in settings: {}".format(iface))
            remove_status(iface)

def wg_server():
    """Handles WireGuard Server initialization."""
    if not os.path.exists(SERVER_FILE):
        sys.exit(0)

    lock_file = open(LOCK_SERVER, "w")
    try:
        fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        sys.exit(0)

    C = {}
    with open(SERVER_FILE, "r") as f:
        for line in f:
            line = line.strip()
            if "=" in line:
                k, v = line.split("=", 1)
                if k:
                    C[k] = v

    IFACE = C.get("interface", "")
    if not valid_iface(IFACE):
        sys.exit(1)

    cleanup_confs("wireguard-server", IFACE)

    if C.get("status") != "on":
        wg_stop(IFACE)
        sys.exit(0)

    conf = "/etc/wireguard/{}.conf".format(IFACE)
    conf_content = []
    conf_content.append("# Managed-by: wireguard-server")
    conf_content.append("[Interface]")
    conf_content.append("PrivateKey = {}".format(fix_key(C.get("private_key", ""))))

    for key, label in [("address", "Address"), ("listen_port", "ListenPort"), 
                       ("mtu", "MTU"), ("post_up", "PostUp"), ("post_down", "PostDown")]:
        if C.get(key):
            conf_content.append("{} = {}".format(label, C[key]))

    if os.path.exists(PEERS_FILE):
        with open(PEERS_FILE, "r") as f:
            for line in f:
                parts = line.strip().split()
                if not parts:
                    continue

                EN = parts[0]
                NAME = parts[1] if len(parts) > 1 else ""
                PRIVKEY = parts[2] if len(parts) > 2 else ""
                PUBKEY = parts[3] if len(parts) > 3 else ""
                ALLOWED = parts[4] if len(parts) > 4 else ""
                ALLOC = parts[5] if len(parts) > 5 else ""
                KEEP = parts[6] if len(parts) > 6 else ""

                if EN != "1":
                    continue
                if not NAME or NAME.startswith("#"):
                    continue

                conf_content.append("")
                conf_content.append("# Peer: {}".format(NAME))
                conf_content.append("[Peer]")
                conf_content.append("PublicKey = {}".format(PUBKEY))

                allowed_final = ALLOC
                if ALLOWED != "-":
                    allowed_final += ", {}".format(ALLOWED)

                conf_content.append("AllowedIPs = {}".format(allowed_final))

                if re.match(r'^[0-9]+$', KEEP):
                    conf_content.append("PersistentKeepalive = {}".format(KEEP))

    with open(conf, "w") as cf:
        cf.write("\n".join(conf_content) + "\n")

    os.chmod(conf, 0600)
    wg_stop(IFACE)

    with open(LOG, "a") as log_file:
        subprocess.call(["wg-quick", "up", IFACE], stdout=log_file, stderr=log_file)

    create_status(IFACE, "server")

def restart_single_interface(iface):
    """Restarts only ONE interface (server or client), without affecting the others."""
    if not valid_iface(iface):
        log("Invalid interface for isolated restart: {}".format(iface))
        return

    # Check if the interface belongs to the server
    server_iface = None
    if os.path.exists(SERVER_FILE):
        with open(SERVER_FILE, 'r') as f:
            for line in f:
                line = line.strip()
                if line.startswith("interface="):
                    server_iface = line.split("=", 1)[1].strip()
                    break

    if server_iface == iface:
        log("Restarting only the server interface {}".format(iface))
        wg_server()
        return

    # Otherwise, treat it as a client interface
    if not os.path.exists(CLIENTS_FILE):
        log("Interface {} not found (no server.settings or clients.settings).".format(iface))
        return

    with open(CLIENTS_FILE, "r") as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue

            EN = parts[0]
            NAME = parts[1] if len(parts) > 1 else ""
            IFACE = parts[2] if len(parts) > 2 else ""

            if not NAME or NAME.startswith("#"):
                continue
            if IFACE != iface:
                continue

            conf = "/etc/wireguard/{}.conf".format(IFACE)

            if EN == "0":
                log("Client interface {} is disabled, stopping only.".format(IFACE))
                wg_stop(IFACE)
                if os.path.exists(conf):
                    try:
                        os.remove(conf)
                    except OSError:
                        pass
                return

            PRIV = parts[3] if len(parts) > 3 else ""
            END = parts[5] if len(parts) > 5 else ""
            END_PUB = parts[6] if len(parts) > 6 else ""
            ALLOWED = parts[7] if len(parts) > 7 else ""
            ADDR = parts[8] if len(parts) > 8 else ""
            KEEP = parts[9] if len(parts) > 9 else ""

            host = END.split(":")[0]
            if not wait_dns(host):
                msg = "DNS resolution failure: could not resolve host '{}' (endpoint field of interface {} in clients.settings)".format(host, IFACE)
                log(msg)
                print(msg)
                return

            conf_content = []
            conf_content.append("# Managed-by: wireguard-client")
            conf_content.append("[Interface]")
            conf_content.append("PrivateKey = {}".format(fix_key(PRIV)))
            if ADDR:
                conf_content.append("Address = {}".format(ADDR))
            conf_content.append("")
            conf_content.append("# Peer: {}".format(NAME))
            conf_content.append("[Peer]")
            conf_content.append("PublicKey = {}".format(fix_key(END_PUB)))
            if ALLOWED:
                conf_content.append("AllowedIPs = {}".format(ALLOWED))
            if KEEP:
                conf_content.append("PersistentKeepalive = {}".format(KEEP))
            conf_content.append("Endpoint = {}".format(END))

            with open(conf, "w") as cf:
                cf.write("\n".join(conf_content) + "\n")

            os.chmod(conf, 0600)
            wg_stop(IFACE)

            with open(LOG, "a") as log_file:
                subprocess.call(["wg-quick", "up", IFACE], stdout=log_file, stderr=log_file)

            create_status(IFACE, "client")
            log("Client interface {} restarted individually.".format(IFACE))
            return

    log("Interface {} not found in clients.settings.".format(iface))

def show_help():
    help_text = """
WireGuard OpenFW Manager CLI

Usage:
  {0} [action] [options]

Actions:
  --client            Initializes and synchronizes client connections.
  --server            Initializes and synchronizes the server service.
  --stop              Stops all connections (or a specific one if provided).
  --restart           Restarts all connections or components.
  --help, -h          Displays this help message.

Options (can be combined with --stop / --restart):
  --interface=<name>  Applies the action only to a specific interface (e.g., wg0).
  --all-peers         Applies the action to all configured clients.
  --server            Applies the action only to the WireGuard server.

Examples:
  {0} --client
  {0} --server
  {0} --stop
  {0} --stop --interface=wg0
  {0} --stop --all-peers
  {0} --restart
  {0} --restart --interface=wg0
""".format(sys.argv[0])
    print(help_text)

def main():
    if len(sys.argv) < 2:
        show_help()
        sys.exit(0)

    args = sys.argv[1:]

    # Flags to assist in parsing stop/restart options
    target_iface = None
    stop_peers = False
    stop_srv = False

    for arg in args:
        if arg.startswith("--interface="):
            target_iface = arg.split("=", 1)[1].strip()
        elif arg == "--all-peers":
            stop_peers = True
        elif arg == "--server" and "--stop" in args:
            stop_srv = True

    # Stop action
    if "--stop" in args:
        if target_iface:
            wg_stop(target_iface)
        else:
            if not stop_peers and not stop_srv:
                stop_peers = True
                stop_srv = True

            if stop_peers:
                stop_all_peers()
            if stop_srv:
                stop_server()
        sys.exit(0)

    # Restart action
    if "--restart" in args:
        if target_iface:
            restart_single_interface(target_iface)
        else:
            stop_all_peers()
            stop_server()
            time.sleep(1)
            wg_client()
            wg_server()
        sys.exit(0)

    # Standard direct calls (CGI and autostart)
    if "--client" in args and not stop_srv:
        wg_client()
    elif "--server" in args and not stop_srv:
        wg_server()
    elif "--help" in args or "-h" in args:
        show_help()
    else:
        show_help()

if __name__ == "__main__":
    main()