#!/usr/bin/env python3

import random # For simulation, not cryptography!
import bisect

import nacl.encoding
import nacl.signing
import network

# A relay descriptor is a dict containing:
#  epoch: epoch id
#  idkey: a public identity key
#  onionkey: a public onion key
#  addr: a network address
#  bw: bandwidth
#  flags: relay flags
#  vrfkey: a VRF public key (Single-Pass Walking Onions only)
#  sig: a signature over the above by the idkey
class RelayDescriptor:
    def __init__(self, descdict):
        self.descdict = descdict

    def __str__(self, withsig = True):
        res = "RelayDesc [\n"
        for k in ["epoch", "idkey", "onionkey", "addr", "bw", "flags",
                    "vrfkey", "sig"]:
            if k in self.descdict:
                if k == "idkey" or k == "onionkey":
                    res += "  " + k + ": " + self.descdict[k].encode(encoder=nacl.encoding.HexEncoder).decode("ascii") + "\n"
                elif k == "sig":
                    if withsig:
                        res += "  " + k + ": " + nacl.encoding.HexEncoder.encode(self.descdict[k]).decode("ascii") + "\n"
                else:
                    res += "  " + k + ": " + str(self.descdict[k]) + "\n"
        res += "]\n"
        return res

    def sign(self, signingkey):
        serialized = self.__str__(False)
        signed = signingkey.sign(serialized.encode("ascii"))
        self.descdict["sig"] = signed.signature


def verify_relaydesc(desc):
    serialized = desc.__str__(False)
    desc.descdict["idkey"].verify(serialized.encode("ascii"), desc.descdict["sig"])


# A consensus is a dict containing:
#  epoch: epoch id
#  numrelays: total number of relays
#  totbw: total bandwidth of relays
#  relays: list of relay descriptors (Vanilla Onion Routing only)
#  sigs: list of signatures from the dirauths
class Consensus:
    def __init__(self, epoch, relays):
        relays = [ d for d in relays if d.descdict['epoch'] == epoch ]
        self.consdict = dict()
        self.consdict['epoch'] = epoch
        self.consdict['numrelays'] = len(relays)
        self.consdict['totbw'] = sum([ d.descdict['bw'] for d in relays ])
        self.consdict['relays'] = relays

    def __str__(self, withsigs = True):
        res = "Consensus [\n"
        for k in ["epoch", "numrelays", "totbw"]:
            if k in self.consdict:
                res += "  " + k + ": " + str(self.consdict[k]) + "\n"
        for r in self.consdict['relays']:
            res += str(r)
        if withsigs and ('sigs' in self.consdict):
            for s in self.consdict['sigs']:
                res += "  sig: " + nacl.encoding.HexEncoder.encode(s).decode("ascii") + "\n"
        res += "]\n"
        return res

    def sign(self, signingkey, index):
        """Use the given signing key to sign the consensus, storing the
        result in the sigs list at the given index."""
        serialized = self.__str__(False)
        signed = signingkey.sign(serialized.encode("ascii"))
        if 'sigs' not in self.consdict:
            self.consdict['sigs'] = []
        if index >= len(self.consdict['sigs']):
            self.consdict['sigs'].extend([None] * (index+1-len(self.consdict['sigs'])))
        self.consdict['sigs'][index] = signed.signature

    def bw_cdf(self):
        """Create the array of cumulative bandwidth values from a consensus.
        The array (cdf) will have the same length as the number of relays
        in the consensus.  cdf[0] = 0, and cdf[i] = cdf[i-1] + relay[i-1].bw."""
        cdf = [0]
        for r in self.consdict['relays']:
            cdf.append(cdf[-1]+r.descdict['bw'])
        # Remove the last item, which should be the sum of all the bws
        cdf.pop()
        print('cdf=', cdf)
        return cdf

    def select_weighted_relay(self, cdf):
        """Use the cdf generated by bw_cdf to select a relay with
        probability proportional to its bw weight."""
        val = random.randint(1, self.consdict['totbw'])
        idx = bisect.bisect_left(cdf, val)
        return self.consdict['relays'][idx-1]



def verify_consensus(consensus, verifkeylist):
    """Use the given list of verification keys to check the
    signatures on the consensus."""
    serialized = consensus.__str__(False)
    for i, vk in enumerate(verifkeylist):
        vk.verify(serialized.encode("ascii"), consensus.consdict['sigs'][i])


class DirAuthNetMsg(network.NetMsg):
    """The subclass of NetMsg for messages to and from directory
    authorities."""

class DirAuthUploadDescMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for uploading a relay
    descriptor."""

    def __init__(self, desc):
        self.desc = desc

class DirAuthDelDescMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for deleting a relay
    descriptor."""

    def __init__(self, desc):
        self.desc = desc

class DirAuthGetConsensusMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the consensus."""

class DirAuthConsensusMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the consensus."""

    def __init__(self, consensus):
        self.consensus = consensus

class DirAuthGetENDIVEMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the ENDIVE."""

class DirAuthENDIVEMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the ENDIVE."""

    def __init__(self, endive):
        self.endive = endive


