walkingonions/client.py

#!/usr/bin/env python3

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

import network
import dirauth
import relay

import nacl.hash

class VanillaCreatedExtendedHandler:
    """A handler for VanillaCreatedCircuitCell and
    VanillaExtendedCircuitCell cells."""

    def __init__(self, channelmgr, ntor, expecteddesc):
        self.channelmgr = channelmgr
        self.ntor = ntor
        self.expecteddesc = expecteddesc
        self.onionkey = expecteddesc.descdict['onionkey']
        self.idkey = expecteddesc.descdict['idkey']

    def received_cell(self, circhandler, cell):
        secret = self.ntor.verify(cell.ntor_reply, self.onionkey, self.idkey)
        enckey = nacl.hash.sha256(secret + b'upstream')
        deckey = nacl.hash.sha256(secret + b'downstream')
        circhandler.add_crypt_layer(enckey, deckey)
        if len(circhandler.circuit_descs) == 0:
            # This was a VanillaCreatedCircuitCell
            circhandler.replace_celltype_handler(
                    relay.VanillaCreatedCircuitCell, None)
        else:
            # This was a VanillaExtendedCircuitCell
            circhandler.replace_celltype_handler(
                    relay.VanillaExtendedCircuitCell, None)
        circhandler.circuit_descs.append(self.expecteddesc)

        # Are we done building the circuit?
        if len(circhandler.circuit_descs) == 3:
            # Yes!
            return

        nexthop = None
        while nexthop is None:
            nexthop = self.channelmgr.relaypicker.pick_weighted_relay()
            if nexthop.descdict['addr'] in \
                    [ desc.descdict['addr'] \
                        for desc in circhandler.circuit_descs ]:
                nexthop = None

        # Construct the VanillaExtendCircuitCell
        ntor = relay.NTor(self.channelmgr.perfstats)
        ntor_request = ntor.request()
        circextendmsg = relay.VanillaExtendCircuitCell(
                nexthop.descdict['addr'], ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.VanillaExtendedCircuitCell,
                VanillaCreatedExtendedHandler(self.channelmgr, ntor, nexthop))

        # Send the cell
        circhandler.send_cell(circextendmsg)

class TelescopingCreatedHandler:
    """A handler for TelescopingCreatedCircuitCell cells; this will only always
    communicate with the client's guard."""

    def __init__(self, channelmgr, ntor):
        self.channelmgr = channelmgr
        self.ntor = ntor
        self.onionkey = self.channelmgr.guard.snipdict['onionkey']
        self.idkey = self.channelmgr.guard.snipdict['idkey']

    def received_cell(self, circhandler, cell):
        logging.debug("Received cell in TelescopingCreatedHandler")

        secret = self.ntor.verify(cell.ntor_reply, self.onionkey, self.idkey)
        enckey = nacl.hash.sha256(secret + b'upstream')
        deckey = nacl.hash.sha256(secret + b'downstream')
        circhandler.add_crypt_layer(enckey, deckey)

        circhandler.replace_celltype_handler(relay.TelescopingCreatedCircuitCell, None)

        nexthopidx = None
        while nexthopidx is None:
            nexthopidx = self.channelmgr.relaypicker.pick_weighted_relay_index()
            #print("WARNING: Unimplemented! Need to check if this idx is in the list of circhandlers idxs")
            # TODO verify we don't need to do the above

        # Construct the TelescopingExtendCircuitCell
        ntor = relay.NTor(self.channelmgr.perfstats)
        ntor_request = ntor.request()
        circextendmsg = relay.TelescopingExtendCircuitCell(
                nexthopidx, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.TelescopingExtendedCircuitCell,
                TelescopingExtendedHandler(self.channelmgr, ntor))

        # Send the cell
        circhandler.send_cell(circextendmsg)


