#!/usr/bin/python3

import argparse
import logging
import ipaddress
import time
import os
import subprocess
import json
import sys
import re

parser = argparse.ArgumentParser(description="Goliath BGP Synchroniser", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("-d", "--delta", action='store_true', help="Takes a JSON based delta and synchronises it to GoBGP.")
parser.add_argument("-f", "--full", action='store_true', help="Run a test using the specified file.")
parser.add_argument("-r", "--routes", action='store_true', help="Makes the script return a list of routes one would expect.")
parser.add_argument("-c", "--clear", action='store_true', help="Clears the current table.")

parser.add_argument("-l", "--level", type=int, help="What level information is emitted to.")

DELTA=None
FULL=False
ROUTES=False
CLEAR=False
LEVEL=1

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

    if args['delta'] is not None:
        DELTA=True
    elif args['full'] is not None:
        FULL=True
    elif args['routes'] is not None:
        ROUTES=True
    
    if args['level'] is not None:
        LEVEL = args['level']

    if args['clear'] is not None:
        CLEAR = args['clear']

logger = logging.getLogger("Gol-BGP-Syn")

def get_given_prefixes():
    input_lines = sys.stdin.readlines()

    routes = []

    for line in input_lines:
        address, mask, small, large, asn = line.strip().split(',')

        routes += [{
            'address': address,
            'mask': mask,
            'small': [] if small == 'NONE' else small.split(' '),
            'large': [] if large == 'NONE' else large.split(' '),
            'asn': asn
        }]

    return routes

def get_given_deltas():
    input_lines = sys.stdin.readlines()

    deletions = []
    insertions = []

    for line in input_lines:
        operation, address, mask, small, large, asn = line.strip().split(',')

        to_add = [{
            'address': address,
            'mask': mask,
            'small': [] if small == 'NONE' else small.split(' '),
            'large': [] if large == 'NONE' else large.split(' '),
            'asn': asn
        }]

        if operation == 'DELETE':
            deletions += to_add
        if operation == 'INSERT':
            insertions += to_add
    
    return insertions, deletions

def get_current_prefixes():
    return {}

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:
                logger.debug(f'''Smaller than /25 detected -  {x["address"]}/{x["mask"]} - Converting to 25s''')
                for subnet in prefix.subnets(new_prefix=25):
                    clean_prefixes[str(subnet)] = {
                        'prefix': subnet,
                        'small': [] if x['small'] == 'NONE' else x['small'],
                        'large': [] if x['large'] == 'NONE' else x['large'],
                        'asn': x['asn'],
                        'version': 'ipv4' if subnet.version == 4 else 'ipv6'
                    }
            else:
                clean_prefixes[str(prefix)] = {
                    'prefix': prefix,
                    'small': [] if x['small'] == 'NONE' else x['small'],
                    'large': [] if x['large'] == 'NONE' else x['large'],
                    'asn': x['asn'],
                    'version': 'ipv4' if prefix.version == 4 else 'ipv6'
                }
        
    return clean_prefixes

def check_deletable(local_list, live_list):
    deletion_list = []
    # If it is in live, but not local - It's a candidate for deletion.
    for live in live_list.items():
        if live not in local_list:
            deletion_list += [live]
    
    return deletion_list

def enact_deletable(deletable):
    
    successful = []
    failed = []

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

        if command.returncode == 0:
            successful +=  [f'''DELETE {str(delete['prefix'])}''']
        else:
            failed += [f'''DELETE {str(delete['prefix'])}''']
    
    return successful, failed

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

    successful = []
    failed = []

    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)

            if command.returncode == 0:
                successful += [f'''INSERT {str(insert['prefix'])}''']
            else:
                failed += [f'''INSERT {str(insert['prefix'])}''']
    
    return successful, failed

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'))
        
    successful = []
    failed = []

    withdrawals = []

    for d in deletable:
        found_at_all = False
        
        for i in insertable:
            if str(d['prefix']) == str(i['prefix']):
                found_at_all = True
        
        if not found_at_all:
            withdrawals += [d]

    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)

            if command.returncode == 0:
                successful +=  [f'''DELETE {str(withdraw['prefix'])}''']
            else:
                failed += [f'''DELETE {str(withdraw['prefix'])}''']
    
    return successful, failed

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 main():
    if CLEAR is True:
        command = subprocess.run([
            f'''/usr/bin/gobgp/gobgp''',
            f'--port=5005{LEVEL}',
            'global',
            'rib',
            'del',
            'all'
        ])

        if command.returncode == 0:
            print('success')
        else:
            print('failure')

    if DELTA is True:
        large_community_lookup, small_community_lookup = get_communities()

        to_insert, to_delete = get_given_deltas()

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

        successful_delete, failed_delete = enact_withdraw(insertable, deletable)
        successful_insert, failed_insert = enact_insertable(insertable, large_community_lookup, small_community_lookup)

        print(json.dumps({
            'error': failed_delete + failed_insert,
            'success': successful_delete + successful_insert
        }) ) 

    elif FULL is True:
        expected_prefix_list = sanitise_prefixes(get_given_prefixes())
        running_prefix_list = get_current_prefixes()

        deletable = check_deletable(expected_prefix_list.keys(), running_prefix_list.keys())
        insertable = check_insertable(expected_prefix_list.keys(), running_prefix_list.keys())

        #enact_deletable(deletable)
        #enact_insertable(insertable)

    elif ROUTES is True:
        print("")

main()