walkingonions/relay.py

#!/usr/bin/env python3

import random # For simulation, not cryptography!
import math
import sys

import nacl.utils
import nacl.signing
import nacl.public
import nacl.hash
import nacl.bindings

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. Sent by
    clients to relays."""


class RelayConsensusMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for returning the consensus from
    relays to clients, in response to RelayGetConsensusMsg."""

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


class RelayGetConsensusDiffMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for fetching the consensus, if the
    requestor already has the previous consensus. Sent by clients to
    relays."""


class RelayConsensusDiffMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for returning the consensus, if the
    requestor already has the previous consensus.  We don't _actually_
    produce the diff at this time; we just charge fewer bytes for this
    message. Sent by relays to clients in response to
    RelayGetConsensusDiffMsg."""

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

    def size(self):
        if network.symbolic_byte_counters:
            return super().size()
        return math.ceil(RelayConsensusMsg(self.consensus).size() \
                            * network.P_Delta)


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. This is the container
    class for all RelayCell messages."""

    def __init__(self, circuitid, cell):
        self.circid = circuitid
        self.cell = cell

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

    def size(self):
        # circuitids are 4 bytes
        return 4 + self.cell.size()


class RelayCell(RelayNetMsg):
    """All cells (which are sent inside a CircuitCellMsg, and so do not
    need their own circuitid) should be a subclass of this class."""


class StringCell(RelayCell):
    """Send an arbitrary string as a cell."""
    def __init__(self, str):
        self.data = str

    def __str__(self):
        return self.data.__str__()


class CloseCell(RelayCell):
    """Close the circuit this cell was sent on.  It should be sent
    _unencrypted_ (not within an EncryptedCell), and relays that receive
    one should forward it along the adjacent circuit, then close both
    the circuit it was received on and the adjacent one."""



# It is intentional that VanillaCreateCircuitMsg is a RelayNetMsg and
# not a RelayCell.  This is the message that _creates_ the circuit, so
# it can't be sent as a cell _within_ the circuit.
class VanillaCreateCircuitMsg(RelayNetMsg):
    """The message for requesting circuit creation in Vanilla Onion
    Routing."""

    def __init__(self, circid, ntor_request):
        self.circid = circid
        self.ntor_request = ntor_request


class VanillaCreatedCircuitCell(RelayCell):
    """The message for responding to circuit creation in Vanilla Onion
    Routing."""

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


class VanillaExtendCircuitCell(RelayCell):
    """The message for requesting circuit extension in Vanilla Onion
    Routing."""

    def __init__(self, hopaddr, ntor_request):
        self.hopaddr = hopaddr
        self.ntor_request = ntor_request


class VanillaExtendedCircuitCell(RelayCell):
    """The message for responding to circuit extension in Vanilla Onion
    Routing."""

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

# It is intentional that TelescopingCreateCircuitMsg is a RelayNetMsg and
# not a RelayCell.  This is the message that _creates_ the circuit, so
# it can't be sent as a cell _within_ the circuit.
class TelescopingCreateCircuitMsg(RelayNetMsg):
    """The message for requesting circuit creation in Telescoping Onion
    Routing."""

    def __init__(self, circid, ntor_request):
        self.circid = circid
        self.ntor_request = ntor_request

class TelescopingCreatedCircuitCell(RelayCell):
    """The message for responding to circuit creation in Telescoping Walking
    Onions."""

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

class TelescopingExtendCircuitCell(RelayCell):
    """The message for requesting circuit extension in Telescoping Walking
    Onionss."""

    def __init__(self, idx, ntor_request):
        self.idx = idx
        self.ntor_request = ntor_request

class TelescopingExtendedCircuitCell(RelayCell):
    """The message for responding to circuit extension in Telescoping Walking
    Onions."""

    def __init__(self, ntor_reply, snip):
        self.ntor_reply = ntor_reply
        self.snip = snip


class EncryptedCell(RelayCell):
    """Send a message encrypted with a symmetric key.  In this
    implementation, the encryption is not really done.  A hash of the
    key is stored with the message so that it can be checked at
    decryption time."""

    def __init__(self, key, msg):
        self.keyhash = nacl.hash.sha256(key)
        self.plaintext = msg

    def decrypt(self, key):
        keyhash = nacl.hash.sha256(key)
        if keyhash != self.keyhash:
            raise ValueError("EncryptedCell key mismatch")
        return self.plaintext

    def size(self):
        # Current Tor actually has no overhead for encryption
        return self.plaintext.size()


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


class Sphinx:
    """Implement the public-key reblinding technique based on Sphinx.
    This does a few more public-key operations than it would strictly
    need to if we were using a group implementation that (unlike nacl)
    supported the operations we needed directly.  The biggest issue is
    that nacl insists the high bit is set on private keys, which means
    we can't just multiply private keys together to get a new private
    key, and do a single DH operation with that resulting key; we have
    to perform a linear number of DH operations instead, per node in the
    circuit, so a quadratic number of DH operations total."""

    @staticmethod
    def makeblindkey(shared_secret, domain_separator, perfstats):
        """Create a Sphinx reblinding key (a PrivateKey) out of a shared
        secret and a domain separator (both bytestrings).  The domain
        separator is just a constant bytestring like b'data' or
        b'circuit' for the data-protecting and circuit-protecting
        public-key elements respectively."""
        rawkey = nacl.hash.sha256(domain_separator + shared_secret,
                encoder=nacl.encoding.RawEncoder)
        perfstats.keygens += 1

        # The PrivateKey constructor does the Curve25519 pinning of
        # certain bits of the key to 0 and 1
        return nacl.public.PrivateKey(rawkey)

    @staticmethod
    def reblindpubkey(blindkey, pubkey, perfstats):
        """Create a Sphinx reblinded PublicKey out of a reblinding key
        (output by makeblindkey) and a (possibly already reblinded)
        PublicKey."""
        new_pubkey =  nacl.bindings.crypto_scalarmult(bytes(blindkey),
                bytes(pubkey))
        perfstats.dhs += 1

        return nacl.public.PublicKey(new_pubkey)

    @staticmethod
    def client(client_privkey, blindkey_list, server_pubkey,
            domain_separator, is_last, perfstats):
        """Given the client's PrivateKey, a (possibly empty) list of
        reblinding keys, and the server's PublicKey, produce the shared
        secret and the new blinding key (to add to the list).  The
        domain separator is as above.  If is_last is true, don't bother
        creating the new blinding key, since this is the last iteration,
        and we won't be using it."""
        reblinded_server_pubkey = server_pubkey
        for blindkey in blindkey_list:
            reblinded_server_pubkey = Sphinx.reblindpubkey(blindkey,
                    reblinded_server_pubkey, perfstats)

        sharedsecret = nacl.public.Box(client_privkey,
                reblinded_server_pubkey).shared_key()
        perfstats.dhs += 1

        if is_last:
            blindkey = None
        else:
            blindkey = Sphinx.makeblindkey(sharedsecret,
                    domain_separator, perfstats)

        return sharedsecret, blindkey

    @staticmethod
    def server(client_pubkey, server_privkey, domain_separator, is_last,
            perfstats):
        """Given the client's PublicKey and the server's PrivateKey,
        produce the shared secret and the new reblinded client
        PublicKey.  The domain separator is as above.  If is_last is
        True, don't bother generating the new PublicKey, since we're the
        last server in the chain, and won't be using it."""
        sharedsecret = nacl.public.Box(server_privkey,
                client_pubkey).shared_key()
        perfstats.dhs += 1

        if is_last:
            blinded_pubkey = None
        else:
            blindkey = Sphinx.makeblindkey(sharedsecret, domain_separator,
                    perfstats)

            blinded_pubkey = Sphinx.reblindpubkey(blindkey, client_pubkey,
                    perfstats)

        return sharedsecret, blinded_pubkey


