Define abstract bases in message and structure to reduce docstrings

This commit is contained in:
Lee Miller 2023-09-08 19:28:58 +03:00
parent fe508c176b
commit 5ddaa3ce65
Signed by untrusted user: lee.miller
GPG Key ID: 4F97A5EA88F4AB63
2 changed files with 41 additions and 13 deletions

View File

@ -4,11 +4,28 @@ import base64
import hashlib
import struct
import time
from abc import ABC, abstractmethod
from . import shared, structure
class Header():
class IMessage(ABC):
"""A base for typical message"""
@abstractmethod
def __repr__(self):
"""Make a printable form"""
@abstractmethod
def to_bytes(self):
"""Serialize to bytes the full message"""
@classmethod
@abstractmethod
def from_message(cls, m):
"""Parse from message"""
class Header(structure.IStructure):
"""Message header structure"""
def __init__(self, command, payload_length, payload_checksum):
self.command = command
@ -24,7 +41,6 @@ class Header():
base64.b16encode(self.payload_checksum).decode())
def to_bytes(self):
"""Serialize to bytes"""
b = b''
b += shared.magic_bytes
b += self.command.ljust(12, b'\x00')
@ -34,7 +50,6 @@ class Header():
@classmethod
def from_bytes(cls, b):
"""Parse from bytes"""
magic_bytes, command, payload_length, payload_checksum = struct.unpack(
'>4s12sL4s', b)
@ -46,7 +61,7 @@ class Header():
return cls(command, payload_length, payload_checksum)
class Message():
class Message(structure.IStructure):
"""Common message structure"""
def __init__(self, command, payload):
self.command = command
@ -61,7 +76,6 @@ class Message():
base64.b16encode(self.payload_checksum).decode())
def to_bytes(self):
"""Serialize to bytes"""
b = Header(
self.command, self.payload_length, self.payload_checksum
).to_bytes()
@ -70,7 +84,6 @@ class Message():
@classmethod
def from_bytes(cls, b):
"""Parse from bytes"""
h = Header.from_bytes(b[:24])
payload = b[24:]
@ -114,6 +127,7 @@ class Version():
base64.b16encode(self.nonce).decode(), self.user_agent)
def to_bytes(self):
"""Serialize to bytes"""
payload = b''
payload += struct.pack('>I', self.protocol_version)
payload += struct.pack('>Q', self.services)
@ -131,6 +145,7 @@ class Version():
@classmethod
def from_bytes(cls, b):
"""Parse from bytes"""
m = Message.from_bytes(b)
payload = m.payload
@ -162,7 +177,7 @@ class Version():
return cls(host, port, protocol_version, services, nonce, user_agent)
class Inv():
class Inv(IMessage):
"""The inv message payload"""
def __init__(self, vectors):
self.vectors = set(vectors)
@ -198,7 +213,7 @@ class Inv():
return cls(vectors)
class GetData():
class GetData(IMessage):
"""The getdata message payload"""
def __init__(self, vectors):
self.vectors = set(vectors)
@ -234,7 +249,7 @@ class GetData():
return cls(vectors)
class Addr():
class Addr(IMessage):
"""The addr message payload"""
def __init__(self, addresses):
self.addresses = addresses

View File

@ -6,11 +6,24 @@ import logging
import socket
import struct
import time
from abc import ABC, abstractmethod
from . import shared
class VarInt():
class IStructure(ABC):
"""A base for typical structure"""
@abstractmethod
def to_bytes(self):
"""Serialize to bytes"""
@classmethod
@abstractmethod
def from_bytes(cls, b):
"""Parse from bytes"""
class VarInt(IStructure):
"""varint object"""
def __init__(self, n):
self.n = n
@ -81,7 +94,7 @@ class Object():
nonce, expires_time, object_type, version, stream_number, payload)
def to_bytes(self):
"""Serialize to bytes"""
"""Serialize to bytes object payload"""
payload = b''
payload += self.nonce
payload += struct.pack('>QL', self.expires_time, self.object_type)
@ -147,7 +160,7 @@ class Object():
return hashlib.sha512(self.to_bytes()[8:]).digest()
class NetAddrNoPrefix():
class NetAddrNoPrefix(IStructure):
"""Network address"""
def __init__(self, services, host, port):
self.services = services
@ -180,7 +193,7 @@ class NetAddrNoPrefix():
return cls(services, host, port)
class NetAddr():
class NetAddr(IStructure):
"""Network address with time and stream"""
def __init__(self, services, host, port, stream=shared.stream):
self.stream = stream