#!/usr/bin/env python3

import argparse
import socket
import time
import json
import select
import struct
import math
import sys
import copy

# DEBUG = True
DEBUG = False

parser = argparse.ArgumentParser(description='route packets')
parser.add_argument('asn', type=int, help="AS Number")
parser.add_argument('networks', metavar='networks',
                    type=str, nargs='+', help="networks")
args = parser.parse_args()

##########################################################################################

# Message Fields
TYPE = "type"
SRCE = "src"
DEST = "dst"
MESG = "msg"
TABL = "table"

# Message Types
DATA = "data"
DUMP = "dump"
UPDT = "update"
RVKE = "revoke"
NRTE = "no route"

# Update Message Fields
NTWK = "network"
NMSK = "netmask"
ORIG = "origin"
LPRF = "localpref"
APTH = "ASPath"
SORG = "selfOrigin"

# internal route info
CUST = "cust"
PEER = "peer"
PROV = "prov"


##########################################################################################

class Router:

    routes = None
    updates = None
    relations = None
    sockets = None

    def __init__(self, networks):
        self.routes = []
        self.updates = []
        self.revokes = []
        self.relations = {}
        self.sockets = {}
        for relationship in networks:
            network, relation = relationship.split("-")
            if DEBUG:
                print("Starting socket for", network, relation)
            self.sockets[network] = socket.socket(
                socket.AF_UNIX, socket.SOCK_SEQPACKET)
            self.sockets[network].setblocking(0)
            self.sockets[network].connect(network)
            self.relations[network] = relation
        return

    def ip_to_binary_string(self, ip):
        return "".join([bin(int(x)+256)[3:] for x in ip.split('.')])

    def get_match_len(self, mask_len, ip1_in_binary, ip2_in_binary):
        match_len = 0
        for i in range(mask_len):
            if ip1_in_binary[i] == ip2_in_binary[i]:
                match_len += 1
            else:
                return 0
        return match_len

    def get_mask_length(self, mask):
        mask_len = 0
        for c in self.ip_to_binary_string(mask):
            if c == "1":
                mask_len += 1
        return mask_len

    def lookup_routes(self, daddr):
        """ Lookup all valid routes for an address """
        outroutes = []
        longest_match_len = 0
        daddr_in_binary = self.ip_to_binary_string(daddr)
        for route in self.routes:
            mask_len = self.get_mask_length(route["netmask"])
            match_len = self.get_match_len(
                mask_len, daddr_in_binary, self.ip_to_binary_string(route["network"]))
            if match_len > longest_match_len:
                outroutes.clear()
                longest_match_len = match_len
                outroutes.append(route)
            elif match_len == longest_match_len and not match_len == 0:
                outroutes.append(route)
        return outroutes

    def get_shortest_as_path(self, routes):
        """ select the route with the shortest AS Path """
        outroutes = []
        min_as_path = sys.maxsize
        for route in routes:
            if len(route["ASPath"]) < min_as_path:
                outroutes.clear()
                min_as_path = len(route["ASPath"])
                outroutes.append(route)
            elif len(route["ASPath"]) == min_as_path:
                outroutes.append(route)
        return outroutes

    def get_highest_preference(self, routes):
        """ select the route with the highest localpref """
        outroutes = []
        highest_pref = 0
        for route in routes:
            if route["localpref"] > highest_pref:
                outroutes.clear()
                highest_pref = route["localpref"]
                outroutes.append(route)
            elif route["localpref"] == highest_pref:
                outroutes.append(route)
        return outroutes

    def get_self_origin(self, routes):
        """ select self originating routes """
        outroutes = []
        for route in routes:
            if route["selfOrigin"]:
                outroutes.append(route)
        if outroutes:
            return outroutes
        else:
            return routes

    def get_origin_routes(self, routes):
        """ select origin routes: IGP > EGP > UNK """
        EGP_routes = []
        IGP_routes = []
        UNK_routes = []
        for route in routes:
            route_origin = route["origin"]
            if route_origin == "IGP":
                IGP_routes.append(route)
            elif route_origin == "EGP":
                EGP_routes.append(route)
            elif route_origin == "UNK":
                UNK_routes.append(route)
        if IGP_routes:
            return IGP_routes
        elif EGP_routes:
            return EGP_routes
        elif UNK_routes:
            return UNK_routes

    def get_lowest_ip(self, routes):
        outroutes = []
        lowest_ip_num = sys.maxsize
        lowest_ip = ""

        for route in routes:
            route_ip_num = self.ip_to_binary_string(route["peer"])
            if int(route_ip_num, 2) < lowest_ip_num:
                outroutes.clear()
                lowest_ip_num = int(route_ip_num, 2)
                lowest_ip = route["peer"]
                outroutes.append(route)
            elif int(route_ip_num, 2) == lowest_ip_num:
                outroutes.append(route)
        return outroutes

    def filter_relationships(self, srcif, routes):
        """ Don't allow Peer->Peer, Peer->Prov, or Prov->Peer forwards """
        outroutes = []
        for route in routes:
            if self.relations[srcif] == "cust" or self.relations[route["peer"]] == "cust":
                outroutes.append(route)
        return outroutes

    def get_route(self, srcif, daddr):
        """	Select the best route for a given address	"""
        peer = None
        routes = self.lookup_routes(daddr)

        # Rules go here
        if routes:
            # 1. Highest Preference
            if len(routes) > 1:
                routes = self.get_highest_preference(routes)
            # 2. Self Origin
            if len(routes) > 1:
                routes = self.get_self_origin(routes)
            # 3. Shortest ASPath
            if len(routes) > 1:
                routes = self.get_shortest_as_path(routes)
            # 4. EGP > IGP > UNK
            if len(routes) > 1:
                routes = self.get_origin_routes(routes)
            # 5. Lowest IP Address
            if len(routes) > 1:
                routes = self.get_lowest_ip(routes)

        # Final check: enforce peering relationships
        routes = self.filter_relationships(srcif, routes)
        if routes:
            peer = routes[0]["peer"]
        return self.sockets[peer] if peer else None

    def forward(self, srcif, packet):
        """	Forward a data packet	"""
        daddr = packet["dst"]
        route_socket = self.get_route(srcif, daddr)
        if route_socket:
            route_socket.sendall(json.dumps(packet).encode())
        else:
            self.send_no_route_message(srcif, packet)

        # else:
        #     no_route_packet = {"src": srcif, "dst": packet["src"],
        #                        "type": "no route", "msg": {}}
        #     self.sockets[srcif].sendall(
        #         json.dumps(no_route_packet).encode())

        # {'type': 'data', 'src': '172.168.0.25', 'dst': '192.168.0.25'}
        # dst_in_routes = False
        # for route in self.routes:
        #  if packet["dst"] in route.values():

        # dst_in_routes = True
        # if not dst_in_routes:
        #     self.send_no_route_message(srcif, packet)
        # else:

    def send_no_route_message(self, srcif, packet):
        src_address = srcif[:-1] + "1"
        no_route_packet = {"src": src_address, "dst": packet["src"],
                           "type": "no route", "msg": {}}
        self.sockets[srcif].sendall(json.dumps(no_route_packet).encode())

    def can_aggregate(self, new_route, route):
        mask_len = self.get_mask_length(route["netmask"])
        new_route_binary = self.ip_to_binary_string(new_route["network"])
        route_binary = self.ip_to_binary_string(route["network"])

        return new_route["netmask"] == route["netmask"] and \
            self.get_match_len(mask_len - 1, new_route_binary, route_binary) == mask_len - 1 and \
            new_route["peer"] == route["peer"] and \
            new_route["localpref"] == route["localpref"] and \
            new_route["selfOrigin"] == route["selfOrigin"] and \
            new_route["ASPath"] == route["ASPath"] and \
            new_route["origin"] == route["origin"]

    def update_netmask(self, ip, to_add):
        ip_in_binary = self.ip_to_binary_string(ip)
        updated_binary = list(ip_in_binary)
        for i in range(len(updated_binary) - 1, -1, -1):
            if updated_binary[i] == "1":
                if to_add:
                    updated_binary[i + 1] = "1"
                    break
                else:
                    updated_binary[i] = "0"
                    break
        updated_binary = "".join(updated_binary)
        return ".".join([str(int(updated_binary[i:i + 8], 2))
                         for i in range(0, len(updated_binary), 8)])

    def coalesce(self):
        """	coalesce any routes that are right next to each other	"""
        new_route = self.routes[len(self.routes) - 1]
        for route in self.routes:
            if new_route is not route and self.can_aggregate(new_route, route):
                new_route["netmask"] = self.update_netmask(
                    new_route["netmask"], False)
                new_route["network"] = route["network"]
                self.routes.remove(route)
                self.coalesce()

    def update(self, srcif, packet):
        """	handle update packets	"""
        # save a copy of the announcement in case you need it later
        self.updates.append(packet)
        # add an entry to your forwarding table
        message = packet["msg"]
        network = message["network"]
        netmask = message["netmask"]
        localpref = message["localpref"]
        selfOrigin = message["selfOrigin"]
        as_path = message["ASPath"]
        origin = message["origin"]
        self.routes.append(
            {"network": network, "netmask": netmask, "peer": srcif, "localpref": localpref,
             "selfOrigin": selfOrigin, "ASPath": as_path, "origin": origin})
        self.coalesce()
        # send copies of the announcement to neighboring routers
        for sock in self.sockets:
            if not sock == srcif:
                if self.relations[srcif] == "cust" or self.relations[sock] == "cust":
                    packet_copy = copy.deepcopy(packet)
                    packet_copy["src"] = sock[:-1] + "1"
                    packet_copy["dst"] = sock
                    if args.asn not in packet["msg"]["ASPath"]:
                        packet_copy["msg"]["ASPath"].append(args.asn)
                    self.sockets[sock].sendall(
                        json.dumps(packet_copy).encode())

    # def disaggregate(self, route_removal, route):
    #     route["netmask"] = self.update_netmask(route["netmask"], True)
    #     mask_len = self.get_mask_length(route["netmask"])
    #     route_in_binary = self.ip_to_binary_string(route["network"])
    #     route0 = list(route_in_binary)[mask_len - 1] = "0"
    #     route0 = ".".join([str(int(route0[i:i + 8], 2))
    #                        for i in range(0, len(route0), 8)])
    #     route1 = list(route_in_binary)[mask_len - 1] = "1"
    #     route1 = ".".join([str(int(route1[i:i + 8], 2))
    #                        for i in range(0, len(route1), 8)])
    #     new_route0 = copy.deepcopy(route)
    #     new_route0["network"] = route0
    #     new_route1 = copy.deepcopy(route)
    #     new_route1["network"] = route1

    #     self.revoke_route(route_removal, new_route0)
    #     self.revoke_route(route_removal, new_route1)

    # def revoke_route(self, route_removal, route):
    #     if route_removal["network"] == route["network"] and \
    #             route_removal["netmask"] == route["netmask"]:
    #         self.routes.remove(route)
    #     else:
    #         mask_len = self.get_mask_length(route["netmask"])
    #         removal_mask_len = self.get_mask_length(route_removal["netmask"])
    #         route_removal_binary = self.ip_to_binary_string(
    #             route_removal["network"])
    #         route_binary = self.ip_to_binary_string(route["network"])

    #         if self.get_match_len(mask_len, route_binary, route_removal_binary) == mask_len and \
    #                 mask_len < removal_mask_len:
    #             self.disaggregate(route_removal, route)

    def is_update_revoked(self, update):
        for revoke in self.revokes:
            if update["src"] == revoke["src"] and \
                    update["msg"]["network"] == revoke["network"]:
                return True
        return False

    # revoke packet
    def disaggregate(self, packet):
        self.routes.clear()

        for revoke_packet in packet["msg"]:
            # for update in self.updates:
            #     if update["src"] == packet["src"] and \
            #             update["msg"]["network"] == revoke_packet["network"] and \
            #             update["msg"]["netmask"] == revoke_packet["netmask"]:
            #         # self.updates.remove(update)
            #         print("hi")
            self.revokes.append(
                {"src": packet["src"], "network": revoke_packet["network"], "netmask": revoke_packet["network"]})

        for update in self.updates:
            # revoke is a list of {"network": "<network prefix>", "netmask": "<associated CIDR netmask>"}
            if update["type"] == UPDT and not self.is_update_revoked(update):
                message = update["msg"]
                network = message["network"]
                netmask = message["netmask"]
                localpref = message["localpref"]
                selfOrigin = message["selfOrigin"]
                as_path = message["ASPath"]
                origin = message["origin"]
                src = update["src"]
                self.routes.append(
                    {"network": network, "netmask": netmask, "peer": src, "localpref": localpref,
                        "selfOrigin": selfOrigin, "ASPath": as_path, "origin": origin})
                self.coalesce()

    def revoke(self, packet):
        """	handle revoke packets	"""
        # save a copy of the announcement in case you need it later
        self.updates.append(packet)
        # remove the networks from the list of routes
        # for route_removal in packet["msg"]:
        #     for route in self.routes:
        #         if packet["src"] == route["src"]:
        #             self.revoke_route(route_removal, route)
        self.disaggregate(packet)
        # send copies of the announcement to neighboring routers
        for sock in self.sockets:
            if not sock == packet["src"]:
                if self.relations[packet["src"]] == "cust" or self.relations[sock] == "cust":
                    packet_copy = copy.deepcopy(packet)
                    packet_copy["src"] = sock[:-1] + "1"
                    packet_copy["dst"] = sock
                    self.sockets[sock].sendall(
                        json.dumps(packet_copy).encode())

    def dump(self, packet):
        """	handles dump table requests	"""
        routing_table_formatted = []
        for route in self.routes:
            routing_table_formatted.append({"network": route["network"],
                                            "netmask": route["netmask"], "peer": route["peer"]})
        message = {"src": packet["dst"], "dst": packet["src"],
                   "type": "table", "msg": routing_table_formatted}
        self.sockets[message["dst"]].sendall(json.dumps(message).encode())

