#!/usr/bin/env python3

import argparse, socket, json, select, math, operator

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

##########################################################################################
# received incoming data
messages = []
forwarding_table = []


##########################################################################################
######### UTILITY FUNCTIONS #########
# convert ip of neighbor to router (*.*.*.2 to *.*.*.1)
def neighborToRouterIP(neighbor):
    return '.'.join(neighbor.split('.')[:-1] + ["1"])


# convert ip of router to neighbor (*.*.*.1 to *.*.*.2)
def routerToNeighborIP(router):
    return '.'.join(router.split('.')[:-1] + ["2"])


# converts an ip to its value by int
def ipToValue(ip):
    return int(ip.replace('.', ''))


# converts a netmask number to its IP representation (Ex: /20 --> 255.255.240.0)
def netmaskNumToIP(num):
    full = math.floor(num / 8)
    part = num % 8
    nmIP = ""
    partIPstr = ""
    remainder = math.ceil(full - 1)

    for f in range(full):
        nmIP += "255."

    for p in range(part):
        partIPstr += "1"

    partIPint = int(partIPstr.ljust(8, '0'), 2)
    partIP = str(partIPint)
    nmIP += partIP

    if remainder < 4:
        for r in range(remainder):
            nmIP += ".0"

    return nmIP


# convert ip address to binary representation of four 8-bit segments
def binl(ip):
    ipl = ip.split(".")
    binl = []
    for n in ipl:
        # turn into binary
        binip = bin(int(n))[2:]
        # add padding for 8-bit segments
        strip = str(binip).zfill(8)
        binl.append(strip)
    return binl


# calculate the combination for network/netmask addresses
def combine_submask(netmask_ip, dest_ip):
    # print("OG Netmask: ", netmask_ip)
    # print("OG Dest: ", dest_ip)
    combined = []
    # separate bits into four groups of 8-bits
    for b in range(4):
        # calculate bitwise and for each bit individually
        finpart = ""
        for s in range(8):
            finpart += str(int(netmask_ip[b][s]) & int(dest_ip[b][s]))
        # print("Part: ", finpart)
        combined.append(finpart)
    # print("Combined: ", combined)
    return combined


# match combined network/netmask addresses with corresponding neighbor network
def mask(dest):
    result = []
    for e in forwarding_table:
        nmbinl = binl(e['netmask'])
        destbinl = binl(dest)
        subresult = combine_submask(nmbinl, destbinl)
        result.append(subresult)
    return result


# convert list of binary ips(in list form) to list of dec ips (in list form)
def ipl(binl):
    ipl = []
    for b in range(len(binl)):
        binp = []
        for p in range(4):
            binp.append(int(binl[b][p], 2))
        ipl.append(binp)
    return ipl


# see if a specific network matches a
# specific destination that has already been masked
def matchIps(networkipl, destipl):
    for i in range(4):
        if networkipl[i] == 0:
            return True
        elif networkipl[i] != destipl[i]:
            return False


# encodes dictionary as bytes in json format
def dictToByte(dictionary):
    dictj = json.dumps(dictionary)
    dictb = dictj.encode("utf-8")
    return dictb


# converts ip address string to integer list of size four
def ipToIntList(strl):
    return list(map(int, strl.split('.')))


# returns the length of the netmask
def countNetmask(nmstr):
    nmBin = binl(nmstr)
    count = 0
    for x in range(len(nmBin)):
        for y in nmBin[x]:
            if int(y) == 1:
                count += 1
    return count


