#!/usr/bin/env python3

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

import network
import dirauth
import relay

import cell as simcell
import msg as simmsg

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(
                    simcell.VanillaCreatedCircuitCell, None)
        else:
            # This was a VanillaExtendedCircuitCell
            circhandler.replace_celltype_handler(
                    simcell.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 = simcell.VanillaExtendCircuitCell(
                nexthop.descdict['addr'], ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                simcell.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
        if type(self.channelmgr.guard) is dirauth.RelayDescriptor:
            guardd = self.channelmgr.guard.descdict
        else:
            guardd = self.channelmgr.guard.snipdict
        self.onionkey = guardd["onionkey"]
        self.idkey = guardd["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(simcell.TelescopingCreatedCircuitCell, None)

        circhandler.circuit_descs.append(self.channelmgr.guard)

        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 = simcell.TelescopingExtendCircuitCell(
                nexthopidx, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                simcell.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(
                    simcell.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:
            # Yes!
            return

        nexthopidx = self.channelmgr.relaypicker.pick_weighted_relay_index()

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

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                simcell.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):
        # We should only get one simcell.SinglePassCreatedCircuitCell per
        # circuit
        circhandler.replace_celltype_handler(simcell.SinglePassCreatedCircuitCell, None)

        # The circuit always starts with the guard
        circhandler.circuit_descs.append(self.channelmgr.guard)

        # Process each layer of the message
        blinding_keys = []
        while cell is not None:
            lasthop = circhandler.circuit_descs[-1]
            if type(lasthop) is dirauth.RelayDescriptor:
                lasthopd = lasthop.descdict
            else:
                lasthopd = lasthop.snipdict
            onionkey = lasthopd["onionkey"]
            idkey = lasthopd["idkey"]
            pathselkey = lasthopd["pathselkey"]
            if cell.enc is None:
                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)
                cell = None
            else:
                secret = self.ntor.verify(cell.ntor_reply, onionkey,
                        idkey, b'circuit')
                enckey = nacl.hash.sha256(secret + b'upstream')
                deckey = nacl.hash.sha256(secret + b'downstream')
                createdkey = nacl.hash.sha256(secret + b'created')
                circhandler.add_crypt_layer(enckey, deckey)
                (snip, vrfout, nextlayer) = cell.enc.decrypt(createdkey)
                # Check the signature on the SNIP
                dirauth.SNIP.verify(snip, self.channelmgr.consensus,
                        network.thenetwork.dirauthkeys()[0],
                        self.channelmgr.perfstats)
                # Compute the index, check the VRF, ensure the SNIP is
                # the correct one
                pathsel_rand, next_blindkey = relay.Sphinx.client(
                        self.client_key, blinding_keys,
                        onionkey, b'pathsel',
                        nextlayer is None, self.channelmgr.perfstats)
                if nextlayer is not None:
                    blinding_keys.append(next_blindkey)

                try:
                    index = int.from_bytes(relay.VRF.check_output(pathselkey,
                                        pathsel_rand, vrfout,
                                        self.channelmgr.perfstats)[:4],
                                    'big', signed=False)
                except ValueError as e:
                    circhandler.close()
                    raise ValueError(str(e.args) + str(lasthopd))

                indexrange = snip.snipdict["range"]
                if index < indexrange[0] or index >= indexrange[1]:
                    logging.error("Incorrect SNIP received")

                circhandler.circuit_descs.append(snip)
                cell = nextlayer


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
        self.circuit_build_begin = 0

    def get_consensus_from_fallbackrelay(self):
        """Download a fresh consensus from a random fallbackrelay."""
        fb = network.thenetwork.getfallbackrelay()
        logging.debug("Chose fallback %s", fb)
        if network.thenetwork.womode == network.WOMode.VANILLA:
            if self.consensus is not None and \
                    len(self.consensus.consdict['relays']) > 0:
                self.send_msg(simmsg.RelayGetConsensusDiffMsg(), fb.netaddr)
            else:
                self.send_msg(simmsg.RelayGetConsensusMsg(), fb.netaddr)
        else:
            self.send_msg(simmsg.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

        # Ensure we have the current descriptor for the guard
        # Note that self.guard may be a RelayDescriptor or a SNIP,
        # depending on how we got it
        if type(self.guard) is dirauth.RelayDescriptor:
            guardepoch = self.guard.descdict["epoch"]
        else:
            guardepoch = self.guard.snipdict["epoch"]
        if guardepoch != network.thenetwork.getepoch():
            guardchannel = self.get_channel_to(self.guardaddr)
            guardchannel.send_msg(simmsg.RelayGetDescMsg())

        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 = simmsg.VanillaCreateCircuitMsg(circid, ntor_request)

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

        # Send the message
        guardchannel.send_msg(circcreatemsg)

        if guardchannel.is_circuit_open(circid):
            self.perfstats.circuit_building_time = time.perf_counter() - self.circuit_build_begin
            logging.debug("Client %s ended circuit building after %s", self.myaddr, str(self.perfstats.circuit_building_time))

        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 = simmsg.TelescopingCreateCircuitMsg(circid, ntor_request)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                simcell.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 guardchannel.is_circuit_open(circid):
            self.perfstats.circuit_building_time = time.perf_counter() - self.circuit_build_begin
            logging.debug("Client %s ended circuit building after %s", self.myaddr, str(self.perfstats.circuit_building_time))
        else:
            logging.debug("Circuit was already closed, not sending bytes. circid: " + str(circid))
            return None

        guard = circhandler.circuit_descs[0]
        if type(guard) is dirauth.RelayDescriptor:
            guardd = guard.descdict
        else:
            guardd = guard.snipdict
        if guardd["addr"] == circhandler.circuit_descs[2].snipdict["addr"]:
            logging.debug("circuit in a loop")
            circhandler.close()
            circhandler = 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_pathsel_key = nacl.public.PrivateKey.generate()
        self.perfstats.keygens += 1

        # Construct the SinglePassCreateCircuitMsg
        ntor = relay.NTor(self.perfstats)
        ntor_request = ntor.request()
        circcreatemsg = simmsg.SinglePassCreateCircuitMsg(circid, ntor_request,
                client_pathsel_key.public_key)

        # Set up the reply handler
        circhandler.replace_celltype_handler(
                simcell.SinglePassCreatedCircuitCell,
                SinglePassCreatedHandler(self, ntor, client_pathsel_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 guardchannel.is_circuit_open(circid):
            self.perfstats.circuit_building_time = time.perf_counter() - self.circuit_build_begin
            logging.debug("Client %s ended circuit building after %s", self.myaddr, str(self.perfstats.circuit_building_time))
        else:
            logging.debug("Circuit was already closed, not sending bytes. circid: " + str(circid))
            return None

        # In Single-Pass Walking Onions, we need to check whether the
        # circuit got into a loop (guard equals exit); each node will
        # refuse to extend to itself, so this is the only possible loop
        # in a circuit of length 3
        guard = circhandler.circuit_descs[0]
        if type(guard) is dirauth.RelayDescriptor:
            guardd = guard.descdict
        else:
            guardd = guard.snipdict
        if guardd["addr"] == circhandler.circuit_descs[2].snipdict["addr"]:
            logging.debug("circuit in a loop")
            circhandler.close()
            circhandler = None

        return circhandler

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

        # Start tracking the time for building a circuit
        self.circuit_build_begin = time.perf_counter()
        logging.debug("%s began circuit building after %s", self.myaddr, str(self.circuit_build_begin))

        # 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, simmsg.RelayConsensusMsg) or \
                isinstance(msg, simmsg.RelayConsensusDiffMsg):
            self.relaypicker = dirauth.Consensus.verify(msg.consensus,
                    network.thenetwork.dirauthkeys(), self.perfstats)
            self.consensus = msg.consensus
        elif isinstance(msg, simmsg.RelayDescMsg):
            dirauth.RelayDescriptor.verify(msg.desc, self.perfstats)
            self.guard = msg.desc
        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(cell, simcell.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(simmsg.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(simmsg.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()