import logging import msg class Connection: def __init__(self, peer = None): """Create a Connection object with the given peer.""" self.peer = peer def closed(self): logging.debug("connection closed with %s", self.peer) self.peer = None def close(self): logging.debug("closing connection with %s", self.peer) self.peer.closed() self.peer = None class ClientConnection(Connection): """The parent class of client-side network connections. Subclass this class to do anything more elaborate than just passing arbitrary NetMsgs, which then get ignored. Use subclasses of this class when the server required no per-connection state, such as just fetching consensus documents.""" def __init__(self, peer): """Create a ClientConnection object with the given peer. The peer must have a received(client, msg) method.""" self.peer = peer self.perfstats = None def sendmsg(self, netmsg): assert(isinstance(netmsg, msg.NetMsg)) msgsize = netmsg.size() self.perfstats.bytes_sent += msgsize self.peer.received(self, netmsg) def reply(self, netmsg): assert(isinstance(netmsg, msg.NetMsg)) msgsize = netmsg.size() self.perfstats.bytes_received += msgsize self.receivedfromserver(netmsg) class ServerConnection(Connection): """The parent class of server-side network connections.""" def __init__(self): self.peer = None def sendmsg(self, netmsg): assert(isinstance(netmsg, msg.NetMsg)) self.peer.received(netmsg) def received(self, client, netmsg): logging.debug("received %s from client %s", netmsg, client) class DirAuthConnection(ClientConnection): """The subclass of Connection for connections to directory authorities.""" def __init__(self, peer): super().__init__(peer) def uploaddesc(self, desc): """Upload our RelayDescriptor to the DirAuth.""" self.sendmsg(msg.DirAuthUploadDescMsg(desc)) def getconsensus(self): self.consensus = None self.sendmsg(msg.DirAuthGetConsensusMsg()) return self.consensus def getconsensusdiff(self): self.consensus = None self.sendmsg(msg.DirAuthGetConsensusDiffMsg()) return self.consensus def getendive(self): self.endive = None self.sendmsg(msg.DirAuthGetENDIVEMsg()) return self.endive def getendivediff(self): self.endive = None self.sendmsg(msg.DirAuthGetENDIVEDiffMsg()) return self.endive def receivedfromserver(self, message): if isinstance(message, msg.DirAuthConsensusMsg): self.consensus = message.consensus elif isinstance(message, msg.DirAuthConsensusDiffMsg): self.consensus = message.consensus elif isinstance(message, msg.DirAuthENDIVEMsg): self.endive = message.endive elif isinstance(message, msg.DirAuthENDIVEDiffMsg): self.endive = message.endive else: raise TypeError('Not a server-originating DirAuthNetMsg', message)