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 RelayNetMsg(network.NetMsg):
    """The subclass of NetMsg for messages between relays and either
    relays or clients."""


class RelayGetConsensusMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for fetching the consensus."""


class RelayConsensusMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for returning the consensus."""

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


class RelayRandomHopMsg(RelayNetMsg):
    """A message used for testing, that hops from relay to relay
    randomly until its TTL expires."""

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

    def __str__(self):
        return "RandomHop TTL=%d" % self.ttl


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

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


class RelayFallbackTerminationError(Exception):
    """An exception raised when someone tries to terminate a fallback
    relay."""


class CircuitHandler:
    """A class for managing sending and receiving encrypted cells on a
    particular circuit."""

    def __init__(self, channel, circid):
        self.channel = channel
        self.circid = circid
        self.send_cell = self.channel_send_cell
        self.received_cell = self.channel_received_cell

    def channel_send_cell(self, cell):
        """Send a cell on this circuit."""
        self.channel.send_msg(CircuitCellMsg(self.circid, cell))

    def channel_received_cell(self, cell, peeraddr, peer):
        """A cell has been received on this circuit.  Forward it to the
        channel's received_cell callback."""
        self.channel.cellhandler.received_cell(self.circid, cell, peeraddr, peer)


class Channel(network.Connection):
    """A class representing a channel between a relay and either a
    client or a relay, transporting cells from various circuits."""
    def __init__(self):
        super().__init__()
        # The CellRelay managing this Channel
        self.cellhandler = None
        # The Channel at the other end
        self.peer = None
        # The function to call when the connection closes
        self.closer = lambda: 0
        # The next circuit id to use on this channel.  The party that
        # opened the channel uses even numbers; the receiving party uses
        # odd numbers.
        self.next_circid = None
        # A map for CircuitHandlers to use for each open circuit on the
        # channel
        self.circuithandlers = dict()

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

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

    def new_circuit(self):
        """Allocate a new circuit on this channel, returning the new
        circuit's id."""
        circid = self.next_circid
        self.next_circid += 2
        self.circuithandlers[circid] = CircuitHandler(self, circid)
        return circid

    def new_circuit_with_circid(self, circid):
        """Allocate a new circuit on this channel, with the circuit id
        received from our peer."""
        self.circuithandlers[circid] = CircuitHandler(self, circid)

    def send_cell(self, circid, cell):
        """Send the given message on the given circuit, encrypting or
        decrypting as needed."""
        self.circuithandlers[circid].send_cell(cell)

    def send_raw_cell(self, circid, cell):
        """Send the given message, tagged for the given circuit id. No
        encryption or decryption is done."""
        self.send_msg(CircuitCellMsg(self.circid, self.cell))

    def send_msg(self, msg):
        """Send the given NetMsg on the channel."""
        self.peer.received(self.cellhandler.myaddr, msg)

    def received(self, peeraddr, msg):
        """Callback when a message is received from the network."""
        if isinstance(msg, CircuitCellMsg):
            circid, cell = msg.circid, msg.cell
            self.circuithandlers[circid].received_cell(cell, peeraddr, self.peer)
        else:
            self.cellhandler.received_msg(msg, peeraddr, self.peer)


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

    def __init__(self, myaddr, dirauthaddrs):
        # A dictionary of Channels to other hosts, indexed by NetAddr
        self.channels = dict()
        self.myaddr = myaddr
        self.dirauthaddrs = dirauthaddrs
        self.consensus = None

    def terminate(self):
        """Close all connections we're managing."""
        while self.channels:
            channelitems = iter(self.channels.items())
            addr, channel = next(channelitems)
            print('closing channel', addr, channel)
            channel.close()

    def add_channel(self, channel, peeraddr):
        """Add the given channel to the list of channels we are
        managing.  If we are already managing a channel to the same
        peer, close it first."""
        if peeraddr in self.channels:
            self.channels[peeraddr].close()

        channel.cellhandler = self
        self.channels[peeraddr] = channel
        channel.closer = lambda: self.channels.pop(peeraddr)

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

        # Create the new channel
        newchannel = network.thenetwork.connect(self.myaddr, addr)
        self.channels[addr] = newchannel
        newchannel.closer = lambda: self.channels.pop(addr)
        newchannel.cellhandler = self

        return newchannel

    def received_msg(self, msg, peeraddr, peer):
        """Callback when a NetMsg not specific to a circuit is
        received."""
        print("CellHandler: Node %s received msg %s from %s" % (self.myaddr, msg, peeraddr))

    def received_cell(self, circid, cell, peeraddr, peer):
        """Callback with a circuit-specific cell is received."""
        print("CellHandler: Node %s received cell on circ %d: %s from %s" % (self.myaddr, circid, cell, peeraddr))

    def send_msg(self, msg, peeraddr):
        """Send a message to the peer with the given address."""
        channel = self.get_channel_to(peeraddr)
        channel.send_msg(msg)

    def send_cell(self, circid, cell, peeraddr):
        """Send a cell on the given circuit to the peer with the given
        address."""
        channel = self.get_channel_to(peeraddr)
        channel.send_cell(circid, cell)


