From 1b5bda2eb4e8c0c0198d76b757db2a44fcb372a8 Mon Sep 17 00:00:00 2001 From: stcb <21@stcb.cc> Date: Sun, 23 Mar 2025 21:32:00 +0200 Subject: [PATCH] Semi-auto, bad sizes and need adjustments --- protocol_prototype/IcingProtocol.drawio | 132 ++++----- protocol_prototype/cli.py | 113 ++++++++ protocol_prototype/crypto_utils.py | 81 ++++++ protocol_prototype/main.py | 16 -- protocol_prototype/messages.py | 126 +++++++++ protocol_prototype/protocol.py | 354 ++++++++++++++++++++++++ protocol_prototype/transmission.py | 100 +++++++ 7 files changed, 840 insertions(+), 82 deletions(-) create mode 100644 protocol_prototype/cli.py create mode 100644 protocol_prototype/crypto_utils.py delete mode 100644 protocol_prototype/main.py create mode 100644 protocol_prototype/messages.py create mode 100644 protocol_prototype/protocol.py create mode 100644 protocol_prototype/transmission.py diff --git a/protocol_prototype/IcingProtocol.drawio b/protocol_prototype/IcingProtocol.drawio index 836e00d..edefb16 100644 --- a/protocol_prototype/IcingProtocol.drawio +++ b/protocol_prototype/IcingProtocol.drawio @@ -1,6 +1,6 @@ - + - + @@ -47,7 +47,7 @@ - + @@ -143,7 +143,7 @@ - + @@ -165,7 +165,7 @@ - + @@ -260,197 +260,197 @@ - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - - + + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + diff --git a/protocol_prototype/cli.py b/protocol_prototype/cli.py new file mode 100644 index 0000000..293566b --- /dev/null +++ b/protocol_prototype/cli.py @@ -0,0 +1,113 @@ +import sys +from protocol import IcingProtocol + +RED = "\033[91m" +GREEN = "\033[92m" +YELLOW = "\033[93m" +BLUE = "\033[94m" +RESET = "\033[0m" + + +def main(): + protocol = IcingProtocol() + + print(f"{YELLOW}\n======================================") + print(" Icing Protocol - Manual CLI Demo ") + print("======================================\n" + RESET) + + print(f"Listening on port: {protocol.local_port}") + print(f"Your identity public key (hex): {protocol.identity_pubkey.hex()}") + print("\nAvailable commands:") + print(" set_peer_identity ") + print(" connect ") + print(" generate_ephemeral_keys") + print(" send_ping") + print(" send_handshake") + print(" respond_ping ") + print(" generate_ecdhe # formerly respond_handshake") + print(" auto_responder ") + print(" show_state") + print(" exit\n") + + while True: + try: + line = input("Cmd> ").strip() + except EOFError: + break + if not line: + continue + + parts = line.split() + cmd = parts[0].lower() + + if cmd == "exit": + protocol.stop() + sys.exit(0) + + elif cmd == "show_state": + protocol.show_state() + + elif cmd == "set_peer_identity": + if len(parts) != 2: + print(f"{RED}Usage: set_peer_identity {RESET}") + continue + protocol.set_peer_identity(parts[1]) + + elif cmd == "connect": + if len(parts) != 2: + print(f"{RED}Usage: connect {RESET}") + continue + try: + port = int(parts[1]) + protocol.connect_to_peer(port) + except ValueError: + print(f"{RED}Invalid port.{RESET}") + + elif cmd == "generate_ephemeral_keys": + protocol.generate_ephemeral_keys() + + elif cmd == "send_ping": + protocol.send_ping_request() + + elif cmd == "send_handshake": + protocol.send_handshake() + + elif cmd == "respond_ping": + if len(parts) != 3: + print(f"{RED}Usage: respond_ping {RESET}") + continue + try: + idx = int(parts[1]) + ac = int(parts[2]) + protocol.respond_to_ping(idx, ac) + except ValueError: + print(f"{RED}Index and answer_code must be integers.{RESET}") + + elif cmd == "generate_ecdhe": + if len(parts) != 2: + print(f"{RED}Usage: generate_ecdhe {RESET}") + continue + try: + idx = int(parts[1]) + protocol.generate_ecdhe(idx) + except ValueError: + print(f"{RED}Index must be an integer.{RESET}") + + elif cmd == "auto_responder": + if len(parts) != 2: + print(f"{RED}Usage: auto_responder {RESET}") + continue + arg = parts[1].lower() + if arg == "on": + protocol.enable_auto_responder(True) + elif arg == "off": + protocol.enable_auto_responder(False) + else: + print(f"{RED}Usage: auto_responder {RESET}") + + else: + print(f"{RED}[ERROR]{RESET} Unknown command: {cmd}") + + +if __name__ == "__main__": + main() diff --git a/protocol_prototype/crypto_utils.py b/protocol_prototype/crypto_utils.py new file mode 100644 index 0000000..de0ce7c --- /dev/null +++ b/protocol_prototype/crypto_utils.py @@ -0,0 +1,81 @@ +import os +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.hazmat.primitives.asymmetric.ec import ECDH +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import utils +from cryptography.exceptions import InvalidSignature + + +def generate_identity_keys(): + """ + Generate an ECDSA (P-256) identity key pair. + Return (private_key, public_key_bytes). + public_key_bytes is raw x||y each 32 bytes (uncompressed minus the 0x04 prefix). + """ + private_key = ec.generate_private_key(ec.SECP256R1()) + public_numbers = private_key.public_key().public_numbers() + + x_bytes = public_numbers.x.to_bytes(32, byteorder='big') + y_bytes = public_numbers.y.to_bytes(32, byteorder='big') + pubkey_bytes = x_bytes + y_bytes # 64 bytes + + return private_key, pubkey_bytes + + +def load_peer_identity_key(pubkey_bytes: bytes): + """ + Given 64 bytes (x||y) for P-256, return a cryptography public key object. + """ + if len(pubkey_bytes) != 64: + raise ValueError("Peer identity pubkey must be exactly 64 bytes (x||y).") + x_int = int.from_bytes(pubkey_bytes[:32], byteorder='big') + y_int = int.from_bytes(pubkey_bytes[32:], byteorder='big') + public_numbers = ec.EllipticCurvePublicNumbers(x_int, y_int, ec.SECP256R1()) + return public_numbers.public_key() + + +def sign_data(private_key, data: bytes) -> bytes: + """ + Sign 'data' with ECDSA using P-256 private key. + Returns DER-encoded signature (variable length, up to ~70-72 bytes). + """ + signature = private_key.sign(data, ec.ECDSA(hashes.SHA256())) + return signature + + +def verify_signature(public_key, signature: bytes, data: bytes) -> bool: + """ + Verify DER-encoded ECDSA signature with the given public key. + Return True if valid, False otherwise. + """ + try: + public_key.verify(signature, data, ec.ECDSA(hashes.SHA256())) + return True + except InvalidSignature: + return False + + +def get_ephemeral_keypair(): + """ + Generate ephemeral ECDH keypair (P-256). + Return (private_key, pubkey_bytes). + """ + private_key = ec.generate_private_key(ec.SECP256R1()) + numbers = private_key.public_key().public_numbers() + x_bytes = numbers.x.to_bytes(32, 'big') + y_bytes = numbers.y.to_bytes(32, 'big') + return private_key, x_bytes + y_bytes # 64 bytes + + +def compute_ecdh_shared_key(private_key, peer_pubkey_bytes: bytes) -> bytes: + """ + Given a local ECDH private_key and the peer's ephemeral pubkey (64 bytes), + compute the shared secret. + """ + x_int = int.from_bytes(peer_pubkey_bytes[:32], 'big') + y_int = int.from_bytes(peer_pubkey_bytes[32:], 'big') + peer_public_numbers = ec.EllipticCurvePublicNumbers(x_int, y_int, ec.SECP256R1()) + peer_public_key = peer_public_numbers.public_key() + shared_key = private_key.exchange(ec.ECDH(), peer_public_key) + return shared_key diff --git a/protocol_prototype/main.py b/protocol_prototype/main.py deleted file mode 100644 index 5596b44..0000000 --- a/protocol_prototype/main.py +++ /dev/null @@ -1,16 +0,0 @@ -# This is a sample Python script. - -# Press Shift+F10 to execute it or replace it with your code. -# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings. - - -def print_hi(name): - # Use a breakpoint in the code line below to debug your script. - print(f'Hi, {name}') # Press Ctrl+F8 to toggle the breakpoint. - - -# Press the green button in the gutter to run the script. -if __name__ == '__main__': - print_hi('PyCharm') - -# See PyCharm help at https://www.jetbrains.com/help/pycharm/ diff --git a/protocol_prototype/messages.py b/protocol_prototype/messages.py new file mode 100644 index 0000000..9ebcb72 --- /dev/null +++ b/protocol_prototype/messages.py @@ -0,0 +1,126 @@ +import os +import struct +import time +import zlib + + +def crc32_of(data: bytes) -> int: + """ + Compute CRC-32 of 'data'. + """ + return zlib.crc32(data) & 0xffffffff + + +# ----------------------------------------------------------------------------- +# Ping +# ----------------------------------------------------------------------------- + +def build_ping_request(version: int = 0) -> bytes: + """ + Build a Ping request: + - Nonce (32 bytes) + - Version (1 byte) + - CRC-32 (4 bytes) + Total = 37 bytes + """ + nonce = os.urandom(32) + partial = nonce + struct.pack("!B", version) + crc_val = crc32_of(partial) + return partial + struct.pack("!I", crc_val) + + +def parse_ping_request(data: bytes): + """ + Parse a Ping request (37 bytes). + Return (nonce, version) or None if invalid. + """ + if len(data) != 37: + return None + nonce = data[:32] + version = data[32] + crc_in = struct.unpack("!I", data[33:37])[0] + partial = data[:33] + crc_calc = crc32_of(partial) + if crc_calc != crc_in: + return None + return (nonce, version) + + +def build_ping_response(version: int, answer_code: int) -> bytes: + """ + Build a Ping response: + - Timestamp (8 bytes) + - Version (1 byte) + - Answer code (1 byte) + - CRC-32 (4 bytes) + Total = 14 bytes + """ + timestamp = struct.pack("!d", time.time()) + partial = timestamp + struct.pack("!B", version) + struct.pack("!B", answer_code) + crc_val = crc32_of(partial) + return partial + struct.pack("!I", crc_val) + + +def parse_ping_response(data: bytes): + """ + Parse a Ping response (14 bytes). + Return (timestamp, version, answer_code) or None if invalid. + """ + if len(data) != 14: + return None + timestamp = struct.unpack("!d", data[:8])[0] + version = data[8] + answer_code = data[9] + crc_in = struct.unpack("!I", data[10:14])[0] + partial = data[:10] + crc_calc = crc32_of(partial) + if crc_calc != crc_in: + return None + return (timestamp, version, answer_code) + + +# ----------------------------------------------------------------------------- +# Handshake +# ----------------------------------------------------------------------------- + +def build_handshake_message(ephemeral_pubkey: bytes, ephemeral_signature: bytes, timestamp: float) -> bytes: + """ + Build a handshake message: + - ephemeral_pubkey (64 bytes) + - ephemeral_signature (72 bytes, DER + zero-pad) + - timestamp (8 bytes) + - CRC-32 (4 bytes) + Total = 148 bytes + """ + if len(ephemeral_pubkey) != 64: + raise ValueError("ephemeral_pubkey must be 64 bytes") + if len(ephemeral_signature) > 72: + raise ValueError("ephemeral_signature too large") + + sig_padded = ephemeral_signature.ljust(72, b'\x00') + ts_bytes = struct.pack("!d", timestamp) + partial = ephemeral_pubkey + sig_padded + ts_bytes + crc_val = crc32_of(partial) + return partial + struct.pack("!I", crc_val) + + +def parse_handshake_message(data: bytes): + """ + Parse a handshake message (148 bytes). + Return (ephemeral_pubkey, ephemeral_signature, timestamp) or None if invalid. + """ + if len(data) != 148: + return None + ephemeral_pubkey = data[:64] + sig_padded = data[64:136] + ts_bytes = data[136:144] + crc_in = struct.unpack("!I", data[144:148])[0] + + partial = data[:144] + crc_calc = crc32_of(partial) + if crc_calc != crc_in: + return None + + ephemeral_signature = sig_padded.rstrip(b'\x00') + timestamp = struct.unpack("!d", ts_bytes)[0] + return (ephemeral_pubkey, ephemeral_signature, timestamp) diff --git a/protocol_prototype/protocol.py b/protocol_prototype/protocol.py new file mode 100644 index 0000000..3e928d7 --- /dev/null +++ b/protocol_prototype/protocol.py @@ -0,0 +1,354 @@ +import random +import time +import threading +from typing import List, Dict, Any + +from crypto_utils import ( + generate_identity_keys, + load_peer_identity_key, + sign_data, + verify_signature, + get_ephemeral_keypair, + compute_ecdh_shared_key +) +from messages import ( + build_ping_request, parse_ping_request, + build_ping_response, parse_ping_response, + build_handshake_message, parse_handshake_message +) +import transmission + +# ANSI colors for pretty printing +RED = "\033[91m" +GREEN = "\033[92m" +YELLOW = "\033[93m" +BLUE = "\033[94m" +RESET = "\033[0m" + + +class IcingProtocol: + def __init__(self): + # Identity keys + self.identity_privkey, self.identity_pubkey = generate_identity_keys() + + # Peer identity (public key) for verifying ephemeral signatures + self.peer_identity_pubkey_obj = None + self.peer_identity_pubkey_bytes = None + + # Ephemeral keys (our side) + self.ephemeral_privkey = None + self.ephemeral_pubkey = None + + # Store the last computed shared secret (hex) from ECDH + self.shared_secret = None + + # Track open connections + self.connections = [] + + # A random listening port + self.local_port = random.randint(30000, 40000) + + # Inbound messages are stored for manual or auto handling + # Each entry: { 'type': str, 'raw': bytes, 'parsed': Any, 'connection': PeerConnection } + self.inbound_messages: List[Dict[str, Any]] = [] + + # Simple dictionary to track protocol flags + self.state = { + "ping_sent": False, + "ping_received": False, + "handshake_sent": False, + "handshake_received": False, + } + + # Auto-responder toggle + self.auto_responder = False + + # Start the listener + self.server_listener = transmission.ServerListener( + host="127.0.0.1", + port=self.local_port, + on_new_connection=self.on_new_connection, + on_data_received=self.on_data_received + ) + self.server_listener.start() + + # ------------------------------------------------------------------------- + # Transport callbacks + # ------------------------------------------------------------------------- + + def on_new_connection(self, conn: transmission.PeerConnection): + print(f"{GREEN}[IcingProtocol]{RESET} New incoming connection.") + self.connections.append(conn) + + def on_data_received(self, conn: transmission.PeerConnection, data: bytes): + """ + Called whenever data arrives on any open PeerConnection. + We'll parse and store the message, then handle automatically if auto_responder is on. + """ + # Print data size in bits, not bytes + bits_count = len(data) * 8 + print(f"{GREEN}[RECV]{RESET} {bits_count} bits from peer: {data.hex()[:60]}{'...' if len(data.hex())>60 else ''}") + + # Attempt to parse Ping request + if len(data) == 37: + parsed = parse_ping_request(data) + if parsed: + nonce, version = parsed + self.state["ping_received"] = True + index = len(self.inbound_messages) + msg = { + "type": "PING_REQUEST", + "raw": data, + "parsed": {"nonce": nonce, "version": version}, + "connection": conn + } + self.inbound_messages.append(msg) + print(f"{YELLOW}[NOTICE]{RESET} Stored inbound PING request (nonce={nonce.hex()}) at index={index}.") + + if self.auto_responder: + # Schedule an automatic response after 2 seconds + threading.Timer(2.0, self._auto_respond_ping, args=(index,)).start() + + return + + # Attempt to parse Ping response + if len(data) == 14: + parsed = parse_ping_response(data) + if parsed: + ts, version, answer_code = parsed + index = len(self.inbound_messages) + msg = { + "type": "PING_RESPONSE", + "raw": data, + "parsed": {"timestamp": ts, "version": version, "answer_code": answer_code}, + "connection": conn + } + self.inbound_messages.append(msg) + print(f"{YELLOW}[NOTICE]{RESET} Stored inbound PING response (answer_code={answer_code}) at index={index}.") + return + + # Attempt to parse handshake + if len(data) == 148: + parsed = parse_handshake_message(data) + if parsed: + ephemeral_pub, ephemeral_sig, ts = parsed + self.state["handshake_received"] = True + index = len(self.inbound_messages) + msg = { + "type": "HANDSHAKE", + "raw": data, + "parsed": { + "ephemeral_pub": ephemeral_pub, + "ephemeral_sig": ephemeral_sig, + "timestamp": ts + }, + "connection": conn + } + self.inbound_messages.append(msg) + print(f"{YELLOW}[NOTICE]{RESET} Stored inbound HANDSHAKE at index={index}. ephemeral_pub={ephemeral_pub.hex()[:20]}...") + + if self.auto_responder: + # Schedule an automatic handshake "response" after 2 seconds + threading.Timer(2.0, self._auto_respond_handshake, args=(index,)).start() + + return + + # Otherwise, unrecognized + index = len(self.inbound_messages) + msg = { + "type": "UNKNOWN", + "raw": data, + "parsed": None, + "connection": conn + } + self.inbound_messages.append(msg) + print(f"{RED}[WARNING]{RESET} Unrecognized or malformed message stored at index={index}.") + + # ------------------------------------------------------------------------- + # Auto-responder helpers + # ------------------------------------------------------------------------- + + def _auto_respond_ping(self, index: int): + """ + Called by a Timer to respond automatically to a PING_REQUEST after 2s. + """ + print(f"{BLUE}[AUTO]{RESET} Delayed responding to PING at index={index}") + self.respond_to_ping(index, answer_code=0) + self.show_state() + + def _auto_respond_handshake(self, index: int): + """ + Called by a Timer to handle inbound HANDSHAKE automatically. + 1) Generate ephemeral keys if not already set + 2) Compute ECDH with the inbound ephemeral pub (generate_ecdhe) + 3) Send our handshake back + 4) Show state + """ + print(f"{BLUE}[AUTO]{RESET} Delayed ECDH process for HANDSHAKE at index={index}") + + # 1) Generate ephemeral keys if we don't have them + if not self.ephemeral_privkey or not self.ephemeral_pubkey: + self.generate_ephemeral_keys() + + # 2) Compute ECDH from inbound ephemeral pub + self.generate_ecdhe(index) + + # 3) Send our handshake to the peer + self.send_handshake() + + # 4) Show final state + self.show_state() + + # ------------------------------------------------------------------------- + # Public Methods + # ------------------------------------------------------------------------- + + def connect_to_peer(self, port: int): + conn = transmission.connect_to_peer("127.0.0.1", port, self.on_data_received) + self.connections.append(conn) + print(f"{GREEN}[IcingProtocol]{RESET} Outgoing connection to port {port} established.") + + def set_peer_identity(self, peer_pubkey_hex: str): + pubkey_bytes = bytes.fromhex(peer_pubkey_hex) + self.peer_identity_pubkey_obj = load_peer_identity_key(pubkey_bytes) + self.peer_identity_pubkey_bytes = pubkey_bytes + print(f"{GREEN}[IcingProtocol]{RESET} Stored peer identity pubkey (hex={peer_pubkey_hex[:16]}...).") + + def generate_ephemeral_keys(self): + self.ephemeral_privkey, self.ephemeral_pubkey = get_ephemeral_keypair() + print(f"{GREEN}[IcingProtocol]{RESET} Generated ephemeral key pair: pubkey={self.ephemeral_pubkey.hex()[:16]}...") + + def send_ping_request(self): + if not self.connections: + print(f"{RED}[ERROR]{RESET} No active connections.") + return + pkt = build_ping_request(version=0) + self._send_packet(self.connections[0], pkt, "PING_REQUEST") + self.state["ping_sent"] = True + + def send_handshake(self): + """ + Build and send a handshake message with ephemeral keys. + """ + if not self.connections: + print(f"{RED}[ERROR]{RESET} No active connections.") + return + if not self.ephemeral_privkey or not self.ephemeral_pubkey: + print(f"{RED}[ERROR]{RESET} Ephemeral keys not generated. Call 'generate_ephemeral_keys' first.") + return + + ephemeral_signature = sign_data(self.identity_privkey, self.ephemeral_pubkey) + ts_now = time.time() + pkt = build_handshake_message(self.ephemeral_pubkey, ephemeral_signature, ts_now) + self._send_packet(self.connections[0], pkt, "HANDSHAKE") + self.state["handshake_sent"] = True + + def enable_auto_responder(self, enable: bool): + self.auto_responder = enable + print(f"{BLUE}[AUTO]{RESET} Auto responder set to {enable}.") + + # ------------------------------------------------------------------------- + # Manual Responses + # ------------------------------------------------------------------------- + + def respond_to_ping(self, index: int, answer_code: int): + """ + Manually respond to an inbound PING_REQUEST in inbound_messages[index]. + """ + if index < 0 or index >= len(self.inbound_messages): + print(f"{RED}[ERROR]{RESET} Invalid index {index}.") + return + msg = self.inbound_messages[index] + if msg["type"] != "PING_REQUEST": + print(f"{RED}[ERROR]{RESET} inbound_messages[{index}] is not a PING_REQUEST.") + return + + version = msg["parsed"]["version"] + conn = msg["connection"] + resp = build_ping_response(version, answer_code) + self._send_packet(conn, resp, "PING_RESPONSE") + print(f"{BLUE}[MANUAL]{RESET} Responded to ping with answer_code={answer_code}.") + + def generate_ecdhe(self, index: int): + """ + Formerly 'respond_to_handshake' - this verifies the inbound ephemeral signature + and computes the ECDH shared secret, storing it in self.shared_secret. + It does NOT send a handshake back. + """ + if index < 0 or index >= len(self.inbound_messages): + print(f"{RED}[ERROR]{RESET} Invalid index {index}.") + return + msg = self.inbound_messages[index] + if msg["type"] != "HANDSHAKE": + print(f"{RED}[ERROR]{RESET} inbound_messages[{index}] is not a HANDSHAKE.") + return + + ephemeral_pub = msg["parsed"]["ephemeral_pub"] + ephemeral_sig = msg["parsed"]["ephemeral_sig"] + + # Verify ephemeral signature + if not self.peer_identity_pubkey_obj: + print(f"{RED}[ERROR]{RESET} Peer identity not set, cannot verify ephemeral signature.") + return + ok = verify_signature(self.peer_identity_pubkey_obj, ephemeral_sig, ephemeral_pub) + if not ok: + print(f"{RED}[ERROR]{RESET} Ephemeral signature is invalid.") + return + print(f"{GREEN}[OK]{RESET} Ephemeral signature verified successfully.") + + # If we have ephemeral_privkey, compute ECDH shared key + if self.ephemeral_privkey: + shared = compute_ecdh_shared_key(self.ephemeral_privkey, ephemeral_pub) + self.shared_secret = shared.hex() + print(f"{GREEN}[OK]{RESET} Derived ECDH shared key = {self.shared_secret}") + else: + print(f"{YELLOW}[WARN]{RESET} No ephemeral_privkey available, cannot compute ECDH shared key.") + + # ------------------------------------------------------------------------- + # Utility + # ------------------------------------------------------------------------- + + def _send_packet(self, conn: transmission.PeerConnection, data: bytes, label: str): + bits_count = len(data) * 8 + print(f"{BLUE}[SEND]{RESET} {label} -> {bits_count} bits: {data.hex()[:60]}{'...' if len(data.hex())>60 else ''}") + conn.send(data) + + def show_state(self): + print(f"\n{YELLOW}=== Global State ==={RESET}") + print(f"Listening Port: {self.local_port}") + print(f"Identity PubKey: {self.identity_pubkey.hex()[:16]}... (64 bytes)") + if self.peer_identity_pubkey_bytes: + print(f"Peer Identity PubKey: {self.peer_identity_pubkey_bytes.hex()[:16]}... (64 bytes)") + else: + print("Peer Identity PubKey: [None]") + + print("\nEphemeral Keys:") + if self.ephemeral_pubkey: + print(f" ephemeral_pubkey={self.ephemeral_pubkey.hex()[:16]}...") + else: + print(" ephemeral_pubkey=[None]") + + print(f"\nShared Secret: {self.shared_secret if self.shared_secret else '[None]'}") + + print("\nProtocol Flags:") + for k, v in self.state.items(): + print(f" {k}: {v}") + + print("\nAuto Responder:", self.auto_responder) + + print("\nActive Connections:") + for i, c in enumerate(self.connections): + print(f" [{i}] Alive={c.alive}") + + print("\nInbound Message Queue:") + for i, m in enumerate(self.inbound_messages): + print(f" [{i}] type={m['type']} len={len(m['raw'])} bytes") + print() + + def stop(self): + self.server_listener.stop() + for c in self.connections: + c.close() + self.connections.clear() + self.inbound_messages.clear() + print(f"{RED}[STOP]{RESET} Protocol stopped.") diff --git a/protocol_prototype/transmission.py b/protocol_prototype/transmission.py new file mode 100644 index 0000000..35f3a21 --- /dev/null +++ b/protocol_prototype/transmission.py @@ -0,0 +1,100 @@ +import socket +import threading +from typing import Callable + +class PeerConnection: + """ + Represents a live, two-way connection to a peer. + We keep a socket open, read data in a background thread, + and can send data from the main thread at any time. + """ + def __init__(self, sock: socket.socket, on_data_received: Callable[['PeerConnection', bytes], None]): + self.sock = sock + self.on_data_received = on_data_received + self.alive = True + + self.read_thread = threading.Thread(target=self.read_loop, daemon=True) + self.read_thread.start() + + def read_loop(self): + while self.alive: + try: + data = self.sock.recv(4096) + if not data: + break + self.on_data_received(self, data) + except OSError: + break + self.alive = False + self.sock.close() + print("[PeerConnection] Connection closed.") + + def send(self, data: bytes): + if not self.alive: + print("[PeerConnection.send] Cannot send, connection not alive.") + return + try: + self.sock.sendall(data) + except OSError: + print("[PeerConnection.send] Send failed, connection might be closed.") + self.alive = False + + def close(self): + self.alive = False + try: + self.sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + self.sock.close() + + +class ServerListener(threading.Thread): + """ + A thread that listens on a given port. When a new client connects, + it creates a PeerConnection for that client. + """ + def __init__(self, host: str, port: int, + on_new_connection: Callable[[PeerConnection], None], + on_data_received: Callable[[PeerConnection, bytes], None]): + super().__init__(daemon=True) + self.host = host + self.port = port + self.on_new_connection = on_new_connection + self.on_data_received = on_data_received + self.server_socket = None + self.stop_event = threading.Event() + + def run(self): + self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.server_socket.bind((self.host, self.port)) + self.server_socket.listen(5) + self.server_socket.settimeout(1.0) + print(f"[ServerListener] Listening on {self.host}:{self.port}") + + while not self.stop_event.is_set(): + try: + client_sock, addr = self.server_socket.accept() + print(f"[ServerListener] Accepted connection from {addr}") + conn = PeerConnection(client_sock, self.on_data_received) + self.on_new_connection(conn) + except socket.timeout: + pass + except OSError: + break + + if self.server_socket: + self.server_socket.close() + + def stop(self): + self.stop_event.set() + if self.server_socket: + self.server_socket.close() + + +def connect_to_peer(host: str, port: int, + on_data_received: Callable[[PeerConnection, bytes], None]) -> PeerConnection: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + print(f"[connect_to_peer] Connected to {host}:{port}") + conn = PeerConnection(sock, on_data_received) + return conn