# [{'network': '192.168.2.0', 'netmask': '255.255.255.0', 'peer': '192.168.0.2'},
# {'network': '192.168.3.0', 'netmask': '255.255.255.0', 'peer': '192.168.0.2'},
# {'network': '172.169.0.0', 'netmask': '255.255.0.0', 'peer': '172.168.0.2'}]

    def handle_packet(self, srcif, packet):
        """	dispatches a packet """
        #! TODO
        # DATA = "data"
        # DUMP = "dump"
        # UPDT = "update"
        # RVKE = "revoke"
        # NRTE = "no route"
        message_type = packet["type"]
        if message_type == UPDT:
            self.update(srcif, packet)
        if message_type == DATA:
            self.forward(srcif, packet)
        if message_type == DUMP:
            self.dump(packet)
        if message_type == RVKE:
            self.revoke(packet)
        return False

    def send_error(self, conn, msg):
        """ Send a no_route error message """
        return

    # # test
    # def ip_address_to_binary(self, ip):
    #     split_ip = ip.split(".")
    #     eight_bit_nums = []
    #     for i in split_ip:
    #         i_in_binary = format(int(i), '08b')
    #         eight_bit_nums.append(i_in_binary.split("b")[0])
    #     binary_ip_str = ""
    #     for i in eight_bit_nums:
    #         binary_ip_str += i
    #     return binary_ip_str

    # # def binary_to_ip(self, binary):

    # def bitwise_addition(self, b1, b2):
    #     for i in range(len(b1)):
    #         result = ""
    #         result += str(int(b1[i]) and int(b2[i]))
    #     return result

    def run(self):
        while True:
            socks = select.select(self.sockets.values(), [], [], 0.1)[0]
            for conn in socks:
                try:
                    k = conn.recv(65535)
                except:
                    # either died on a connection reset, or was SIGTERM's by parent
                    return
                if k:
                    for sock in self.sockets:
                        if self.sockets[sock] == conn:
                            srcif = sock
                    msg = json.loads(k)
                    if not self.handle_packet(srcif, msg):
                        self.send_error(conn, msg)
                else:
                    return
        return


if __name__ == "__main__":
    router = Router(args.networks)
    router.run()
