diff --git a/get_cert.py b/get_cert.py index 1df6cb0..75c7281 100755 --- a/get_cert.py +++ b/get_cert.py @@ -15,29 +15,26 @@ logging.basicConfig(level=logging.INFO, logger = logging.getLogger('certo') -# The following two functions from +# The following inspired by: # https://stackoverflow.com/questions/17667903/python-socket-receive-large-amount-of-data +# https://www.binarytides.com/receive-full-data-with-the-recv-socket-function-in-python/ def recv_msg(sock, timeout): # Read message length and unpack it into an integer - raw_msglen = recvall(sock, 1024, timeout) - if not raw_msglen: - return None - msglen = struct.unpack('>I', raw_msglen)[0] - # Read the message data - return recvall(sock, msglen, timeout) - - -def recvall(sock, n, timeout): - # Helper function to recv n bytes or return None if EOF is hit data = b'' begin = time.time() - while len(data) < n: - packet = sock.recv(n - len(data)) - if not packet: - return None - data += packet - if time.time() - begin > timeout: + n = 1024 + while len(data) < n or time.time() - begin < timeout: + try: + data += sock.recv(n - len(data)) + logger.debug("Partial: %r" % data) + except socket.timeout: + pass + if b'Ready to start TLS\r\n' in data: break + time.sleep(0.1) + if len(data) == 0: + return None + logger.debug("Data: %r" % data) return data