class NTor:
    """A class implementing the ntor one-way authenticated key agreement
    scheme.  The details are not exactly the same as either the ntor
    paper or Tor's implementation, but it will agree on keys and have
    the same number of public key operations."""

    def __init__(self, perfstats):
        self.perfstats = perfstats
        # Only used for Single-Pass Walking Onions; it is the sequence
        # of blinding keys used by Sphinx
        self.blinding_keys = []

    def request(self):
        """Create the ntor request message: X = g^x."""
        self.client_ephem_key = nacl.public.PrivateKey.generate()
        self.perfstats.keygens += 1
        return self.client_ephem_key.public_key

    @staticmethod
    def reply(onion_privkey, idpubkey, client_pubkey, perfstats,
            sphinx_domainsep=None):
        """The server calls this static method to produce the ntor reply
        message: (Y = g^y, B = g^b, A = H(M, "verify")) and the shared
        secret S = H(M, "secret") for M = (X^y,X^b,ID,B,X,Y). If
        sphinx_domainsep is not None, also compute and return the Sphinx
        reblinded client request to pass to the next server."""
        server_ephem_key = nacl.public.PrivateKey.generate()
        perfstats.keygens += 1
        xykey = nacl.public.Box(server_ephem_key, client_pubkey).shared_key()
        xbkey = nacl.public.Box(onion_privkey, client_pubkey).shared_key()
        perfstats.dhs += 2
        M = xykey + xbkey + \
                idpubkey.encode(encoder=nacl.encoding.RawEncoder) + \
                onion_privkey.public_key.encode(encoder=nacl.encoding.RawEncoder) + \
                server_ephem_key.public_key.encode(encoder=nacl.encoding.RawEncoder)
        A = nacl.hash.sha256(M + b'verify', encoder=nacl.encoding.RawEncoder)
        S = nacl.hash.sha256(M + b'secret', encoder=nacl.encoding.RawEncoder)

        if sphinx_domainsep is not None:
            blindkey = Sphinx.makeblindkey(S, sphinx_domainsep, perfstats)
            blinded_client_pubkey = Sphinx.reblindpubkey(blindkey,
                    client_pubkey, perfstats)
            return ((server_ephem_key.public_key, onion_privkey.public_key, A),
                    S), blinded_client_pubkey
        else:
            return ((server_ephem_key.public_key, onion_privkey.public_key, A),
                    S)

    def verify(self, reply, onion_pubkey, idpubkey, sphinx_domainsep=None):
        """The client calls this method to verify the ntor reply
        message, passing the onion and id public keys for the server
        it's expecting to be talking to.  If sphinx_domainsep is not
        None, also compute the reblinding key so that the client can
        reuse this same NTor object for the next server.  Returns the
        shared secret on success, or raises ValueError on failure."""
        server_ephem_pubkey, server_onion_pubkey, authtag = reply
        if onion_pubkey != server_onion_pubkey:
            raise ValueError("NTor onion pubkey mismatch")
        # We use the blinding keys if present; if they're not present
        # (because we're not in Single-Pass Walking Onions), the loops
        # are just empty anyway, so everything will work in the usual
        # unblinded way.
        reblinded_server_ephem_pubkey = server_ephem_pubkey
        for blindkey in self.blinding_keys:
            reblinded_server_ephem_pubkey = Sphinx.reblindpubkey(blindkey,
                    reblinded_server_ephem_pubkey, self.perfstats)
        xykey = nacl.public.Box(self.client_ephem_key,
                reblinded_server_ephem_pubkey).shared_key()
        reblinded_onion_pubkey = onion_pubkey
        for blindkey in self.blinding_keys:
            reblinded_onion_pubkey = Sphinx.reblindpubkey(blindkey,
                    reblinded_onion_pubkey, self.perfstats)
        xbkey = nacl.public.Box(self.client_ephem_key,
                reblinded_onion_pubkey).shared_key()
        self.perfstats.dhs += 2
        M = xykey + xbkey + \
                idpubkey.encode(encoder=nacl.encoding.RawEncoder) + \
                onion_pubkey.encode(encoder=nacl.encoding.RawEncoder) + \
                server_ephem_pubkey.encode(encoder=nacl.encoding.RawEncoder)
        Acheck = nacl.hash.sha256(M + b'verify', encoder=nacl.encoding.RawEncoder)
        S = nacl.hash.sha256(M + b'secret', encoder=nacl.encoding.RawEncoder)
        if Acheck != authtag:
            raise ValueError("NTor auth mismatch")

        if sphinx_domainsep is not None:
            blindkey = Sphinx.makeblindkey(S, sphinx_domainsep,
                    self.perfstats)
            self.blinding_keys.append(blindkey)
        return S