class TelescopingExtendedHandler:
    """A handler for TelescopingExtendedCircuitCell cells."""

    def __init__(self, channelmgr, ntor):
        self.channelmgr = channelmgr
        self.ntor = ntor

    def received_cell(self, circhandler, cell):
        logging.debug("Received cell in TelescopingExtendedHandler")

        # Validate the SNIP
        dirauth.SNIP.verify(cell.snip, self.channelmgr.consensus,
                network.thenetwork.dirauthkeys()[0],
                self.channelmgr.perfstats)

        onionkey = cell.snip.snipdict['onionkey']
        idkey = cell.snip.snipdict['idkey']

        secret = self.ntor.verify(cell.ntor_reply, onionkey, idkey)
        enckey = nacl.hash.sha256(secret + b'upstream')
        deckey = nacl.hash.sha256(secret + b'downstream')
        circhandler.add_crypt_layer(enckey, deckey)

        circhandler.replace_celltype_handler(
                    relay.TelescopingExtendedCircuitCell, None)
        circhandler.circuit_descs.append(cell.snip)

        # Are we done building the circuit?
        logging.warning("we may need another circhandler structure for snips")
        if len(circhandler.circuit_descs) == 3:
            logging.debug("Circuit is long enough; exiting.")
            # Yes!
            return

        nexthopidx = None
        while nexthopidx is None:
            # Relays make sure that when the extend to a relay, they are not
            # extending to themselves. So here, we just need to make sure that
            # this ID is not the same as the guard ID, to protect against the
            # guard and exit being the same relay
            print("WARNING: Unimplemented! Need to check that sampled ID is not the same as the guard.")
            nexthopidx = self.channelmgr.relaypicker.pick_weighted_relay_index()
            logging.warning("Unimplemented! Need to check if this idx is in the list of circhandlers idxs")

        # Construct the VanillaExtendCircuitCell
        ntor = relay.NTor(self.channelmgr.perfstats)
        ntor_request = ntor.request()
        circextendmsg = relay.TelescopingExtendCircuitCell(
                nexthopidx, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.TelescopingExtendedCircuitCell,
                TelescopingExtendedHandler(self.channelmgr, ntor))

        # Send the cell
        circhandler.send_cell(circextendmsg)

class SinglePassCreatedHandler:
    """A handler for SinglePassCreatedCircuitCell cells."""

    def __init__(self, channelmgr, ntor, client_key):
        self.channelmgr = channelmgr
        self.ntor = ntor
        self.client_key = client_key

    def received_cell(self, circhandler, cell):
        logging.debug("Received cell in SinglePassCreatedHandler")
        sys.exit("not yet implemented")


