1
0
walkingonions-boosted/connection.py
2022-03-17 17:05:34 +01:00

102 lines
3.1 KiB
Python

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)