class CellRelay(CellHandler):
    """The subclass of CellHandler for relays."""

    def __init__(self, myaddr, dirauthaddrs):
        super().__init__(myaddr, dirauthaddrs)

    def get_consensus(self):
        """Download a fresh consensus from a random dirauth."""
        a = random.choice(self.dirauthaddrs)
        c = network.thenetwork.connect(self, a)
        self.consensus = c.getconsensus()
        dirauth.Consensus.verify(self.consensus, network.thenetwork.dirauthkeys())
        c.close()

    def received_msg(self, msg, peeraddr, peer):
        """Callback when a NetMsg not specific to a circuit is
        received."""
        print("CellRelay: Node %s received msg %s from %s" % (self.myaddr, msg, peeraddr))
        if isinstance(msg, RelayRandomHopMsg):
            if msg.ttl > 0:
                # Pick a random next hop from the consensus
                nexthop = random.choice(self.consensus.consdict['relays'])
                nextaddr = nexthop.descdict['addr']
                self.send_msg(RelayRandomHopMsg(msg.ttl-1), nextaddr)
        elif isinstance(msg, RelayGetConsensusMsg):
            self.send_msg(RelayConsensusMsg(self.consensus), peeraddr)
        else:
            return super().received_msg(msg, peeraddr, peer)

    def received_cell(self, circid, cell, peeraddr, peer):
        """Callback with a circuit-specific cell is received."""
        print("CellRelay: Node %s received cell on circ %d: %s from %s" % (self.myaddr, circid, cell, peeraddr))
        return super().received_cell(circid, cell, peeraddr, peer)


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

    def __init__(self, dirauthaddrs, bw, flags):
        # 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.cellhandler = CellRelay(self.netaddr, dirauthaddrs)

        # Initially, we're not a fallback relay
        self.is_fallbackrelay = False

        self.uploaddesc()

    def terminate(self):
        """Stop this relay."""

        if self.is_fallbackrelay:
            # Fallback relays must not (for now) terminate
            raise RelayFallbackTerminationError(self)

        # Stop listening for epoch ticks
        network.thenetwork.wantepochticks(self, False, end=True)
        network.thenetwork.wantepochticks(self, False)

        # Tell the dirauths we're going away
        self.uploaddesc(False)

        # Close connections to other relays
        self.cellhandler.terminate()

        # Stop listening to our own bound port
        self.close()

    def set_is_fallbackrelay(self, isfallback = True):
        """Set this relay to be a fallback relay (or unset if passed
        False)."""
        self.is_fallbackrelay = isfallback

    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'.
        self.cellhandler.get_consensus()

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

    def uploaddesc(self, upload=True):
        # Upload the descriptor for the epoch to come, or delete a
        # previous upload if upload=False
        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)
        dirauth.RelayDescriptor.verify(desc)

        if upload:
            descmsg = dirauth.DirAuthUploadDescMsg(desc)
        else:
            # Note that this relies on signatures being deterministic;
            # otherwise we'd need to save the descriptor we uploaded
            # before so we could tell the airauths to delete the exact
            # one
            descmsg = dirauth.DirAuthDelDescMsg(desc)

        # Upload them
        for a in self.cellhandler.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 Channels and return the peer half
        to the peer."""

        # Create the linked pair
        if peer is self.netaddr:
            # A self-loop?  We'll allow it.
            peerchannel = Channel()
            peerchannel.peer = peerchannel
            peerchannel.next_circid = 2
            return peerchannel

        peerchannel = Channel()
        ourchannel = Channel()
        peerchannel.peer = ourchannel
        peerchannel.next_circid = 2
        ourchannel.peer = peerchannel
        ourchannel.next_circid = 1

        # Add our channel to the CellRelay
        self.cellhandler.add_channel(ourchannel, peer)

        return peerchannel

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

    # The fallback relays are a hardcoded list of about 5% of the
    # relays, used by clients for bootstrapping
    numfallbackrelays = int(numrelays * 0.05) + 1
    fallbackrelays = random.sample(relays, numfallbackrelays)
    for r in fallbackrelays:
        r.set_is_fallbackrelay()
    network.thenetwork.setfallbackrelays(fallbackrelays)

    # Tick the epoch
    network.thenetwork.nextepoch()

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

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

    relays[3].cellhandler.send_msg(RelayRandomHopMsg(30), relays[5].netaddr)

    # See what channels exist and do a consistency check
    for r in relays:
        print("%s: %s" % (r.netaddr, [ str(k) for k in r.cellhandler.channels.keys()]))
        raddr = r.netaddr
        for ad, ch in r.cellhandler.channels.items():
            if ch.peer.cellhandler.myaddr != ad:
                print('address mismatch:', raddr, ad, ch.peer.cellhandler.myaddr)

            if ch.peer.cellhandler.channels[raddr].peer is not ch:
                print('asymmetry:', raddr, ad, ch, ch.peer.cellhandler.channels[raddr].peer)

    # Stop some relays
    relays[3].terminate()
    del relays[3]
    relays[5].terminate()
    del relays[5]
    relays[7].terminate()
    del relays[7]

    # Tick the epoch
    network.thenetwork.nextepoch()

    print(dirauth.DirAuth.consensus)

    # See what channels exist and do a consistency check
    for r in relays:
        print("%s: %s" % (r.netaddr, [ str(k) for k in r.cellhandler.channels.keys()]))
        raddr = r.netaddr
        for ad, ch in r.cellhandler.channels.items():
            if ch.peer.cellhandler.myaddr != ad:
                print('address mismatch:', raddr, ad, ch.peer.cellhandler.myaddr)

            if ch.peer.cellhandler.channels[raddr].peer is not ch:
                print('asymmetry:', raddr, ad, ch, ch.peer.cellhandler.channels[raddr].peer)

    channel = relays[3].cellhandler.get_channel_to(relays[5].netaddr)
    circid = channel.new_circuit()
    peerchannel = relays[5].cellhandler.get_channel_to(relays[3].netaddr)
    peerchannel.new_circuit_with_circid(circid)
    relays[3].cellhandler.send_cell(circid, network.StringNetMsg("test"), relays[5].netaddr)