from flask import Flask, request, jsonify
import re
import subprocess
import argparse
import logging
import ipaddress
import time
import os
import subprocess
import json
import sys
import re
import jsonrpclib
import ssl
from threading import Thread

# This allows us to send things to the Arista's
try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    # Legacy Python that doesn't verify HTTPS certificates by default
    pass
else:
    # Handle target environment that doesn't support HTTPS verification
    ssl._create_default_https_context = _create_unverified_https_context

# Make flask application
app = Flask(__name__)


def load_configuration():
    CONFIG = {
        'endpoints': [
            { "user": "", "pw": "", "hostname": "" }
        ]
    }

    if os.path.exists('/etc/acl-forward/config.json'):
        CONFIG = json.load(open('/etc/acl-forward/config.json'))
    else:    
        os.makedirs('/etc/acl-forward/', exist_ok=True)
        
        with open('/etc/acl-forward/config.json', 'w') as new_config_file:
            try:
                new_config_file.write(json.dumps(CONFIG, indent=2))
            except IOError as e:
                print("Failed to create blank configuration file")

    return CONFIG

# Process ACL param
def process_input(input_lines):

    hosts = []
    excludes = []
    includes = []

    for line in input_lines:
        op, address, mask = line.strip().split(',')

        if op == 'ADD':
            includes = includes + [ ipaddress.IPv4Network(f'''{address}/{mask}''') ]
        else:
            excludes = excludes + [ ipaddress.IPv4Network(f'''{address}/{mask}''') ]
    
    return list(dict.fromkeys(ipaddress.collapse_addresses(includes))), list(dict.fromkeys(ipaddress.collapse_addresses(excludes)))

# Process adds and subs
def process_diff(includes, excludes):
    # We need to actually keep collapsing until there is no changes, which is... gross
    diffed = []
    diffed_last = 0

    if len(excludes) == 0:
        return includes
    
    # Do the first pass
    for include in includes:
        for exclude in excludes:
            if exclude.subnet_of(include):
                diffed = diffed + list(include.address_exclude(exclude))
            else:
                diffed = diffed + [ include ]

    while len(diffed) != diffed_last:
        for include in includes:
            for exclude in excludes:
                if exclude.subnet_of(include):
                    diffed = diffed + list(include.address_exclude(exclude))
                else:
                    diffed = diffed + [ include ]
        
        diffed_last = len(diffed)

    return list(dict.fromkeys(ipaddress.collapse_addresses(diffed)))

# Build command to send to Aristas
def build_acl_command(prefixes):
    cmd = [
        "enable",
        "configure session",
        "no ip access-list creatia-acl",
        "ip access-list creatia-acl",
        "counters per-entry"
    ]

    for prefix in prefixes:
        cmd = cmd + [ f"permit ip any {prefix.with_prefixlen}" ]

    return cmd + ["commit"]

# Send command off to Aristas
def execute_command(CONFIG, command):
    
    # Build up our thread-pool and reserve outcomes
    thread_outcomes = {}
    threads = []

    # Define individual thread
    def command_thread(command, arista):
        try:            
            url = "{protocol}://{user}:{pw}@{hostname}/command-api".format(protocol="https", user=arista["user"], pw=arista["pw"], hostname=arista["hostname"])
            
            server = jsonrpclib.Server(url)
        
            server.runCmds(1, command)

            print(f"Sent command to {arista['hostname']}")

            thread_outcomes[arista["hostname"]] = True
        except:
            print(f"Failed to send command to {arista['hostname']}")
            
            thread_outcomes[arista["hostname"]] = False

    # Start a thread for each arista
    for arista in CONFIG['endpoints']:
        thread = Thread(target=command_thread, args=[command, arista])
        
        thread.start()

        threads.append(thread)

    # Wait for threads to finish
    for thread in threads:
        thread.join()

    return thread_outcomes

# Route that takes external input
@app.route('/acl', methods=['POST'])
def acl():
    CONFIG = load_configuration()

    # Get data from request
    data = request.get_json()
    acl = data.get('acl')

    # Validate request payload
    if acl is None or len(acl) == 0:
        return jsonify({"error": "Invalid request"}), 400
    else:
        includes, excludes = process_input(acl)

        prefixes = process_diff(includes, excludes)

        command = build_acl_command(prefixes)

        sys.stdout.flush()
        
        return jsonify({"result": execute_command(CONFIG, command)}), 200

if __name__ == '__main__':
    app.run(debug=False)