class VanillaExtendCircuitHandler:
    """A handler for VanillaExtendCircuitCell cells.  It allocates a new
    circuit id on the Channel to the requested next hop, connects the
    existing and new circuits together, and forwards a
    VanillaCreateCircuitMsg to the next hop."""

    def received_cell(self, circhandler, cell):
        # Remove ourselves from handling a second
        # VanillaExtendCircuitCell on this circuit
        circhandler.replace_celltype_handler(VanillaExtendCircuitCell, None)

        # Allocate a new circuit id to the requested next hop
        channelmgr = circhandler.channel.channelmgr
        nexthopchannel = channelmgr.get_channel_to(cell.hopaddr)
        newcircid, newcirchandler = nexthopchannel.new_circuit()

        # Connect the existing and new circuits together
        circhandler.adjacent_circuit_handler = newcirchandler
        newcirchandler.adjacent_circuit_handler = circhandler

        # Set up a handler for when the VanillaCreatedCircuitCell comes
        # back
        newcirchandler.replace_celltype_handler(
                VanillaCreatedCircuitCell,
                VanillaCreatedRelayHandler())

        # Forward a VanillaCreateCircuitMsg to the next hop
        nexthopchannel.send_msg(
                VanillaCreateCircuitMsg(newcircid, cell.ntor_request))

