walkingonions/dirauth.py

#!/usr/bin/env python3

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(self):
        serialized = self.__str__(False)
        self.descdict["idkey"].verify(serialized.encode("ascii"), self.descdict["sig"])

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 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 of the number of
    # dirauths that saw that descriptor.
    uploadeddescs = dict()

    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)
        self.consensus = None
        self.endive = None
        network.thenetwork.wantepochticks(self, 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 newepoch(self, epoch):
        print('New epoch', epoch, 'for', self)

    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
            else:
                DirAuth.uploadeddescs[epoch][descstr] += 1
        elif isinstance(msg, DirAuthGetConsensusMsg):
            client.sendmsg(DirAuthConsensusMsg(self.consensus))
        elif isinstance(msg, DirAuthGetENDIVEMsg):
            client.sendmsg(DirAuthENDIVEMsg(self.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()