#!/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 BGP Synchroniser", 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.")

LEVEL=2
TEST=False

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

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

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


# 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:
        to_add = []
        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 check_insertable(local_list, live_list):
    insertion_list = []

    for local in local_list:
        if local not in live_list:
            insertion_list += [local]

    return insertion_list

def enact_insertable(insertable, large_communities, small_communities):

    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'))

    for insert in insertable:
        accurate_large = False
        accurate_small = False

        if str(insert['prefix']) in large_communities and len(set(insert['large'])) == len(set(large_communities[str(insert['prefix'])])) and set(insert['large']) == set(large_communities[str(insert['prefix'])]):
            accurate_large = True
        elif len(insert['large']) == 0 and str(insert['prefix']) not in large_communities:
            accurate_large = True

        if str(insert['prefix']) in small_communities and len(set(insert['small'])) == len(set(small_communities[str(insert['prefix'])])) and set(insert['small']) == set(small_communities[str(insert['prefix'])]):
            accurate_small = True
        elif len(insert['small']) == 0 and str(insert['prefix']) not in small_communities:
            accurate_small = True

        if str(insert['prefix']) not in searchable or (not accurate_small or not accurate_large):

            command = subprocess.run([
                f'''/usr/bin/gobgp/gobgp''',
                f'--port=5005{LEVEL}',
                'global',
                'rib',
                '-a',
                insert['version'],
                'add',
                str(insert['prefix']),
                'large-community',
                ','.join(insert['large']).strip(',')
            ] if len(insert['small']) == 0 else [
                f'''/usr/bin/gobgp/gobgp''',
                f'--port=5005{LEVEL}',
                'global',
                'rib',
                '-a',
                insert['version'],
                'add',             
                str(insert['prefix']),
                'community',
                ','.join(insert['small']).strip(','),
                'large-community',
                ','.join(insert['large']).strip(',')
            ], capture_output=True)

def enact_withdraw(insertable, deletable):

    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'))

    chunk_divisor = math.floor(len(deletable) / 4)
    chunk_calculates = [list(deletable)[i:i + chunk_divisor] for i in range(0, len(deletable), chunk_divisor)]
    chunk_withdrawals = { i: [] for i in range(0, len(chunk_calculates)) }

    # We need to thread withdrawal calculation to ensure speed
    insert_prefix_dictionary = {str(i['prefix']) for i in insertable} # Calculate this once

    def calculate_withdrawals(tid, deletes):
        chunk_withdrawals[tid] += [d for d in deletes if str(d['prefix']) not in insert_prefix_dictionary]

    threads = []

    for i in range(0, len(chunk_calculates)):
        t = Thread(target=calculate_withdrawals, args=(i, chunk_calculates[i]))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()

    withdrawals = []

    for values in chunk_withdrawals.values():
        withdrawals = withdrawals + values

    for withdraw in withdrawals:

        if str(withdraw['prefix']) in searchable:
            command = subprocess.run([
                f'''/usr/bin/gobgp/gobgp''',
                f'--port=5005{LEVEL}',
                'global',
                'rib',
                '-a',
                withdraw['version'],
                'del',
                str(withdraw['prefix'])
            ], capture_output=True)

def get_communities():

    known_prefixes = subprocess.run([
        f'''/usr/bin/gobgp/gobgp''',
        f'--port=5005{LEVEL}',
        'global',
        'rib'
    ], capture_output=True)


    large_community_lookup = {}
    small_community_lookup = {}

    large_prefix_match_re = re.compile("((?:\d+\.){3}\d+\/\d+).+: \[ ((?:(?:\d+:){2}\d+(?:, )*)*)")
    large_pairs = large_prefix_match_re.findall(known_prefixes.stdout.decode('utf-8'))

    for couple in large_pairs:
        large_community_lookup[couple[0]] = couple[1].replace(" ", '').split(',')

    small_prefix_match_re = re.compile("((?:\d+\.){3}\d+\/\d+).+Communities:((?: \d+:\d+(?:\,)*)*)")
    small_pairs = small_prefix_match_re.findall(known_prefixes.stdout.decode('utf-8'))

    for couple in small_pairs:
        small_community_lookup[couple[0]] = couple[1].replace(" ", '').split(',')

    return large_community_lookup, small_community_lookup

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()

    # 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 parental communities - ONLY for prefixes that are still announced or have been recently announced. We want to drop old communities but still maintain them within Creatia.
    parent_community_rows = threadsafe_multiquery('''
        SELECT
            subnet_id,
            string_agg(community::text, \' \') as communities
        FROM
            prefixes
        LEFT JOIN
            prefix_community ON prefix_id = prefixes.id
        LEFT 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 LIKE '7578:%%' OR community LIKE '65500:%%')
            AND community LIKE '%%99'
            AND community != '7578:99'
            AND is_large = 0
        GROUP BY
            subnet_id
    ''', community_prepared_params)

    parent_community_lookup = { x['subnet_id']: x for x in parent_community_rows }
    
    # Retrieve self communities
    self_community_rows = threadsafe_multiquery('''
        SELECT
            prefix_id,
            string_agg(community::text, \' \') as communities
        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 LIKE '7578:%%' OR community LIKE '65500:%%')
            AND community LIKE '%%99'
            AND community != '7578:99'
            AND is_large = 0
        GROUP BY
            prefix_id
    ''', community_prepared_params)

    self_community_lookup = { x['prefix_id']: x for x in self_community_rows }

    # Retrieve parental mercury - Only actually relevant IF the parent is still announced as it could've possibly shifted towards a normally announced prefix...
    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 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 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_delete = []
    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 only requires calculation if it's not a deleteable
            large_communities = []
            small_communities = []

            if level_1:
                large_communities = [ '7578:666:661' ]

            if in_mercury:
                large_communities = large_communities + [ '7578:666:664' ]
            
            if level_2:
                large_communities = large_communities + [ '7578:666:662' ]

            if prefix['id'] in parent_community_lookup:
                small_communities = parent_community_lookup[prefix['id']]['communities'].split(' ')

            if prefix['id'] in self_community_lookup:
                small_communities = small_communities + self_community_lookup[prefix['id']]['communities'].split(' ')
            
            # Remove duplicate communities
            small_communities = list(set(small_communities))

            # Insertion
            to_insert = to_insert + [{
                'address': prefix['address'],
                'mask': prefix['mask'],
                'small': small_communities,
                'large': large_communities,
                'asn': prefix['asn']
            }]
        else:
            # Deletion
            to_delete = to_delete + [{
                'address': prefix['address'],
                'mask': prefix['mask'],
                'small': [],
                'large': [],
                'asn': prefix['asn']
            }]

    # Run our mainline update
    large_community_lookup, small_community_lookup = get_communities()

    insertable = sanitise_prefixes(to_insert).values()
    deletable = sanitise_prefixes(to_delete).values()

    enact_withdraw(insertable, deletable)

    enact_insertable(insertable, large_community_lookup, small_community_lookup)

    end_time = time.time()
    print(f"Execution time: {(end_time - start_time) * 1000:.3f} ms")

main()