class TelescopingExtendCircuitHandler:
    """A handler for TelescopingExtendCircuitCell cells.  It allocates a new
    circuit id on the Channel to the requested next hop, connects the
    existing and new circuits together, and forwards a
    TelescopingCreateCircuitMsg to the next hop."""

    def __init__(self, relaypicker, current_relay_idkey):
        self.relaypicker = relaypicker
        self.current_relay_idkey = current_relay_idkey

    def received_cell(self, circhandler, cell):
        # Remove ourselves from handling a second
        # TelescopingExtendCircuitCell on this circuit
        circhandler.replace_celltype_handler(TelescopingExtendCircuitCell, None)

        # Find the SNIP corresponding to the  index sent by the client
        next_snip = self.relaypicker.pick_relay_by_uniform_index(cell.idx)

        # Check to make sure that we aren't extending to ourselves. If we are,
        # close the circuit.
        if next_snip.snipdict["idkey"] == self.current_relay_idkey:
            print("ERR: Client requested extending the circuit to a relay already in the path; aborting. my circid: " + str(circhandler.circid))
            circhandler.close()
            return


        # Allocate a new circuit id to the requested next hop
        channelmgr = circhandler.channel.channelmgr
        nexthopchannel = channelmgr.get_channel_to(next_snip.snipdict["addr"])
        newcircid, newcirchandler = nexthopchannel.new_circuit()

        # Connect the existing and new circuits together
        circhandler.adjacent_circuit_handler = newcirchandler
        newcirchandler.adjacent_circuit_handler = circhandler


        # Set up a handler for when the TelescopingCreatedCircuitCell comes
        # back
        newcirchandler.replace_celltype_handler(
                TelescopingCreatedCircuitCell,
                TelescopingCreatedRelayHandler(next_snip))

        # Forward a TelescopingCreateCircuitMsg to the next hop
        nexthopchannel.send_msg(
                TelescopingCreateCircuitMsg(newcircid, cell.ntor_request))

class VanillaCreatedRelayHandler:
    """Handle a VanillaCreatedCircuitCell received by a _relay_ that
    recently received a VanillaExtendCircuitCell from a client, and so
    forwarded a VanillaCreateCircuitCell to the next hop."""

    def received_cell(self, circhandler, cell):
        # Remove ourselves from handling a second
        # VanillaCreatedCircuitCell on this circuit
        circhandler.replace_celltype_handler(VanillaCreatedCircuitCell, None)

        # Just forward a VanillaExtendedCircuitCell back towards the
        # client
        circhandler.adjacent_circuit_handler.send_cell(
            VanillaExtendedCircuitCell(cell.ntor_reply))

