walkingonions/relay.py

#!/usr/bin/env python3

import nacl.utils
from nacl.signing import SigningKey
from nacl.public import PrivateKey, Box

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 = SigningKey.generate()
        self.onionkey = PrivateKey.generate()

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

        self.uploaddesc()

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