class DirAuthConnection(network.ClientConnection):
    """The subclass of Connection for connections to directory
    authorities."""

    def __init__(self, peer = None):
        super().__init__(peer)

    def uploaddesc(self, desc):
        """Upload our RelayDescriptor to the DirAuth."""
        self.sendmsg(DirAuthUploadDescMeg(desc))

    def getconsensus(self):
        self.consensus = None
        self.sendmsg(DirAuthGetConsensusMsg())
        return self.consensus

    def getENDIVE(self):
        self.endive = None
        self.sendmsg(DirAuthGetENDIVEMsg())
        return self.endive

    def receivedfromserver(self, msg):
        if isinstance(msg, DirAuthConsensusMsg):
            self.consensus = msg.consensus
        elif isinstance(msg, DirAuthENDIVEMsg):
            self.endive = msg.endive
        else:
            raise TypeError('Not a server-originating DirAuthNetMsg', msg)
    

class DirAuth(network.Server):
    """The class representing directory authorities."""

    # We simulate the act of computing the consensus by keeping a
    # class-static dict that's accessible to all of the dirauths
    # This dict is indexed by epoch, and the value is itself a dict
    # indexed by the stringified descriptor, with value a pair of (the
    # number of dirauths that saw that descriptor, the descriptor
    # itself).
    uploadeddescs = dict()
    consensus = None
    endive = None

    def __init__(self, me, tot):
        """Create a new directory authority. me is the index of which
        dirauth this one is (starting from 0), and tot is the total
        number of dirauths."""
        self.me = me
        self.tot = tot
        self.name = "Dirauth %d of %d" % (me+1, tot)

        # Create the dirauth signature keypair
        self.sigkey = nacl.signing.SigningKey.generate()

        network.thenetwork.setdirauthkey(me, self.sigkey.verify_key)
        network.thenetwork.wantepochticks(self, True, True)

    def connected(self, client):
        """Callback invoked when a client connects to us. This callback
        creates the DirAuthConnection that will be passed to the
        client."""

        # We don't actually need to keep per-connection state at
        # dirauths, even in long-lived connections, so this is
        # particularly simple.
        return DirAuthConnection(self)

    def generate_consensus(self, epoch):
        """Generate the consensus (and ENDIVE, if using Walking Onions)
        for the given epoch, which should be the one after the one
        that's currently about to end."""
        threshold = int(self.tot/2)+1
        consensusdescs = []
        for numseen, desc in DirAuth.uploadeddescs[epoch].values():
            if numseen >= threshold:
                consensusdescs.append(desc)
        DirAuth.consensus = Consensus(epoch, consensusdescs)


    def epoch_ending(self, epoch):
        # Only dirauth 0 actually needs to generate the consensus
        # because of the shared class-static state, but everyone has to
        # sign it.  Note that this code relies on dirauth 0's
        # epoch_ending callback being called before any of the other
        # dirauths'.
        if self.me == 0:
            self.generate_consensus(epoch+1)
            del DirAuth.uploadeddescs[epoch+1]
        DirAuth.consensus.sign(self.sigkey, self.me)

    def received(self, client, msg):
        if isinstance(msg, DirAuthUploadDescMsg):
            # Check the uploaded descriptor for sanity
            epoch = msg.desc.descdict['epoch']
            if epoch != network.thenetwork.getepoch() + 1:
                return
            # Store it in the class-static dict
            if epoch not in DirAuth.uploadeddescs:
                DirAuth.uploadeddescs[epoch] = dict()
            descstr = str(msg.desc)
            if descstr not in DirAuth.uploadeddescs[epoch]:
                DirAuth.uploadeddescs[epoch][descstr] = (1, msg.desc)
            else:
                DirAuth.uploadeddescs[epoch][descstr] = \
                    (DirAuth.uploadeddescs[epoch][descstr][0]+1,
                     DirAuth.uploadeddescs[epoch][descstr][1])
        elif isinstance(msg, DirAuthDelDescMsg):
            # Check the uploaded descriptor for sanity
            epoch = msg.desc.descdict['epoch']
            if epoch != network.thenetwork.getepoch() + 1:
                return
            # Remove it from the class-static dict
            if epoch not in DirAuth.uploadeddescs:
                return
            descstr = str(msg.desc)
            if descstr not in DirAuth.uploadeddescs[epoch]:
                return
            elif DirAuth.uploadeddescs[epoch][descstr][0] == 1:
                del DirAuth.uploadeddescs[epoch][descstr]
            else:
                DirAuth.uploadeddescs[epoch][descstr] = \
                    (DirAuth.uploadeddescs[epoch][descstr][0]-1,
                     DirAuth.uploadeddescs[epoch][descstr][1])
        elif isinstance(msg, DirAuthGetConsensusMsg):
            client.reply(DirAuthConsensusMsg(DirAuth.consensus))
        elif isinstance(msg, DirAuthGetENDIVEMsg):
            client.reply(DirAuthENDIVEMsg(DirAuth.endive))
        else:
            raise TypeError('Not a client-originating DirAuthNetMsg', msg)

    def closed(self):
        pass


if __name__ == '__main__':
    # Start some dirauths
    numdirauths = 9
    dirauthaddrs = []
    for i in range(numdirauths):
        dirauth = DirAuth(i, numdirauths)
        dirauthaddrs.append(network.thenetwork.bind(dirauth))

    for a in dirauthaddrs:
        print(a,end=' ')
    print()

    network.thenetwork.nextepoch()