class TelescopingCreatedRelayHandler:
    """Handle a TelescopingCreatedCircuitCell received by a _relay_ that
    recently received a TelescopingExtendCircuitCell from a client, and so
    forwarded a TelescopingCreateCircuitCell to the next hop."""

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

    def received_cell(self, circhandler, cell):
        print("LOG: Handle a TelescopingCreatedCircui received by a relay")
        # Remove ourselves from handling a second
        # VanillaCreatedCircuitCell on this circuit
        circhandler.replace_celltype_handler(TelescopingCreatedCircuitCell, None)

        # Just forward a TelescopingExtendedCircuitCell back towards the
        # client
        circhandler.adjacent_circuit_handler.send_cell(
            TelescopingExtendedCircuitCell(cell.ntor_reply, self.next_snip))

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

    class NoCryptLayer:
        def encrypt_msg(self, msg): return msg
        def decrypt_msg(self, msg): return msg

    class CryptLayer:
        def __init__(self, enckey, deckey, next_layer):
            self.enckey = enckey
            self.deckey = deckey
            self.next_layer = next_layer
        def encrypt_msg(self, msg):
            return self.next_layer.encrypt_msg(EncryptedCell(self.enckey, msg))
        def decrypt_msg(self, msg):
            return self.next_layer.decrypt_msg(msg).decrypt(self.deckey)

    def __init__(self, channel, circid):
        self.channel = channel
        self.circid = circid
        # The list of relay descriptors that form the circuit so far
        # (client side only)
        self.circuit_descs = []
        # The dispatch table is indexed by type, and the values are
        # objects with received_cell(circhandler, cell) methods.
        self.cell_dispatch_table = dict()
        # The topmost crypt layer.  This is an object with
        # encrypt_msg(msg) and decrypt_msg(msg) methods that returns the
        # en/decrypted messages respectively.  Messages are encrypted
        # starting with the last layer that was added (the keys for the
        # furthest away relay in the circuit) and are decrypted starting
        # with the first layer that was added (the keys for the guard).
        self.crypt_layer = self.NoCryptLayer()
        # The adjacent CircuitHandler that's connected to this one.  If
        # we get a cell on one, we forward it to the other (if it's not
        # meant for us to handle directly)
        self.adjacent_circuit_handler = None
        # The function to call when this circuit closes
        self.closer = lambda: self.channel.circuithandlers.pop(circid)

    def close(self):
        """Close the circuit.  Sends a CloseCell on the circuit (and its
        adjacent circuit, if present) and closes both."""
        adjcirchandler = self.adjacent_circuit_handler
        self.adjacent_circuit_handler = None
        print("Log: CLosing circuit. circid: " + str(self.circid))
        if adjcirchandler is not None:
            adjcirchandler.adjacent_circuit_handler = None
        self.closer()
        self.channel_send_cell(CloseCell())
        if adjcirchandler is not None:
            adjcirchandler.closer()
            adjcirchandler.channel_send_cell(CloseCell())

    def send_cell(self, cell):
        """Send a cell on this circuit, encrypting as needed."""
        self.channel_send_cell(self.crypt_layer.encrypt_msg(cell))

    def channel_send_cell(self, cell):
        """Send a cell on this circuit directly without encrypting it
        first."""
        self.channel.send_msg(CircuitCellMsg(self.circid, cell))

    def received_cell(self, cell):
        """A cell has been received on this circuit.  Dispatch it
        according to its type."""
        if isinstance(cell, EncryptedCell):
            cell = self.crypt_layer.decrypt_msg(cell)

        print("CircuitHandler: %s received cell %s on circuit %d from %s" % (self.channel.channelmgr.myaddr, cell, self.circid, self.channel.peer.channelmgr.myaddr))

        # If it's still encrypted, it's for sure meant for forwarding to
        # our adjacent hop, which had better exist.
        if isinstance(cell, EncryptedCell):
            self.adjacent_circuit_handler.send_cell(cell)
        else:
            # This is a plaintext cell meant for us.  Handle it
            # according to the table.
            celltype = type(cell)
            if celltype in self.cell_dispatch_table:
                self.cell_dispatch_table[celltype].received_cell(self, cell)

            elif isinstance(cell, StringCell):
                # Default handler; just print the message in the cell
                print("CircuitHandler: %s received '%s' on circuit %d from %s" \
                        % (self.channel.channelmgr.myaddr, cell,
                            self.circid, self.channel.peer.channelmgr.myaddr))

            elif isinstance(cell, CloseCell):
                print("WARNING: Received CloseCell on circuit " + str(self.circid))
                # Forward the CloseCell (without encryption) to the
                # adjacent circuit, if any, and close both this and the
                # adjacent circuit
                adjcirchandler = self.adjacent_circuit_handler
                self.adjacent_circuit_handler = None
                if adjcirchandler is not None:
                    adjcirchandler.adjacent_circuit_handler = None
                self.closer()
                if adjcirchandler is not None:
                    adjcirchandler.closer()
                    adjcirchandler.channel_send_cell(cell)

            else:
                # I don't know how to handle this cell?
                raise ValueError("CircuitHandler: %s received unknown cell type %s on circuit %d from %s" \
                        % (self.channel.channelmgr.myaddr, cell,
                            self.circid, self.channel.peer.channelmgr.myaddr))

    def replace_celltype_handler(self, celltype, handler):
        """Add an object with a received_cell(circhandler, cell) method
        to the cell dispatch table.  It replaces anything that's already
        there.  Passing None as the handler removes the dispatcher for
        that cell type."""
        if handler is None:
            del self.cell_dispatch_table[celltype]
        else:
            self.cell_dispatch_table[celltype] = handler

    def add_crypt_layer(self, enckey, deckey):
        """Add a processing layer to this CircuitHandler so that cells
        we send will get encrypted with the first given key, and cells
        we receive will be decrypted with the other given key."""
        current_crypt_layer = self.crypt_layer
        self.crypt_layer = self.CryptLayer(enckey, deckey, current_crypt_layer)


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 RelayChannelManager managing this Channel
        self.channelmgr = None
        # The Channel at the other end
        self.peer = None
        # The function to call when the connection closes
        self.closer = lambda: None
        # 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):
        # Close each circuithandler we're managing
        while self.circuithandlers:
            chitems = iter(self.circuithandlers.items())
            circid, circhandler = next(chitems)
            print('closing circuit', circid)
            circhandler.close()
        self.closer()
        self.peer = None

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

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

    def is_circuit_open(self, circid):
        is_open = (circid in self.circuithandlers) and (self.circuithandlers[circid] is not None)
        return is_open

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

    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.channelmgr.perfstats.bytes_sent += msg.size()
        self.peer.received(self.channelmgr.myaddr, msg)

    def received(self, peeraddr, msg):
        """Callback when a message is received from the network."""
        print('Channel: %s received %s from %s' % (self.channelmgr.myaddr, msg, peeraddr))
        self.channelmgr.perfstats.bytes_received += msg.size()
        if isinstance(msg, CircuitCellMsg):
            circid, cell = msg.circid, msg.cell
            self.circuithandlers[circid].received_cell(cell)
        else:
            self.channelmgr.received_msg(msg, peeraddr, self)


