From 8fa0465bdf9f650016bb9eef53c5aca525704f6d Mon Sep 17 00:00:00 2001 From: dsc Date: Fri, 3 May 2019 19:49:08 +0200 Subject: [PATCH] Exit early on invalid response --- peer_retreiver.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/peer_retreiver.py b/peer_retreiver.py index c2ba000..5c9b06f 100644 --- a/peer_retreiver.py +++ b/peer_retreiver.py @@ -30,16 +30,22 @@ buckets = [] while 1: buffer = sock.recv(8) + if not buffer: + sys.stderr.write("Invalid response; exiting\n") + break - if buffer.startswith(bytes(LEVIN_SIGNATURE)): - bucket = Bucket.from_buffer(signature=buffer, sock=sock) - buckets.append(bucket) + if not buffer.startswith(bytes(LEVIN_SIGNATURE)): + sys.stderr.write("Invalid response; exiting\n") + break - if bucket.command == 1001: - peers = bucket.get_peers() or [] + bucket = Bucket.from_buffer(signature=buffer, sock=sock) + buckets.append(bucket) - for peer in peers: - print('%s:%d' % (peer['ip'].ipv4, peer['port'].value)) + if bucket.command == 1001: + peers = bucket.get_peers() or [] - sock.close() - break + for peer in peers: + print('%s:%d' % (peer['ip'].ipv4, peer['port'].value)) + + sock.close() + break