#!/usr/bin/env python3

import random # For simulation, not cryptography!
import math

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

import network
import dirauth

class RelayNetMsg(network.NetMsg):
    """The subclass of NetMsg for messages between relays and either
    relays or clients."""


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


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

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


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

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

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


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

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


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

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

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

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

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

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


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

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

    def get_consensus(self):
        """Download a fresh consensus from a random dirauth."""
        a = random.choice(self.dirauthaddrs)
        c = network.thenetwork.connect(self, a)
        self.consensus = c.getconsensus()
        print('consensus downloaded:', self.consensus)
        c.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.cellrelay = self
        self.channels[peeraddr] = channel

    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 connection
        newconn = network.thenetwork.connect(self.myaddr, addr)
        self.channels[addr] = newconn
        newconn.closer = lambda: self.channels.pop(addr)
        newconn.cellrelay = self

        return newconn

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


    def received_cell(self, circid, cell, peeraddr, peer):
        """Callback with a circuit-specific cell is received."""
        print("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."""
        conn = self.get_channel_to(peeraddr)
        conn.send_msg(msg)

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


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

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

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

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

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

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

        self.uploaddesc()

    def epoch_ending(self, epoch):
        # Download the new consensus, which will have been created
        # already since the dirauths' epoch_ending callbacks happened
        # before the relays'.
        self.cellrelay.get_consensus()

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

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

        descmsg = dirauth.DirAuthUploadDescMsg(desc)

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

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

        # Create the linked pair
        peerchannel = Channel()
        ourchannel = Channel()
        peerchannel.peer = ourchannel
        ourchannel.peer = peerchannel

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

        return peerchannel

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

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

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

    # Tick the epoch
    network.thenetwork.nextepoch()

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

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

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

    # See what channels exist
    for r in relays:
        print("%s: %s" % (r.netaddr, [ str(k) for k in r.cellrelay.channels.keys()]))
    #relays[3].cellrelay.send_cell(1, network.StringNetMsg("test"), relays[3].consensus.consdict['relays'][5].descdict['addr'])
    #relays[3].cellrelay.send_cell(2, network.StringNetMsg("cell"), relays[3].consensus.consdict['relays'][6].descdict['addr'])
    #relays[3].cellrelay.send_cell(2, network.StringNetMsg("again"), relays[3].consensus.consdict['relays'][1].descdict['addr'])
    #relays[3].cellrelay.send_cell(2, network.StringNetMsg("and again"), relays[3].consensus.consdict['relays'][5].descdict['addr'])