class ChannelManager:
    """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, perfstats):
        # A dictionary of Channels to other hosts, indexed by NetAddr
        self.channels = dict()
        self.myaddr = myaddr
        self.dirauthaddrs = dirauthaddrs
        self.consensus = None
        self.relaypicker = None
        self.perfstats = perfstats

    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.channelmgr = 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
        print('getting channel from',self.myaddr,'to',addr)
        newchannel = network.thenetwork.connect(self.myaddr, addr,
                self.perfstats)
        print('got channel from',self.myaddr,'to',addr)
        self.channels[addr] = newchannel
        newchannel.closer = lambda: self.channels.pop(addr)
        newchannel.channelmgr = self

        return newchannel

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

    def received_cell(self, circid, cell, peeraddr, channel):
        """Callback with a circuit-specific cell is received."""
        print("ChannelManager: 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 RelayChannelManager(ChannelManager):
    """The subclass of ChannelManager for relays."""

    def __init__(self, myaddr, dirauthaddrs, onionprivkey, idpubkey, perfstats):
        super().__init__(myaddr, dirauthaddrs, perfstats)
        self.onionkey = onionprivkey
        self.idpubkey = idpubkey
        if network.thenetwork.womode != network.WOMode.VANILLA:
            self.endive = None

    def get_consensus(self):
        """Download a fresh consensus (and ENDIVE if using Walking
        Onions) from a random dirauth."""
        a = random.choice(self.dirauthaddrs)
        c = network.thenetwork.connect(self, a, self.perfstats)
        if network.thenetwork.womode == network.WOMode.VANILLA:
            if self.consensus is not None and \
                    len(self.consensus.consdict['relays']) > 0:
                self.consensus = c.getconsensusdiff()
            else:
                self.consensus = c.getconsensus()
            self.relaypicker = dirauth.Consensus.verify(self.consensus,
                    network.thenetwork.dirauthkeys(), self.perfstats)
        else:
            self.consensus = c.getconsensus()
            if self.endive is not None and \
                    len(self.endive.enddict['snips']) > 0:
                self.endive = c.getendivediff()
            else:
                self.endive = c.getendive()
            self.relaypicker = dirauth.ENDIVE.verify(self.endive,
                    self.consensus, network.thenetwork.dirauthkeys(),
                    self.perfstats)
        c.close()

    def received_msg(self, msg, peeraddr, channel):
        """Callback when a NetMsg not specific to a circuit is
        received."""
        print("RelayChannelManager: 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 = self.relaypicker.pick_weighted_relay()
                if network.thenetwork.womode == network.WOMode.VANILLA:
                    nextaddr = nexthop.descdict["addr"]
                else:
                    nextaddr = nexthop.snipdict["addr"]
                self.send_msg(RelayRandomHopMsg(msg.ttl-1), nextaddr)
        elif isinstance(msg, RelayGetConsensusMsg):
            self.send_msg(RelayConsensusMsg(self.consensus), peeraddr)
        elif isinstance(msg, RelayGetConsensusDiffMsg):
            self.send_msg(RelayConsensusDiffMsg(self.consensus), peeraddr)
        elif isinstance(msg, VanillaCreateCircuitMsg):
            # A new circuit has arrived
            circhandler = channel.new_circuit_with_circid(msg.circid)
            # Create the ntor reply
            reply, secret = NTor.reply(self.onionkey, self.idpubkey,
                    msg.ntor_request, self.perfstats)
            # Set up the circuit to use the shared secret
            enckey = nacl.hash.sha256(secret + b'downstream')
            deckey = nacl.hash.sha256(secret + b'upstream')
            circhandler.add_crypt_layer(enckey, deckey)
            # Add a handler for if an Extend Cell arrives (there should
            # be at most one on this circuit).
            circhandler.replace_celltype_handler(
                    VanillaExtendCircuitCell, VanillaExtendCircuitHandler())
            # Send the ntor reply
            self.send_msg(CircuitCellMsg(msg.circid,
                    VanillaCreatedCircuitCell(reply)), peeraddr)
        elif isinstance(msg, TelescopingCreateCircuitMsg):
            # A new circuit has arrived
            circhandler = channel.new_circuit_with_circid(msg.circid)
            # Create the ntor reply
            reply, secret = NTor.reply(self.onionkey, self.idpubkey,
                    msg.ntor_request, self.perfstats)
            # Set up the circuit to use the shared secret
            enckey = nacl.hash.sha256(secret + b'downstream')
            deckey = nacl.hash.sha256(secret + b'upstream')
            circhandler.add_crypt_layer(enckey, deckey)
            # Add a handler for if an Extend Cell arrives (there should
            # be at most one on this circuit).
            circhandler.replace_celltype_handler(
                    TelescopingExtendCircuitCell,
                    TelescopingExtendCircuitHandler(self.relaypicker,
                        self.idpubkey))

            # Send the ntor reply
            self.send_msg(CircuitCellMsg(msg.circid,
                    TelescopingCreatedCircuitCell(reply)), peeraddr)
        else:
            return super().received_msg(msg, peeraddr, channel)

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


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

    def __init__(self, dirauthaddrs, bw, flags):
        # Gather performance statistics
        self.perfstats = network.PerfStats(network.EntType.RELAY)
        self.perfstats.is_bootstrapping = True

        # Create the identity and onion keys
        self.idkey = nacl.signing.SigningKey.generate()
        self.onionkey = nacl.public.PrivateKey.generate()
        self.perfstats.keygens += 2
        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)
        self.perfstats.name = "Relay at %s" % self.netaddr

        # 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 RelayChannelManager connection manager
        self.channelmgr = RelayChannelManager(self.netaddr, dirauthaddrs,
                self.onionkey, self.idkey.verify_key, self.perfstats)

        # 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.channelmgr.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.channelmgr.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, self.perfstats)
        dirauth.RelayDescriptor.verify(desc, self.perfstats)

        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.channelmgr.dirauthaddrs:
            c = network.thenetwork.connect(self, a, self.perfstats)
            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 RelayChannelManager
        self.channelmgr.add_channel(ourchannel, peer)

        return peerchannel

if __name__ == '__main__':
    perfstats = network.PerfStats(network.EntType.NONE)

    network.thenetwork.set_wo_style(network.WOMode.TELESCOPING,
            network.SNIPAuthMode.MERKLE)

    # Start some dirauths
    numdirauths = 9
    dirauthaddrs = []
    for i in range(numdirauths):
        dira = dirauth.DirAuth(i, numdirauths)
        dirauthaddrs.append(dira.netaddr)

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

    print(dirauth.DirAuth.consensus)
    print(dirauth.DirAuth.endive)

    if network.thenetwork.womode == network.WOMode.VANILLA:
        relaypicker = dirauth.Consensus.verify(dirauth.DirAuth.consensus,
                network.thenetwork.dirauthkeys(), perfstats)
    else:
        relaypicker = dirauth.ENDIVE.verify(dirauth.DirAuth.endive,
                dirauth.DirAuth.consensus,
                network.thenetwork.dirauthkeys(), perfstats)
        if network.thenetwork.snipauthmode == \
                network.SNIPAuthMode.THRESHSIG:
            for s in dirauth.DirAuth.endive.enddict['snips']:
                dirauth.SNIP.verify(s, dirauth.DirAuth.consensus,
                        network.thenetwork.dirauthkeys()[0], perfstats)

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

    relays[3].channelmgr.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.channelmgr.channels.keys()]))
        raddr = r.netaddr
        for ad, ch in r.channelmgr.channels.items():
            if ch.peer.channelmgr.myaddr != ad:
                print('address mismatch:', raddr, ad, ch.peer.channelmgr.myaddr)

            if ch.peer.channelmgr.channels[raddr].peer is not ch:
                print('asymmetry:', raddr, ad, ch, ch.peer.channelmgr.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)
    print(dirauth.DirAuth.endive)

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

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

    channel = relays[3].channelmgr.get_channel_to(relays[5].netaddr)
    circid, circhandler = channel.new_circuit()
    peerchannel = relays[5].channelmgr.get_channel_to(relays[3].netaddr)
    peerchannel.new_circuit_with_circid(circid)
    relays[3].channelmgr.send_cell(circid, StringCell("test"), relays[5].netaddr)

    idpubkey = relays[1].idkey.verify_key
    onionpubkey = relays[1].onionkey.public_key
    nt = NTor(perfstats)
    req = nt.request()
    R, S = NTor.reply(relays[1].onionkey, idpubkey, req, perfstats)
    S2 = nt.verify(R, onionpubkey, idpubkey)
    print(S == S2)

    # Test the Sphinx class: DH version (for the path selection keys)

    server1_key = nacl.public.PrivateKey.generate()
    server2_key = nacl.public.PrivateKey.generate()
    server3_key = nacl.public.PrivateKey.generate()
    client_key = nacl.public.PrivateKey.generate()

    # Check that normal DH is working
    ckey = nacl.public.Box(client_key, server1_key.public_key).shared_key()
    skey = nacl.public.Box(server1_key, client_key.public_key).shared_key()
    assert(ckey == skey)

    # Transform the client pubkey with Sphinx as it passes through the
    # servers and record the resulting shared secrets
    blinded_client_pubkey = client_key.public_key
    s1secret, blinded_client_pubkey = Sphinx.server(blinded_client_pubkey,
        server1_key, b'circuit', False, perfstats)
    s2secret, blinded_client_pubkey = Sphinx.server(blinded_client_pubkey,
        server2_key, b'circuit', False, perfstats)
    s3secret, _ = Sphinx.server(blinded_client_pubkey,
        server3_key, b'circuit', True, perfstats)

    # Hopefully matching keys on the client side
    blinding_keys = []
    c1secret, blind_key = Sphinx.client(client_key, blinding_keys,
        server1_key.public_key, b'circuit', False, perfstats)
    blinding_keys.append(blind_key)
    c2secret, blind_key = Sphinx.client(client_key, blinding_keys,
        server2_key.public_key, b'circuit', False, perfstats)
    blinding_keys.append(blind_key)
    c3secret, _ = Sphinx.client(client_key, blinding_keys,
        server3_key.public_key, b'circuit', True, perfstats)

    assert(s1secret == c1secret)
    assert(s2secret == c2secret)
    assert(s3secret == c3secret)
    print('Sphinx DH test successful')

    # End test of Sphinx (DH version)

    # Test the Sphinx class: NTor version (for the path selection keys)

    server1_idkey = nacl.signing.SigningKey.generate().verify_key
    server2_idkey = nacl.signing.SigningKey.generate().verify_key
    server3_idkey = nacl.signing.SigningKey.generate().verify_key
    server1_onionkey = nacl.public.PrivateKey.generate()
    server2_onionkey = nacl.public.PrivateKey.generate()
    server3_onionkey = nacl.public.PrivateKey.generate()
    client_ntor = NTor(perfstats)

    # Client's initial message
    client_pubkey = client_ntor.request()

    # Transform the client pubkey with Sphinx as it passes through the
    # servers and record the resulting shared secrets
    blinded_client_pubkey = client_pubkey
    (s1reply, s1secret), blinded_client_pubkey = NTor.reply(
            server1_onionkey, server1_idkey, blinded_client_pubkey,
            perfstats, b'data')
    (s2reply, s2secret), blinded_client_pubkey = NTor.reply(
            server2_onionkey, server2_idkey, blinded_client_pubkey,
            perfstats, b'data')
    (s3reply, s3secret) = NTor.reply(
            server3_onionkey, server3_idkey, blinded_client_pubkey,
            perfstats)

    # Hopefully matching keys on the client side
    c1secret = client_ntor.verify(s1reply, server1_onionkey.public_key,
            server1_idkey, b'data')
    c2secret = client_ntor.verify(s2reply, server2_onionkey.public_key,
            server2_idkey, b'data')
    c3secret = client_ntor.verify(s3reply, server3_onionkey.public_key,
            server3_idkey)

    assert(s1secret == c1secret)
    assert(s2secret == c2secret)
    assert(s3secret == c3secret)
    print('Sphinx NTor test successful')

    # End test of Sphinx (NTor version)