import logging import sys import struct import socket from io import BytesIO from levin.section import Section from levin.constants import * from levin.exceptions import BadArgumentException from levin.ctypes import * log = logging.getLogger() class Bucket: def __init__(self): self.signature = LEVIN_SIGNATURE if self.signature != LEVIN_SIGNATURE: raise BadArgumentException('Bender\'s nightmare') self.cb = None self.return_data = None self.command = None self.return_code = None self.flags = None self.protocol_version = None self.payload_section = None @classmethod def create_request(cls, command: int, payload: bytes = None, section: Section = None): bucket = cls() bucket.return_data = c_bool(True) bucket.command = c_uint32(command) bucket.return_code = c_uint32(0) bucket.flags = LEVIN_PACKET_REQUEST bucket.protocol_version = LEVIN_PROTOCOL_VER_1 if payload: bucket.cb = c_uint64(len(payload)) bucket.payload_section = payload if section: bucket.payload_section = bytes(section) bucket.cb = c_uint64(len(bucket.payload_section)) return bucket @classmethod def create_response(cls, command: int, payload: bytes, return_code: int): bucket = cls() bucket.cb = c_uint64(len(payload)) bucket.return_data = c_bool(False) bucket.command = c_uint32(command) bucket.return_code = c_uint32(return_code) bucket.flags = LEVIN_PACKET_RESPONSE bucket.protocol_version = LEVIN_PROTOCOL_VER_1 @staticmethod def create_handshake_request( my_port: int = 0, network_id: bytes = None, peer_id: bytes = b'\x41\x41\x41\x41\x41\x41\x41\x41', ): """ Helper function to create a handshake request. Does not require parameters but you can use them to impersonate a legit node. :param my_port: defaults to 0 :param network_id: defaults to mainnet :param peer_id: :param verbose: :return: """ handshake_section = Section.handshake_request(peer_id=peer_id, network_id=network_id, my_port=my_port) bucket = Bucket.create_request(P2P_COMMAND_HANDSHAKE.value, section=handshake_section) log.debug(">> created packet '%s'" % P2P_COMMANDS[bucket.command]) header = bucket.header() body = bucket.payload() return bucket @staticmethod def create_stat_info_request( peer_id: bytes = b"\x41\x41\x41\x41\x41\x41\x41\x41", ): stat_info_section = Section.stat_info_request(peer_id=peer_id) log.debug(stat_info_section.entries["proof_of_trust"].entries.keys()) bucket = Bucket.create_request( P2P_COMMAND_REQUEST_STAT_INFO.value, section=stat_info_section ) log.debug(">> created packet '%s'" % P2P_COMMANDS[bucket.command]) return bucket @classmethod def from_buffer(cls, signature: c_uint64, sock: socket.socket): if isinstance(signature, bytes): signature = c_uint64(signature) # if isinstance(buffer, bytes): # buffer = BytesIO(buffer) bucket = cls() bucket.signature = signature bucket.cb = c_uint64.from_buffer(sock) bucket.return_data = c_bool.from_buffer(sock) bucket.command = c_uint32.from_buffer(sock) bucket.return_code = c_int32.from_buffer(sock) bucket.flags = c_uint32.from_buffer(sock) bucket.protocol_version = c_uint32.from_buffer(sock) if bucket.signature != LEVIN_SIGNATURE: raise IOError("Bender's nightmare missing") if bucket.cb > LEVIN_DEFAULT_MAX_PACKET_SIZE: raise IOError("payload too large") if bucket.command.value not in P2P_COMMANDS: raise IOError("unregonized command: %d" % bucket.command.value) bucket.payload = sock.recv(bucket.cb.value) while len(bucket.payload) < bucket.cb.value: bufsize = 2048 remaining = bucket.cb.value - len(bucket.payload) if remaining < bufsize: bufsize = remaining bucket.payload += sock.recv(bufsize) bucket.payload_section = None log.debug("<< received packet '%s'" % P2P_COMMANDS[bucket.command]) from levin.reader import LevinReader bucket.payload_section = LevinReader(BytesIO(bucket.payload)).read_payload() log.debug("<< parsed packet '%s'" % P2P_COMMANDS[bucket.command]) return bucket def header(self): return b''.join(( bytes(LEVIN_SIGNATURE), bytes(self.cb), b'\x01' if self.return_data else b'\x00', bytes(self.command), bytes(self.return_code), bytes(self.flags), bytes(self.protocol_version) )) def payload(self): return self.payload_section def get_peers(self): # helper function to retreive peerlisting where buckets.command was 1001 if self.command != 1001: raise Exception("Only handshake has peerlisting") peers = [] if 'local_peerlist_new' not in self.payload_section.entries: return for peer in [e.entries for e in self.payload_section.entries['local_peerlist_new']]: if 'adr' not in peer or 'addr' not in peer['adr'].entries: continue addr = peer['adr'].entries['addr'].entries ipv4 = True if "m_ip" in addr else False if not ipv4 and not "addr" in addr: continue if ipv4 and len(addr["m_ip"]) == 4: m_ip, m_port = addr['m_ip'], addr['m_port'] m_ip = c_uint32(m_ip.to_bytes(), endian='big') elif len(addr["addr"]) == 16: m_ip, m_port = addr['addr'], addr["m_port"] m_ip = c_uint64(m_ip, endian="big") else: continue peers.append({ 'ip': m_ip, 'port': m_port }) return peers