class ClientChannelManager(relay.ChannelManager):
    """The subclass of ChannelManager for clients."""

    def __init__(self, myaddr, dirauthaddrs, perfstats):
        super().__init__(myaddr, dirauthaddrs, perfstats)
        self.guardaddr = None
        self.guard = None

    def get_consensus_from_fallbackrelay(self):
        """Download a fresh consensus from a random fallbackrelay."""
        fb = random.choice(network.thenetwork.getfallbackrelays())
        if network.thenetwork.womode == network.WOMode.VANILLA:
            if self.consensus is not None and \
                    len(self.consensus.consdict['relays']) > 0:
                self.send_msg(relay.RelayGetConsensusDiffMsg(), fb.netaddr)
            else:
                self.send_msg(relay.RelayGetConsensusMsg(), fb.netaddr)
        else:
            self.send_msg(relay.RelayGetConsensusMsg(), fb.netaddr)

    def ensure_guard_vanilla(self):
        """Ensure that we have a channel to a guard (Vanilla Onion
        Routing version)."""
        while True:
            if self.guardaddr is None:
                # Pick a guard from the consensus
                self.guard = self.relaypicker.pick_weighted_relay()
                self.guardaddr = self.guard.descdict['addr']

            self.test_guard_connection()

            if self.guardaddr is not None:
                break

        logging.debug('chose guard=%s', self.guardaddr)

    def test_guard_connection(self):
            # Connect to the guard
            try:
                self.get_channel_to(self.guardaddr)
            except network.NetNoServer:
                # Our guard is gone
                self.guardaddr = None
                self.guard = None

    def ensure_guard_walking_onions(self):
        """Ensure we have a channel to a guard (Walking Onions version).

        For the first implementation, we assume an out-of-band mechanism
        that just simply hands us a guard; we don't count the number of
        operations or bandwidth as this operation in practice occurs
        infrequently."""

        while True:
            if self.guardaddr is None:
                #randomly-sample a guard
                logging.warning("Unimplemented! guard should be selected from any relays.")
                self.guard = self.relaypicker.pick_weighted_relay()
                # here, we have a SNIP instead of a relay descriptor
                self.guardaddr = self.guard.snipdict['addr']

            self.test_guard_connection()

            if self.guardaddr is not None:
                break

        logging.debug('chose guard=%s', self.guardaddr)

    def ensure_guard(self):
        """Ensure that we have a channel to a guard."""
        if network.thenetwork.womode == network.WOMode.VANILLA:
            self.ensure_guard_vanilla()
            return
        # At this point, we are either in Telescoping or Single-Pass mode

        self.ensure_guard_walking_onions()

    def new_circuit_vanilla(self):
        """Create a new circuit from this client. (Vanilla Onion Routing
        version)"""

        # Get our channel to the guard
        guardchannel = self.get_channel_to(self.guardaddr)

        # Allocate a new circuit id on it
        circid, circhandler = guardchannel.new_circuit()

        # Construct the VanillaCreateCircuitMsg
        ntor = relay.NTor(self.perfstats)
        ntor_request = ntor.request()
        circcreatemsg = relay.VanillaCreateCircuitMsg(circid, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.VanillaCreatedCircuitCell,
                VanillaCreatedExtendedHandler(self, ntor, self.guard))

        # Send the message
        guardchannel.send_msg(circcreatemsg)

        return circhandler

    def new_circuit_telescoping(self):
        """Create a new circuit from this client. (Telescoping Walking Onions
        version). If an error occurs and the circuit is deleted from the guard
        channel, return None, otherwise, return the circuit handler."""

        # Get our channel to the guard
        guardchannel = self.get_channel_to(self.guardaddr)

        # Allocate a new circuit id on it
        circid, circhandler = guardchannel.new_circuit()

        # Construct the TelescopingCreateCircuitMsg
        ntor = relay.NTor(self.perfstats)
        ntor_request = ntor.request()
        circcreatemsg = relay.TelescopingCreateCircuitMsg(circid, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.TelescopingCreatedCircuitCell,
                TelescopingCreatedHandler(self, ntor))

        # Send the message
        guardchannel.send_msg(circcreatemsg)

        # Check to make sure the circuit is open before sending it- if there
        # was an error when establishing it, the circuit could already be
        # closed.
        if not guardchannel.is_circuit_open(circid):
            logging.debug("Circuit was already closed, not sending bytes. circid: " + str(circid))
            return None

        return circhandler

    def new_circuit_singlepass(self):
        """Create a new circuit from this client. (Single-Pass Walking Onions
        version). If an error occurs and the circuit is deleted from the guard
        channel, return None, otherwise, return the circuit handler."""

        # Get our channel to the guard
        guardchannel = self.get_channel_to(self.guardaddr)

        # Allocate a new circuit id on it
        circid, circhandler = guardchannel.new_circuit()

        # first, create the path-selection key used for Sphinx
        client_key = nacl.public.PrivateKey.generate()

        # Construct the SinglePassCreateCircuitMsg
        ntor = relay.NTor(self.perfstats)
        ntor_request = ntor.request()
        ttl = 2 # TODO set a default for the msg type
        circcreatemsg = relay.SinglePassCreateCircuitMsg(circid, ntor_request,
                client_key.public_key, ttl)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                relay.SinglePassCreatedCircuitCell,
                SinglePassCreatedHandler(self, ntor, client_key))

        # Send the message
        guardchannel.send_msg(circcreatemsg)

        # Check to make sure the circuit is open before sending it- if there
        # was an error when establishing it, the circuit could already be
        # closed.
        if not guardchannel.is_circuit_open(circid):
            logging.debug("Circuit was already closed, not sending bytes. circid: " + str(circid))
            return None

        return circhandler


    def new_circuit(self):
        """Create a new circuit from this client."""
        circhandler = None

        # If an error occured, circhandler will still be None, so we should
        # try again.
        while circhandler is None:
            if network.thenetwork.womode == network.WOMode.VANILLA:
                circhandler = self.new_circuit_vanilla()
            elif network.thenetwork.womode == network.WOMode.TELESCOPING:
                circhandler = self.new_circuit_telescoping()
            elif network.thenetwork.womode == network.WOMode.SINGLEPASS:
                circhandler = self.new_circuit_singlepass()

        return circhandler

    def received_msg(self, msg, peeraddr, channel):
        """Callback when a NetMsg not specific to a circuit is
        received."""
        logging.debug("Client %s received msg %s from %s" % (self.myaddr, msg, peeraddr))
        if isinstance(msg, relay.RelayConsensusMsg) or \
                isinstance(msg, relay.RelayConsensusDiffMsg):
            self.relaypicker = dirauth.Consensus.verify(msg.consensus,
                    network.thenetwork.dirauthkeys(), self.perfstats)
            self.consensus = msg.consensus
        else:
            return super().received_msg(msg, peeraddr, channel)

    def received_cell(self, circid, cell, peeraddr, channel):
        """Callback with a circuit-specific cell is received."""
        logging.debug("Client %s received cell on circ %d: %s from %s" % (self.myaddr, circid, cell, peeraddr))
        if isinstance(msg, relay.CloseCell):
            logging.debug("Log: Client received close cell; closing circuit")
            # TODO close cell
        return super().received_cell(circid, cell, peeraddr, channel)


