From 3cab847c8f6994ec2d8fc82375b2772f081b5ab4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tarjei=20Hus=C3=B8y?= Date: Sat, 21 Apr 2018 12:24:53 -0700 Subject: [PATCH] Fix byte types for python3 --- postgres_mitm.py | 54 ++++++++++++++++++++++++------------------------ 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/postgres_mitm.py b/postgres_mitm.py index 25fc306..93d7ac1 100755 --- a/postgres_mitm.py +++ b/postgres_mitm.py @@ -37,9 +37,9 @@ from collections import namedtuple # Sent by client when requesting TLS connection (this is the magic version # 1234.5679 of the protocol, defined in pgcomm.h) -VERSION_SSL = '\x04\xd2\x16\x2f' -VERSION_3 = '\x00\x03\x00\x00' -SSL_STARTUP_RESPONSE = 'S' +VERSION_SSL = b'\x04\xd2\x16\x2f' +VERSION_3 = b'\x00\x03\x00\x00' +SSL_STARTUP_RESPONSE = b'S' _logger = logging.getLogger(__name__) @@ -105,7 +105,7 @@ def main(): try: cert_file = tempfile.NamedTemporaryFile(delete=False) with cert_file: - cert_file.write(CERTIFICATE) + cert_file.write(CERTIFICATE.encode('utf-8')) while True: client_socket, address = sock.accept() client_handler = ClientConnection(client_socket, target_backend, @@ -263,8 +263,8 @@ class ClientConnection(threading.Thread): def wait_for_auth_request(self): # Format of message is first_5_bytes = self.read_n_bytes_from_client(5) - tag = first_5_bytes[0] - if tag != 'p': + tag = first_5_bytes[0:1] + if tag != b'p': raise Exception("Received non-auth request: %s" % tag) # Bump length with 1 to offset for tag @@ -296,9 +296,9 @@ class ClientConnection(threading.Thread): self.startup_packet = data self.options = parse_options_from_startup_packet(data) _logger.debug('Startup packet processed successfully: %s' % self.options) - auth_reply = 'R%(length)s%(method)s' % { - 'length': struct.pack('!I', 8), - 'method': struct.pack('!I', AUTH_METHODS['AUTH_REQ_PASSWORD']), + auth_reply = b'R%(length)s%(method)s' % { + b'length': struct.pack('!I', 8), + b'method': struct.pack('!I', AUTH_METHODS['AUTH_REQ_PASSWORD']), } _logger.debug('Replying to startup: %s' % repr(auth_reply)) self.socket.send(auth_reply) @@ -311,10 +311,10 @@ class ClientConnection(threading.Thread): password = parse_password_from_authentication_packet(data) if self.connect_to_actual_backend(password): captured_uri = 'postgres://%(user)s:%(password)s@%(host)s:5432/%(database)s' % { - 'user': self.options.get('user', ''), + 'user': self.options.get('user', b''), 'password': password, 'host': self.target_backend, - 'database': self.options.get('database', ''), + 'database': self.options.get('database', b''), } _logger.info('Intercepted auth: %s' % captured_uri) # Switch socket to non-blocking to enable messages to pass in @@ -327,13 +327,13 @@ class ClientConnection(threading.Thread): def connect_to_actual_backend(self, password): self.server_socket = socket.create_connection((self.target_backend, 5432)) - packet = '%(length)s%(version)s' % { - 'length': struct.pack('!I', 8), - 'version': VERSION_SSL, + packet = b'%(length)s%(version)s' % { + b'length': struct.pack('!I', 8), + b'version': VERSION_SSL, } self.server_socket.sendall(packet) data = read_n_bytes_from_socket(self.server_socket, 1) - assert data == 'S' + assert data == b'S' self.server_socket = self.ssl_context.wrap_socket(self.server_socket) self.server_socket.do_handshake() self.server_socket.sendall(self.startup_packet) @@ -343,7 +343,7 @@ class ClientConnection(threading.Thread): 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) + response = create_md5_auth_packet(self.options.get('user', b''), password, salt) self.server_socket.sendall(response) else: _logger.debug('Unsupported backend auth method: %s' % auth_request.method) @@ -357,7 +357,7 @@ class ClientConnection(threading.Thread): def receive_auth_request_from_backend(self): first_9_bytes = self.read_n_bytes_from_server(9) - assert first_9_bytes[0] == 'R' + assert first_9_bytes[0:1] == b'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) @@ -391,17 +391,17 @@ def read_n_bytes_from_socket(sock, n): nbytes = sock.recv_into(view, n) view = view[nbytes:] # slicing views is cheap n -= nbytes - return str(buf) + return buf def create_md5_auth_packet(username, password, salt): pw_and_username = password + username pw_hash = hashlib.md5(pw_and_username).hexdigest() - salted_hash = 'md5' + hashlib.md5(pw_hash + salt).hexdigest() - response = 'p%(length)s%(salted_hash)s\0' % { + salted_hash = 'md5' + hashlib.md5(pw_hash.encode('utf-8') + salt).hexdigest() + response = b'p%(length)s%(salted_hash)s\x00' % { # 32 bytes of digest, four bytes length, 3 bytes for 'md5', one byte terminating null - 'length': struct.pack('!I', 40), - 'salted_hash': salted_hash, + b'length': struct.pack('!I', 40), + b'salted_hash': salted_hash.encode('utf-8'), } return response @@ -409,22 +409,22 @@ def create_md5_auth_packet(username, password, salt): def parse_options_from_startup_packet(data): # format is [\0\0]+\0 raw_key_value_pairs = data[8:] - assert raw_key_value_pairs[-1] == '\0' + assert raw_key_value_pairs[-1] == 0 raw_key_value_pairs = raw_key_value_pairs[0:-1] - assert raw_key_value_pairs.count('\0') % 2 == 0 + assert raw_key_value_pairs.count(0) % 2 == 0 options = {} - key_value_pairs = data[8:].split('\0') + key_value_pairs = data[8:].split(b'\x00') for i in range(0, len(key_value_pairs), 2): key = key_value_pairs[i] value = key_value_pairs[i + 1] - options[key] = value + options[key.decode('utf-8')] = value return options def parse_password_from_authentication_packet(data): - assert data[-1] == '\0' + assert data[-1] == 0 return data[5:-1]