#!/usr/bin/python3

import argparse
import logging
import ipaddress
from dotenv import load_dotenv
import time
import os
import subprocess
import json
import sys
import re
import math
import psycopg2
from psycopg2 import Error
from psycopg2.extras import DictCursor
import threading
from threading import Thread

parser = argparse.ArgumentParser(description="Goliath RIB Prefix Validator", formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument("-l", "--level", type=int, help="What level information is emitted to.")
parser.add_argument("-t", "--test", action='store_true', help="Whether or not this is a unit test run.")
parser.add_argument("-a", "--action", action='store_true', help="Generate deletion commands instead of just listing prefixes.")

LEVEL=2
TEST=False
ACTION=False

if len(vars(parser.parse_args())) != 0:
    args = vars(parser.parse_args())

    if args['test'] is True:
        TEST=True
    
    if args['action'] is True:
        ACTION=True
    
    if args['level'] is not None:
        LEVEL = args['level']

logger = logging.getLogger("Goliath-RIB-Validator")

# Load in the appropriate environment and then develop our DB connection details
if TEST:
    load_dotenv(dotenv_path='/var/www/html/creatia-refactor/.env.testing')
else: 
    if os.path.exists('/var/www/html/creatia-refactor/.env'):
        load_dotenv(dotenv_path='/var/www/html/creatia-refactor/.env')
    else:
        load_dotenv(dotenv_path='/var/www/html/creatia/.env')

db_host=os.getenv("DB_HOST")
db_database=os.getenv("DB_DATABASE")
db_user=os.getenv("DB_USERNAME")
db_password=os.getenv("DB_PASSWORD")
db_port=os.getenv("DB_PORT")

def sanitise_prefixes(prefixes):
    clean_prefixes = {}

    for x in prefixes:
        prefix = ipaddress.IPv4Network(f'''{x["address"]}/{x["mask"]}''')

        if prefix.is_global:
            if int(x['mask']) < 25:
                for subnet in prefix.subnets(new_prefix=25):
                    clean_prefixes[str(subnet)] = {
                        'prefix': subnet,
                        'small': x['small'],
                        'large': x['large'],
                        'asn': x['asn'],
                        'version': 'ipv4'
                    }
            else:
                clean_prefixes[str(prefix)] = {
                    'prefix': prefix,
                    'small': x['small'],
                    'large': x['large'],
                    'asn': x['asn'],
                    'version': 'ipv4'
                }
        
    return clean_prefixes

def get_rib_prefixes():
    """Get all prefixes currently in the RIB"""
    known_prefixes = subprocess.run([
        f'''/usr/bin/gobgp/gobgp''',
        f'--port=5005{LEVEL}',
        'global',
        'rib'
    ], capture_output=True)

    prefix_match_re = re.compile("\*> ((?:\d+[\.\/])+\d+)")
    searchable = prefix_match_re.findall(known_prefixes.stdout.decode('utf-8'))
    
    return set(searchable)

def query_database(query, params=None):
    try:
        # Database connection parameters
        connection = psycopg2.connect(
            host=db_host,
            database=db_database,
            user=db_user,
            password=db_password,
            port=db_port
        )

        connection.autocommit = True
        
        # Create a cursor that returns dictionaries
        cursor = connection.cursor(cursor_factory=DictCursor)
        
        # Execute the query
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)
            
        # Fetch all rows and convert to dictionaries
        records = [dict(row) for row in cursor.fetchall()]
        
        return records
        
    except (Exception, Error) as error:
        print("Error while connecting to PostgreSQL:", error)
        return []
        
    finally:
        if connection:
            cursor.close()
            connection.close()

def threadsafe_multiquery(query, params):
    shared_dictionary = {}

    def dict_write(tid, value):
        shared_dictionary[tid] = value
    
    def query_thread(tid, q, p):
        dict_write(tid, query_database(q, p))

    threads = []

    for i, param_set in enumerate(params):
        t = Thread(target=query_thread, args=(i, query, param_set))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()

    joined_results = []

    for values in shared_dictionary.values():
        joined_results = joined_results + values

    return joined_results