class Client:
    """A class representing a Tor client."""

    def __init__(self, dirauthaddrs):
        # Get a network address for client-side use only (do not bind it
        # to the network)
        self.netaddr = network.NetAddr()
        self.perfstats = network.PerfStats(network.EntType.CLIENT)
        self.perfstats.name = "Client at %s" % self.netaddr
        self.perfstats.is_bootstrapping = True
        self.channelmgr = ClientChannelManager(self.netaddr, dirauthaddrs,
                self.perfstats)

        # Register for epoch tick notifications
        network.thenetwork.wantepochticks(self, True)

    def terminate(self):
        """Quit this client."""

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

        # Close relay connections
        self.channelmgr.terminate()

    def get_consensus(self):
        """Fetch a new consensus."""

        # We're going to want a new consensus from our guard.  In order
        # to get that, we'll need a channel to our guard.  In order to
        # get that, we'll need a guard address.  In order to get that,
        # we'll need a consensus (uh, oh; in that case, fetch the
        # consensus from a fallback relay).


        guardaddr = self.channelmgr.guardaddr
        guardchannel = None
        if guardaddr is not None:
            try:
                guardchannel = self.channelmgr.get_channel_to(guardaddr)
            except network.NetNoServer:
                guardaddr = None

        if guardchannel is None:
            logging.debug("In bootstrapping mode")
            self.channelmgr.get_consensus_from_fallbackrelay()
            logging.debug('client consensus=%s', self.channelmgr.consensus)
            return

        if network.thenetwork.womode == network.WOMode.VANILLA:
            if self.channelmgr.consensus is not None and len(self.channelmgr.consensus.consdict['relays']) > 0:
                guardchannel.send_msg(relay.RelayGetConsensusDiffMsg())
                logging.debug('got consensus diff, client consensus=%s', self.channelmgr.consensus)
                return

        # At this point, we are in one of the following scenarios:
            # 1. This is a walking onions protocol, and the client fetches the
            #    complete consensus each epoch
            # 2. This is Vanilla Onion Routing and the client doesn't have a
            #    consensus and needs to bootstrap it.

        guardchannel.send_msg(relay.RelayGetConsensusMsg())
        logging.debug('client consensus=%s', self.channelmgr.consensus)

    def newepoch(self, epoch):
        """Callback that fires at the start of each epoch"""

        # We'll need a new consensus
        self.get_consensus()

        # If we don't have a guard, pick one and make a channel to it
        self.channelmgr.ensure_guard()

