#!/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 Relay(network.Server):
    """The class representing an onion relay."""

    def __init__(self, dirauthaddrs, bw, flags):
        self.consensus = None
        self.dirauthaddrs = dirauthaddrs

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

        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'.
        a = random.choice(self.dirauthaddrs)
        c = network.thenetwork.connect(self, a)
        self.consensus = c.getconsensus()
        c.close()

    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.dirauthaddrs:
            c = network.thenetwork.connect(self, a)
            c.sendmsg(descmsg)
            c.close()


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
    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))
        Relay(dirauthaddrs, bw, 0)

    # Tick the epoch
    network.thenetwork.nextepoch()

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

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