walkingonions/relay.py

#!/usr/bin/env python3

import random # For simulation, not cryptography!
import math

import nacl.utils
import nacl.signing
import nacl.public

import network
import dirauth

class CircuitCellMsg(network.NetMsg):
    """Send a message tagged with a circuit id."""
    def __init__(self, circuitid, msg):
        self.circid = circuitid
        self.msg = msg

    def __str__(self):
        return "C%d:%s" % (self.circid, self.msg)


class MultiplexedCircuitConnection(network.Connection):
    """A class representing a connection between a relay and either a
    client or a relay, transporting cells from various circuits."""
    def __init__(self):
        super().__init__()
        # The CellRelay managing this MultiplexedCircuitConnection
        self.cellrelay = None
        # The MultiplexedCircuitConnection at the other end
        self.peer = None
        # The function to call when the connection closes
        self.closer = lambda: 0

    def closed(self):
        self.closer()
        self.peer = None

    def close(self):
        if self.peer is not None:
            self.peer.closed()
        self.closed()

    def send_cell(self, circid, msg):
        """Send the given message, tagged for the given circuit id."""
        cell = CircuitCellMsg(circid, msg)
        self.peer.received(self.cellrelay.myaddr, cell)

    def received(self, peeraddr, cell):
        """Callback when a cell is received from the network."""
        circid, msg = cell.circid, cell.msg
        print("received", msg, "on circuit", circid, "from", peeraddr)


class CellRelay:
    """The class that manages the connections to other relays and
    clients.  Relays and clients both use this class to both create
    on-demand connections to relays, to gracefully handle the closing of
    connections, and to handle commands received over the
    connections."""

    def __init__(self, myaddr):
        # A dictionary of MultiplexedCircuitConnections to other hosts,
        # indexed by NetAddr
        self.connections = dict()
        self.myaddr = myaddr

    def get_connection_to(self, addr):
        """Get the MultiplexedCircuitConnection connected to the given
        NetAddr, creating one if none exists right now."""
        if addr in self.connections:
            return self.connections[addr]

        # Create the new connection
        newconn = network.thenetwork.connect(self.myaddr, addr)
        self.connections[addr] = newconn
        newconn.closer = lambda: self.connections.pop(addr)
        newconn.cellrelay = self

        return newconn


class Relay(network.Server):
    """The class representing an onion relay."""

    def __init__(self, dirauthaddrs, bw, flags):
        self.consensus = None
        self.dirauthaddrs = dirauthaddrs

        # Create the identity and onion keys
        self.idkey = nacl.signing.SigningKey.generate()
        self.onionkey = nacl.public.PrivateKey.generate()
        self.name = self.idkey.verify_key.encode(encoder=nacl.encoding.HexEncoder).decode("ascii")

        # Bind to the network to get a network address
        self.netaddr = network.thenetwork.bind(self)

        # Our bandwidth and flags
        self.bw = bw
        self.flags = flags

        # Register for epoch change notification
        network.thenetwork.wantepochticks(self, True, end=True)
        network.thenetwork.wantepochticks(self, True)

        # Create the CellRelay connection manager
        self.cellrelay = CellRelay(self.netaddr)

        self.uploaddesc()

    def epoch_ending(self, epoch):
        # Download the new consensus, which will have been created
        # already since the dirauths' epoch_ending callbacks happened
        # before the relays'.
        a = random.choice(self.dirauthaddrs)
        c = network.thenetwork.connect(self, a)
        self.consensus = c.getconsensus()
        c.close()

    def newepoch(self, epoch):
        self.uploaddesc()

    def uploaddesc(self):
        # Upload the descriptor for the epoch to come
        descdict = dict();
        descdict["epoch"] = network.thenetwork.getepoch() + 1
        descdict["idkey"] = self.idkey.verify_key
        descdict["onionkey"] = self.onionkey.public_key
        descdict["addr"] = self.netaddr
        descdict["bw"] = self.bw
        descdict["flags"] = self.flags
        desc = dirauth.RelayDescriptor(descdict)
        desc.sign(self.idkey)
        desc.verify()

        descmsg = dirauth.DirAuthUploadDescMsg(desc)

        # Upload them
        for a in self.dirauthaddrs:
            c = network.thenetwork.connect(self, a)
            c.sendmsg(descmsg)
            c.close()

    def connected(self, peer):
        """Callback invoked when someone (client or relay) connects to
        us.  Create a pair of linked MultiplexedCircuitConnections and
        return the peer half to the peer."""

        # Create the linked pair
        peerconn = MultiplexedCircuitConnection()
        ourconn = MultiplexedCircuitConnection()
        peerconn.peer = ourconn
        ourconn.peer = peerconn

        return peerconn

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

    # Start some relays
    numrelays = 10
    relays = []
    for i in range(numrelays):
        # Relay bandwidths (at least the ones fast enough to get used)
        # in the live Tor network (as of Dec 2019) are well approximated
        # by (200000-(200000-25000)/3*log10(x)) where x is a
        # uniform integer in [1,2500]
        x = random.randint(1,2500)
        bw = int(200000-(200000-25000)/3*math.log10(x))
        relays.append(Relay(dirauthaddrs, bw, 0))

    # Tick the epoch
    network.thenetwork.nextepoch()

    dirauth.DirAuth.consensus.verify(network.thenetwork.dirauthkeys())

    print('ticked; epoch=', network.thenetwork.getepoch())

    c = relays[3].cellrelay.get_connection_to(relays[3].consensus.consdict['relays'][5].descdict['addr'])

    c.send_cell(1, network.StringNetMsg("test"))
    c.close()
    c2 = relays[3].cellrelay.get_connection_to(relays[3].consensus.consdict['relays'][6].descdict['addr'])
    c = relays[3].cellrelay.get_connection_to(relays[3].consensus.consdict['relays'][5].descdict['addr'])
    c.send_cell(2, network.StringNetMsg("cell"))
    c3 = relays[3].cellrelay.get_connection_to(relays[3].consensus.consdict['relays'][1].descdict['addr'])
    c = relays[3].cellrelay.get_connection_to(relays[3].consensus.consdict['relays'][5].descdict['addr'])
    c.send_cell(3, network.StringNetMsg("again"))
    c.close()