if __name__ == '__main__':
    perfstats = network.PerfStats(network.EntType.NONE)
    totsent = 0
    totrecv = 0
    dirasent = 0
    dirarecv = 0
    relaysent = 0
    relayrecv = 0
    clisent = 0
    clirecv = 0

    if len(sys.argv) < 3:
        print("Must pass in network mode and snip auth mode!")
        print("Network options are vanilla, telescoping, or single-pass.")
        print("SNIP auth options are merkle or threshold.")
        sys.exit(0)

    logging.basicConfig(level=logging.DEBUG)

    network_mode = network.WOMode.string_to_type(sys.argv[1])
    if network_mode == -1:
        print("Not a valid network mode: " + network_mode)
        sys.exit(0)

    snipauth_mode = network.SNIPAuthMode.string_to_type(sys.argv[2])
    if network_mode == -1:
        print("Not a valid SNIP authentication mode: " + snipauth_mode)
        sys.exit(0)

    # Initialize the (non-cryptographic) random seed
    random.seed(1)

    if network_mode == network.WOMode.VANILLA:
        network.thenetwork.set_wo_style(network.WOMode.VANILLA,
            network.SNIPAuthMode.NONE)
    elif network_mode == network.WOMode.TELESCOPING:
        if snipauth_mode == network.SNIPAuthMode.MERKLE:
            network.thenetwork.set_wo_style(network.WOMode.TELESCOPING,
                network.SNIPAuthMode.MERKLE)
        else:
            network.thenetwork.set_wo_style(network.WOMode.TELESCOPING,
                network.SNIPAuthMode.THRESHSIG)
    elif network_mode == network.WOMode.SINGLEPASS:
        if snipauth_mode == network.SNIPAuthMode.MERKLE:
            network.thenetwork.set_wo_style(network.WOMode.SINGLEPASS,
                network.SNIPAuthMode.MERKLE)
        else:
            network.thenetwork.set_wo_style(network.WOMode.SINGLEPASS,
                network.SNIPAuthMode.THRESHSIG)
    else:
        sys.exit("Received unsupported network mode, exiting.")

    # Start some dirauths
    numdirauths = 9
    dirauthaddrs = []
    dirauths = []
    for i in range(numdirauths):
        dira = dirauth.DirAuth(i, numdirauths)
        dirauths.append(dira)
        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.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(), perfstats)

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

    relays[3].channelmgr.send_msg(relay.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)

    # Start some clients
    numclients = 1
    clients = []
    for i in range(numclients):
        clients.append(Client(dirauthaddrs))

    # Tick the epoch
    network.thenetwork.nextepoch()

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

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

    # Pick a bunch of bw-weighted random relays and look at the
    # distribution
    for i in range(100):
        r = relays[0].channelmgr.relaypicker.pick_weighted_relay()
        if network.thenetwork.womode == network.WOMode.VANILLA:
            print("relay",r.descdict["addr"])
        else:
            print("relay",r.snipdict["addr"])

    relays[3].terminate()
    relaysent += relays[3].perfstats.bytes_sent
    relayrecv += relays[3].perfstats.bytes_received
    del relays[3]

    # Tick the epoch
    network.thenetwork.nextepoch()

    circs = []
    for i in range(20):
        circ = clients[0].channelmgr.new_circuit()
        if circ is None:
            sys.exit("ERR: Client unable to create circuits")
        circs.append(circ)
        circ.send_cell(relay.StringCell("hello world circuit %d" % i))

    # Tick the epoch
    network.thenetwork.nextepoch()

    # See what channels exist and do a consistency check
    for r in relays:
        print("%s: %s" % (r.netaddr, [ str(k) + str([ck for ck in r.channelmgr.channels[k].circuithandlers.keys()]) 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)

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

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

            if ch.circuithandlers.keys() != \
                    ch.peer.channelmgr.channels[caddr].circuithandlers.keys():
                print('circuit asymmetry:', caddr, ad, ch.peer.channelmgr.myaddr)

    for c in circs:
        c.close()

    for d in dirauths:
        print(d.perfstats)
        dirasent += d.perfstats.bytes_sent
        dirarecv += d.perfstats.bytes_received
    print("DirAuths sent=%s recv=%s" % (dirasent, dirarecv))
    totsent += dirasent
    totrecv += dirarecv
    for r in relays:
        print(r.perfstats)
        relaysent += r.perfstats.bytes_sent
        relayrecv += r.perfstats.bytes_received
    print("Relays sent=%s recv=%s" % (relaysent, relayrecv))
    totsent += relaysent
    totrecv += relayrecv
    for c in clients:
        print(c.perfstats)
        clisent += c.perfstats.bytes_sent
        clirecv += c.perfstats.bytes_received
    print("Client sent=%s recv=%s" % (clisent, clirecv))
    totsent += clisent
    totrecv += clirecv

    print("Total sent=%s recv=%s" % (totsent, totrecv))