def main():
    start_time = time.time()

    # Get current RIB prefixes
    current_rib_prefixes = get_rib_prefixes()
    print(f"Found {len(current_rib_prefixes)} prefixes currently in RIB")

    # Retrieve information - Segregate it by a modulo of the prefix ID
    prefix_rows = threadsafe_multiquery('''
        SELECT
            p1.id,
            still_announced,
            flag_synproxy,
            forced,
            address,
            mask,
            asn,
            (select count(*) from supernets left join prefixes p2 on p2.id = supernet_id where subnet_id = p1.id and p2.still_announced = 1) as in_supernet
        FROM
            prefixes AS p1
        INNER JOIN
            asns ON p1.asn_id = asns.id
        LEFT JOIN
            prefix_goliath_inline ON p1.id = prefix_goliath_inline.prefix_id
        LEFT JOIN
            prefix_goliath_flags ON p1.id = prefix_goliath_flags.prefix_id
        WHERE
            MOD(p1.id, %s) = %s
            AND family(address) = 4
            AND mask >= 18
            AND mask <= 32
            AND asns.approved = 1
            AND (setup_required IS NULL OR setup_required NOT IN (1, 2))
            AND (p1.deleted_at IS NULL OR p1.deleted_at > (now() - INTERVAL '10 MINUTES')) 
    ''', [
        (4, 0),
        (4, 1),
        (4, 2),
        (4, 3)
    ])

    prefix_rows.sort(key=lambda x: x['mask'])

    # Only calculate a chunked version of the prefix array once for performance
    chunked_prefix_rows = [prefix_rows[i:i + 3000] for i in range(0, len(prefix_rows), 3000)]
    chunked_prefix_count = len(chunked_prefix_rows)

    # Prepared params for communities
    community_prepared_params = []

    for i in chunked_prefix_rows:
        community_prepared_params += [ ( ([j['id'] for j in i],) ) ]    

    # Retrieve self mercury
    self_mercury_lookup = threadsafe_multiquery('''
        SELECT
            prefix_id
        FROM
            prefixes
        INNER JOIN
            prefix_community ON prefix_id = prefixes.id
        WHERE
            prefixes.deleted_at IS NULL
            AND prefix_community.deleted_at IS NULL
            AND prefix_id = ANY(%s)
            AND community = '7578:91'
            AND is_large = 0
    ''', community_prepared_params)

    self_mercury_lookup = { x['prefix_id']: x for x in self_mercury_lookup }

    # Retrieve parental mercury
    parent_mercury_lookup = threadsafe_multiquery('''
        SELECT
            subnet_id
        FROM
            prefixes
        INNER JOIN
            prefix_community ON prefix_id = prefixes.id
        INNER JOIN
            supernets ON supernet_id = prefixes.id
        WHERE
            prefixes.deleted_at IS NULL
            AND (still_announced = 1 OR prefixes.seen_at > (NOW() - INTERVAL '1 DAY'))
            AND prefix_community.deleted_at IS NULL
            AND supernets.subnet_id = ANY(%s)
            AND community = '7578:91'
            AND is_large = 0
    ''', community_prepared_params)

    parent_mercury_lookup = { x['subnet_id']: x for x in parent_mercury_lookup }

    # Retrieve firewall and parental firewalls
    parent_firewall_lookup = threadsafe_multiquery('''
        SELECT
            subnet_id,
            ruleset_id
        FROM
            prefix_goliath_firewall
        INNER JOIN
            supernets ON supernet_id = prefix_goliath_firewall.prefix_id
        INNER JOIN
            prefixes ON prefixes.id = supernet_id
        WHERE
            prefix_id = ANY(%s) AND prefix_goliath_firewall.deleted_at IS NULL
    ''', community_prepared_params)

    parent_firewall_lookup = { x['subnet_id']: x for x in parent_firewall_lookup }

    self_firewall_lookup = threadsafe_multiquery('''
        SELECT
            prefix_id,
            ruleset_id
        FROM
            prefix_goliath_firewall
        WHERE
            prefix_id = ANY(%s) AND prefix_goliath_firewall.deleted_at IS NULL
    ''', community_prepared_params)

    self_firewall_lookup = { x['prefix_id']: x for x in self_firewall_lookup }

    # Format this selected information for usage 
    to_insert = []

    for prefix in prefix_rows:
        in_mercury = (prefix['still_announced'] or prefix['in_supernet'] > 0) and ((prefix['id'] in self_mercury_lookup) or (prefix['id'] in parent_mercury_lookup)) 
        valid_announcement = (prefix['still_announced'] == 1 or prefix['in_supernet'] > 0)
        
        level_1 = (prefix['flag_synproxy'] == 1 or prefix['forced'] == 1) and valid_announcement
        level_2 = ((prefix['id'] in self_firewall_lookup and self_firewall_lookup[prefix['id']]['ruleset_id'] is not None) or (prefix['id'] in parent_firewall_lookup and parent_firewall_lookup[prefix['id']]['ruleset_id'] is not None)) and valid_announcement

        if level_1 or level_2:
            # This prefix should be in the RIB
            to_insert = to_insert + [{
                'address': prefix['address'],
                'mask': prefix['mask'],
                'small': [],  # Not needed for this comparison
                'large': [],  # Not needed for this comparison
                'asn': prefix['asn']
            }]

    # Get all prefixes that should be in the RIB (after /25 splitting)
    should_be_in_rib = sanitise_prefixes(to_insert)
    should_be_in_rib_set = set(str(v['prefix']) for v in should_be_in_rib.values())
    
    print(f"Found {len(should_be_in_rib_set)} prefixes that should be in RIB")

    # Find prefixes that are in RIB but shouldn't be
    prefixes_to_remove = current_rib_prefixes - should_be_in_rib_set
    
    if ACTION:
        print(f"# Found {len(prefixes_to_remove)} prefixes in RIB that should NOT be there")
        print(f"# Generating deletion commands...")
        print("#" + "-" * 50)
        
        # Sort the prefixes for better readability
        sorted_prefixes = sorted(prefixes_to_remove, key=lambda x: ipaddress.IPv4Network(x))
        
        for prefix in sorted_prefixes:
            print(f"/usr/bin/gobgp/gobgp --port=5005{LEVEL} global rib del {prefix}")
    else:
        print(f"\nFound {len(prefixes_to_remove)} prefixes in RIB that should NOT be there:")
        print("-" * 50)
        
        # Sort the prefixes for better readability
        sorted_prefixes = sorted(prefixes_to_remove, key=lambda x: ipaddress.IPv4Network(x))
        
        for prefix in sorted_prefixes:
            print(prefix)

    end_time = time.time()
    print(f"\n# Execution time: {(end_time - start_time) * 1000:.3f} ms" if ACTION else f"\nExecution time: {(end_time - start_time) * 1000:.3f} ms")

if __name__ == "__main__":
    main()