import shutil
import subprocess
import signal
import os
import psycopg2
import datetime

# PostgreSQL database connection details (assuming you have psql.creds file)
with open('psql.creds', 'r') as f:
    config = {}
    for line in f:
        key, value = line.strip().split(' = ')
        config[key] = value

def update_and_reload():
    """
    Retrieves VPN list from the database, restarts rbldnsd, and logs the process.
    Also reports the number of VPN servers with exit IPs updated in the last 7 days.
    """

    # Connect to PostgreSQL database
    conn = psycopg2.connect(
        dbname=config['db_name'],
        user=config['db_user'],
        password=config['db_password'],
        host=config['db_host'],
        port=config['db_port']
    )
    cursor = conn.cursor()

    # Fetch VPN servers with exit IPs updated in the last 7 days
    seven_days_ago = datetime.datetime.now() - datetime.timedelta(days=7)
    cursor.execute("""
        SELECT COUNT(*) FROM ovpn_files 
        WHERE last_exit_ip_check >= %s 
        AND exit_ip IS NOT NULL 
        AND exit_ip ~ '^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$'
    """, (seven_days_ago,))
    updated_servers_count = cursor.fetchone()[0]
    print(f"VPN servers with exit IPs updated in the last 7 days: {updated_servers_count}")

    # Fetch the list of VPN servers from the database (filter for valid IPs and not NULL)
    cursor.execute("""
        SELECT exit_ip, file_name FROM ovpn_files 
        WHERE exit_ip IS NOT NULL 
        AND exit_ip ~ '^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$'
    """)
    vpn_list = cursor.fetchall()

    # Write the VPN list to /etc/vpnlist.txt
    try:
        with open('/etc/vpnlist.txt', 'w') as f:
            for ip, filename in vpn_list:
                f.write(f"{ip} # {filename}\n")
        print("VPN list written to /etc/vpnlist.txt successfully.")
    except IOError as e:
        print(f"Error writing VPN list to /etc/vpnlist.txt: {e}")
        cursor.close()
        conn.close()
        return  # Exit if write fails

    # Find the PID of rbldnsd using ps aux
    try:
        output = subprocess.check_output(["ps", "aux"]).decode()
        for line in output.splitlines():
            if "rbldnsd" in line:
                pid = int(line.split()[1])
                break
        else:
            raise ValueError("rbldnsd process not found in ps aux output.")
        print(f"Found rbldnsd process with PID: {pid}")
    except (subprocess.CalledProcessError, ValueError) as e:
        print(f"Error finding rbldnsd PID: {e}")
        cursor.close()
        conn.close()
        return

    # Stop rbldnsd
    try:
        os.kill(pid, signal.SIGTERM)  # Send a termination signal
        print("Sent termination signal to rbldnsd.")
    except ProcessLookupError:
        print("Error: rbldnsd process not found (might have already stopped).")
        cursor.close()
        conn.close()
        return

    # Start rbldnsd (make sure this matches your actual rbldnsd command)
    try:
        subprocess.run(["rbldnsd", "-4", "-b", "0.0.0.0/5553", "-l", "/var/log/rbldnsd.log",
                        "dnsbl.rizon.net:ip4set:/etc/rbldnsd.db",
                        "dnsbl.rizon.net:ip4set:/etc/vpnlist.txt"], check=True)
        print("rbldnsd restarted successfully.")
    except subprocess.CalledProcessError as e:
        print(f"Error restarting rbldnsd: {e}")

    cursor.close()
    conn.close()

    print("VPN list updated and rbldnsd restarted successfully.")

if __name__ == "__main__":
    update_and_reload()