Onionr/static-data/official-plugins/tor/torpeer.py

67 lines
1.9 KiB
Python

import socks
from gossip.peerset import gossip_peer_set
from logger import log as logging
class HandleRevc:
def __init__(self, sock):
self.sock_recv = sock.recv
self.sock = sock
def recv(self, *args, **kwargs):
try:
got_data = self.sock_recv(*args, **kwargs)
if not len(got_data):
raise ConnectionError("Peer socket returned empty value")
return got_data
except Exception:
try:
gossip_peer_set.remove(self.sock)
except KeyError:
pass
raise
class TorPeer:
def __init__(self, socks_host, socks_port, onion_address):
if not onion_address or onion_address == '.onion':
raise ValueError("Invalid transport address")
if not onion_address.endswith('.onion'):
self.onion_address = onion_address.strip() + '.onion'
self.transport_address = self.onion_address = onion_address
self.socks_host = socks_host
self.socks_port = socks_port
def get_socket(self, connect_timeout) -> socks.socksocket:
s = socks.socksocket()
s.set_proxy(socks.SOCKS4, self.socks_host, self.socks_port, rdns=True)
s.settimeout(connect_timeout)
try:
s.connect((self.onion_address, 80))
except socks.GeneralProxyError:
try:
gossip_peer_set.remove(self)
except KeyError:
pass
else:
logging.debug(
f"Could not create socket to peer {self.transport_address}",
)
raise TimeoutError
mock_recv = HandleRevc(s)
s.recv = mock_recv.recv
return s
def __hash__(self):
return hash(self.transport_address)
def __eq__(self, other):
try:
return self.transport_address == other.transport_address
except AttributeError:
return False