class Router:
    routes = None
    updates = None
    relations = None
    sockets = None

    def __init__(self, networks):
        self.routes = []
        self.updates = []
        # {network IP of neighbor : the kind of neighbor they are (peer, prov, cust)}
        self.relations = {}
        # {network IP of neighbor : corresponding socket}
        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 remove_entries(self, src, msgs):
        """ Finds messages to be revoked and removes from forwarding table"""
        # print("Messages to be removed: ", msgs)
        for msg in msgs:
            for entry in forwarding_table:
                if msg['network'] == entry['network'] and msg['netmask'] == entry['netmask'] and src == entry['peer']:
                    forwarding_table.remove(entry)

    # print("Removed entry: ", entry)

    def lookup_routes(self, daddr):
        """ Lookup all valid routes for an address """
        port = mask(daddr)
        resultdec = ipl(port)
        outroutes = []
        for i in range(len(port)):
            ipstr = forwarding_table[i]['network']
            if matchIps(ipToIntList(ipstr), resultdec[i]):
                outroutes.append(forwarding_table[i])
        return outroutes

    def get_shortest_as_path(self, routes):
        """ select the route with the shortest AS Path """
        outroutes = []
        # {length of AS Path : corresponding entry from forwarding table}
        pathCounts = {}
        for entry in routes:
            currPathLength = len(entry['ASPath'])
            pathCounts[currPathLength] = entry
        minPath = min(pathCounts.keys())

        for pathLength in pathCounts.keys():
            if pathLength == minPath:
                outroutes.append(pathCounts[pathLength])
        return outroutes

    def get_highest_preference(self, routes):
        """ select the route with the highest localpref """
        outroutes = []
        # {localPref : corresponding entry from forwarding table with that localPref #}
        prefRoutes = []
        for entry in routes:
            localPref = entry['localpref']
            # print("LOCALPREF ", localPref)
            prefRoutes.append(localPref)
        highestPref = max(prefRoutes)
        for entry in routes:
            localpref = entry['localpref']
            if localpref == highestPref:
                outroutes.append(entry)
        # print("HIGHEST PREF: ", outroutes)
        return outroutes

    def get_self_origin(self, routes):
        """ select self originating routes """
        outroutes = []
        for entry in routes:
            origin = entry['selfOrigin']
            if origin:
                # print("SELFORIGIN IS TRUE ", entry)
                outroutes.append(entry)
        if not outroutes:
            return routes
        else:
            return outroutes

    def get_origin_routes(self, routes):
        """ select origin routes: IGP > EGP > UNK """
        igpList = []
        egpList = []
        unkList = []
        outroutes = []
        for entry in routes:
            originRoute = entry['origin']
            if originRoute == "IGP":
                igpList.append(entry)
            elif originRoute == "EGP":
                egpList.append(entry)
            elif originRoute == "UNK":
                unkList.append(entry)

        if igpList:
            outroutes = igpList
        elif egpList:
            outroutes = egpList
        elif unkList:
            outroutes = unkList

        return outroutes

    def lowest_ip(self, routes):
        """ return the neighbor router (i.e. the "src" of the update message) with the lowest IP address """
        outroutes = []
        # {entry 'src' ip address : entry}
        iplist = {}
        # {entry 'src' in number form : entry}
        ipvlist = {}
        for entry in routes:
            iplist[entry['peer']] = entry

        for r in iplist.keys():
            # print("IP: ", r)
            ipnum = ipToValue(r)
            ipvlist[ipnum] = iplist[r]

        # print("IP Num List: ", ipvlist)
        minipv = min(ipvlist.keys())

        for m in ipvlist.keys():
            if minipv == m:
                outroutes.append(ipvlist[m])
        return outroutes

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

    def get_route(self, srcif, routes):
        """ Select the best route for a given address   """
        if routes:
            # print("Check route rules to select best route")
            # 1. Highest Preference
            routes = self.get_highest_preference(routes)
            # 2. Self Origin
            routes = self.get_self_origin(routes)
            # 3. Shortest ASPath
            routes = self.get_shortest_as_path(routes)
            # 4. EGP > IGP > UNK
            routes = self.get_origin_routes(routes)
            # 5. Lowest IP Address
            routes = self.lowest_ip(routes)
            # Final check: enforce peering relationships
            # print("Final route check")
            routes = self.filter_relationships(srcif, routes)
        # print("Best route: ", routes[0])
        # return the peer to send the msg through, else return nothing
        if len(routes) == 0:
            return None
        else:
            return routes[0]

    def longest_prefix(self, routes):
        # count netmasks of each route
        outroutes = []
        # {netmask length : list of forwarding table entries with matching length}
        nmCounts = {}
        for entry in routes:
            # print("Entry: ", entry)
            nmlen = countNetmask(entry['netmask'])
            entryList = []
            if nmlen in nmCounts.keys():
                entryList = nmCounts[nmlen]
            entryList.append(entry)
            nmCounts[nmlen] = entryList
        # find max netmask = maxNM
        maxNM = max(nmCounts.keys())
        # print("numCounts: ", nmCounts)
        # print("Max Netmask: ", maxNM)
        # return nmCounts value of list of routes with maxNM key
        for nm in nmCounts:
            if nm == maxNM:
                outroutes = nmCounts[nm]
        return outroutes

    # sends a no route msg to src of packet
    def send_no_route(self, srcif, packet):
        my_ip = neighborToRouterIP(srcif)
        datapkg = {"src": my_ip, "dst": packet['src'], "type": "no route", "msg": {}}
        datapkgb = dictToByte(datapkg)
        try:
            self.sockets[srcif].sendall(datapkgb)
        except socket.error as e:
            print("Socket failed to send: ", e)
            exit(1)

    # sends route msg to src of packet
    def send_one_route(self, packet, route):
        datapkg = {"src": packet['src'], "dst": packet['dst'], "type": "data", "msg": packet['msg']}
        datapkgb = dictToByte(datapkg)
        try:
            self.sockets[route['peer']].sendall(datapkgb)
        except socket.error as e:
            print("Socket failed to send: ", e)
            exit(1)

    # broadcasts messages to respective relation types
    def broadcast(self, announcement, packet):
        srcIP = packet['src']
        srcType = self.relations[srcIP]
        # if broadcast to all neighbors if source is customer
        if srcType == "cust":
            for n in self.sockets.keys():
                if n != srcIP:
                    # creates appropriate python address for announcement source
                    dest = neighborToRouterIP(n)
                    # sets the correct values for announcement src, dst, and ancts
                    announcement['src'] = dest
                    announcement['dst'] = n
                    # print("ANNOUNCEMENT: ", announcement)
                    ancts = json.dumps(announcement)
                    anctb = ancts.encode('utf-8')
                    self.sockets[n].sendall(anctb)
                    # print("Send message to all customers")
        else:
            for r in self.relations.keys():
                if r != srcIP and (self.relations[r] == "cust"):
                    # creates appropriate python address for announcement source
                    dest = neighborToRouterIP(r)
                    # sets the correct values for announcement src, dst, and ancts
                    announcement['src'] = dest
                    announcement['dst'] = r
                    # print("ANNOUNCEMENT: ", announcement)
                    ancts = json.dumps(announcement)
                    anctb = ancts.encode('utf-8')
                    self.sockets[r].sendall(anctb)
                    # print("Send message to all peers and providers")
        # print("Finished broadcasting")

    def forward(self, srcif, packet):
        """ Forward a data packet   """
        # lookup all valid routes for an address
        routes = self.lookup_routes(packet['dst'])
        numRoutes = len(routes)

        # Scenario 1: no route => send a "no route" msg back to the src of the data packet
        if numRoutes == 0:
            # print("Scenario 1")
            self.send_no_route(srcif, packet)
        # Scenario 2: one possible route
        if numRoutes == 1:
            # print("Scenario 2")
            route = routes[0]
            srcType = self.relations[srcif]
            dstIP = route['peer']
            dstType = self.relations[dstIP]
            if srcType != "cust" and dstType != "cust":
                self.send_no_route(srcif, packet)
            else:
                self.send_one_route(packet, route)
        elif numRoutes > 1:
            routes = self.longest_prefix(routes)
            numRoutes = len(routes)
            # Scenario 3: longest prefix
            if numRoutes == 1:
                # print("Scenario 3")
                # if longest prefix returned only one route, make that the route to take
                # print("Routes: ", routes)
                route = routes[0]
                self.send_one_route(packet, route)
            # Scenario 4: tie breakers
            elif numRoutes > 1:
                # print("Scenario 4")
                # if not, then go to get_route to apply the tie-breaker rules
                route = self.get_route(srcif, routes)
                if not route:
                    self.send_no_route(srcif, packet)
                else:
                    self.send_one_route(packet, route)

        return False

    # creates updated list of what networks this new entry will cover
    def update_covered_networks(self, covered_networks, network):
        if network['isOriginal'] != True:
            for entry in network['coveredNetworks']:
                covered_networks.append(entry)
        else:
            covered_networks.append(network)
        return covered_networks

    # checks if two entries in the forwarding table can be combined
    def compare_entries(self, first, sec):
        return first['netmask'] == sec['netmask'] and first['peer'] == sec['peer'] and first['localpref'] == sec[
            'localpref'] and first['origin'] == sec['origin'] and first['selfOrigin'] == sec['selfOrigin'] and first[
                   'ASPath'] == sec['ASPath'] and self.is_adjacent(first, sec)

    # combines 2 entries in the forwarding table
    def combine_entries(self, first, sec):
        # print("COMBINING ENTRIES: ")
        firstnw = ipToValue(first['network'])
        secnw = ipToValue(sec['network'])
        nm = first['netmask']
        nmlen = countNetmask(nm) - 1
        coveredNetworks = []
        first['isCombined'] = True
        sec['isCombined'] = True
        # add covered networks from both networks to new list
        coveredNetworks = self.update_covered_networks(coveredNetworks, first)
        coveredNetworks = self.update_covered_networks(coveredNetworks, sec)

        # creates new entry of combined entries in forwarding table and sets the combined entries to "is combined"
        if firstnw > secnw:
            entry = {"network": sec['network'], "netmask": netmaskNumToIP(nmlen), "peer": sec['peer'],
                     'localpref': sec['localpref'], 'origin': sec['origin'],
                     'selfOrigin': sec['selfOrigin'], 'ASPath': sec['ASPath'], 'isCombined': False, 'isOriginal': False,
                     'coveredNetworks': coveredNetworks}
            # print("COMBINED ENTRY: ", entry)
        else:
            entry = {"network": first['network'], "netmask": netmaskNumToIP(nmlen), "peer": sec['peer'],
                     'localpref': sec['localpref'], 'origin': sec['origin'],
                     'selfOrigin': sec['selfOrigin'], 'ASPath': sec['ASPath'], 'isCombined': False, 'isOriginal': False,
                     'coveredNetworks': coveredNetworks}
            # print("COMBINED ENTRY: ", entry)
        forwarding_table.append(entry)

    # Ex: 192.168.0.0/24 and 192.168.1.0/24
    # check if two entries in the forwarding table are adjacent by first checking if they have identical length of netmask,
    # then masking the networks with the length of the netmask - 1 and compare the masked networks
    def is_adjacent(self, first, sec):
        firstnm = first['netmask']
        secnm = sec['netmask']

        # checking for identical length of netmask
        if firstnm != secnm:
            return False

        # masking the networks with the length of the netmask - 1        
        newNetmaskLength = countNetmask(first['netmask']) - 1
        newNetmaskIP = netmaskNumToIP(newNetmaskLength)
        newNetmaskBinl = binl(newNetmaskIP)

        # compare the masked networks to see if they are identical
        firstMask = combine_submask(binl(first['network']), newNetmaskBinl)
        secMask = combine_submask(binl(sec['network']), newNetmaskBinl)
        return firstMask == secMask

    # 1. find entry in forwarding table that contains the network that is being revoked and remove that entry
    # 2. set all networks that are contained in that entry to have isCombined = false
    def disaggregate(self, src, msgs):
        disaggregated = []
        for msg in msgs:
            active_table = self.filter_combined()
            # print(active_table)
            # check all active entries
            for entry in active_table:
                if entry['isOriginal'] == True:
                    # print("ENTRY IS ORIGINAL")
                    if msg['network'] == entry['network'] and msg['netmask'] == entry['netmask'] and src == entry[
                        'peer']:
                        # remove entry from forwarding table
                        forwarding_table.remove(entry)
                else:
                    # print("ENTRY IS NOT ORIGINAL")
                    # check all networks in covered networks of root network
                    for coveredNetwork in entry['coveredNetworks']:
                        if msg['network'] == coveredNetwork['network'] and msg['netmask'] == coveredNetwork['netmask'] \
                                and src == coveredNetwork['peer']:
                            containingEntry = entry
                            # remove root network from forwarding table and disaggregate covered networks
                            forwarding_table.remove(containingEntry)
                            disaggregated = containingEntry['coveredNetworks']
                            # print("CONTAINING ENTRY: ", containingEntry)
                            # print("DISAGGREGATED LIST STILL EXISTS? ", disaggregated)

            # set all disaggregated entries to be not combined
            for entry in disaggregated:
                entry['isCombined'] = False
                if msg['network'] == entry['network'] and msg['netmask'] == entry['netmask'] and src == entry['peer']:
                    forwarding_table.remove(entry)

        # print(self.filter_combined())

    def coalesce(self):
        """ coalesce any routes that are right next to each other   """
        # print("BEFORE AGGREGATION ", forwarding_table)
        # The 2am Strategy (TM)
        runAgain = False
        # Sort list by network
        active_table = self.filter_combined()
        sorted_table = sorted(active_table, key=operator.itemgetter('network'))
        # Check if any adjacent entries are route neighbors, if so, combine
        for entry1, entry2 in zip(sorted_table, sorted_table[1:]):
            if self.compare_entries(entry1, entry2):
                runAgain = True
                self.combine_entries(entry1, entry2)

        # dump_table = self.filter_combined()
        # print("COALESCED TABLE: ", dump_table)
        # print("AFTER AGGREGATION ", forwarding_table)
        return runAgain

    def update(self, srcif, packet):
        """ handle update packets   """
        # adds entry to forwarding table (updates)
        entry = {"network": packet['msg']['network'], "netmask": packet['msg']['netmask'], "peer": packet['src'],
                 'localpref': packet['msg']['localpref'], 'origin': packet['msg']['origin'],
                 'selfOrigin': packet['msg']['selfOrigin'], 'ASPath': packet['msg']['ASPath'], 'isCombined': False,
                 'isOriginal': True, 'coveredNetworks': []}

        forwarding_table.append(entry)
        runAgain = True
        while runAgain:
            runAgain = self.coalesce()
        self.updates.append(packet)
        # print("Update forwarding table with new packet")
        announcement = {"src": "", "dst": "", "type": "update",
                        "msg": {"network": packet['msg']['network'], "netmask": packet['msg']['netmask'],
                                "localpref": packet['msg']['localpref'], "ASPath": [],
                                "origin": packet['msg']['origin'], "selfOrigin": packet['msg']['selfOrigin']}}
        for a in packet['msg']['ASPath']:
            announcement['msg']['ASPath'].append(a)

        announcement['msg']['ASPath'].append(args.asn)
        self.broadcast(announcement, packet)
        return False

    def revoke(self, packet):
        """ handle revoke packets """
        msg = packet['msg']
        # Handle disaggregation
        # Throw away entire forwarding table and rebuild using saved "update" and "revoke" messages
        self.disaggregate(packet['src'], msg)
        runAgain = True
        while runAgain:
            runAgain = self.coalesce()
        print("REVOKEREVOKEREVOKE")
        announcement = {"src": "", "dst": "", "type": "revoke",
                        "msg": msg}
        self.broadcast(announcement, packet)
        return True

    def filter_combined(self):
        """ Filters out all combined entries """
        dump_table = []
        for entry in forwarding_table:
            if entry['isCombined'] != True:
                dump_table.append(entry)
        return dump_table

    def dump(self, packet):
        """ handles dump table requests """
        # print("DUMPING")
        dump_table = self.filter_combined()
        dumppkg = {'src': packet['dst'], 'dst': packet['src'], 'type': 'table', 'msg': dump_table}
        dumppkgb = dictToByte(dumppkg)
        # print("Incoming packet from: ", packet['src'])
        # print("Forwarding Table: ", dump_table)
        self.sockets[packet['src']].sendall(dumppkgb)
        return True

    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)
                    message = {"type": msg['type'], "dst": msg['dst'], "src": msg['src'], "msg": msg['msg']}
                    messages.append(message)
                    if message['type'] == 'update':
                        self.update(srcif, msg)
                    if message['type'] == 'data':
                        self.forward(srcif, msg)
                    if message['type'] == 'dump':
                        self.dump(msg)
                    if message['type'] == 'revoke':
                        self.revoke(msg)
                else:
                    return


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