Fix byte types for python3

This commit is contained in:
Tarjei Husøy 2018-04-21 12:24:53 -07:00
parent 6199600e14
commit 3cab847c8f

View file

@ -37,9 +37,9 @@ from collections import namedtuple
# Sent by client when requesting TLS connection (this is the magic version # Sent by client when requesting TLS connection (this is the magic version
# 1234.5679 of the protocol, defined in pgcomm.h) # 1234.5679 of the protocol, defined in pgcomm.h)
VERSION_SSL = '\x04\xd2\x16\x2f' VERSION_SSL = b'\x04\xd2\x16\x2f'
VERSION_3 = '\x00\x03\x00\x00' VERSION_3 = b'\x00\x03\x00\x00'
SSL_STARTUP_RESPONSE = 'S' SSL_STARTUP_RESPONSE = b'S'
_logger = logging.getLogger(__name__) _logger = logging.getLogger(__name__)
@ -105,7 +105,7 @@ def main():
try: try:
cert_file = tempfile.NamedTemporaryFile(delete=False) cert_file = tempfile.NamedTemporaryFile(delete=False)
with cert_file: with cert_file:
cert_file.write(CERTIFICATE) cert_file.write(CERTIFICATE.encode('utf-8'))
while True: while True:
client_socket, address = sock.accept() client_socket, address = sock.accept()
client_handler = ClientConnection(client_socket, target_backend, client_handler = ClientConnection(client_socket, target_backend,
@ -263,8 +263,8 @@ class ClientConnection(threading.Thread):
def wait_for_auth_request(self): def wait_for_auth_request(self):
# Format of message is <char tag><int32 len><message> # Format of message is <char tag><int32 len><message>
first_5_bytes = self.read_n_bytes_from_client(5) first_5_bytes = self.read_n_bytes_from_client(5)
tag = first_5_bytes[0] tag = first_5_bytes[0:1]
if tag != 'p': if tag != b'p':
raise Exception("Received non-auth request: %s" % tag) raise Exception("Received non-auth request: %s" % tag)
# Bump length with 1 to offset for tag # Bump length with 1 to offset for tag
@ -296,9 +296,9 @@ class ClientConnection(threading.Thread):
self.startup_packet = data self.startup_packet = data
self.options = parse_options_from_startup_packet(data) self.options = parse_options_from_startup_packet(data)
_logger.debug('Startup packet processed successfully: %s' % self.options) _logger.debug('Startup packet processed successfully: %s' % self.options)
auth_reply = 'R%(length)s%(method)s' % { auth_reply = b'R%(length)s%(method)s' % {
'length': struct.pack('!I', 8), b'length': struct.pack('!I', 8),
'method': struct.pack('!I', AUTH_METHODS['AUTH_REQ_PASSWORD']), b'method': struct.pack('!I', AUTH_METHODS['AUTH_REQ_PASSWORD']),
} }
_logger.debug('Replying to startup: %s' % repr(auth_reply)) _logger.debug('Replying to startup: %s' % repr(auth_reply))
self.socket.send(auth_reply) self.socket.send(auth_reply)
@ -311,10 +311,10 @@ class ClientConnection(threading.Thread):
password = parse_password_from_authentication_packet(data) password = parse_password_from_authentication_packet(data)
if self.connect_to_actual_backend(password): if self.connect_to_actual_backend(password):
captured_uri = 'postgres://%(user)s:%(password)s@%(host)s:5432/%(database)s' % { 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, 'password': password,
'host': self.target_backend, 'host': self.target_backend,
'database': self.options.get('database', ''), 'database': self.options.get('database', b''),
} }
_logger.info('Intercepted auth: %s' % captured_uri) _logger.info('Intercepted auth: %s' % captured_uri)
# Switch socket to non-blocking to enable messages to pass in # 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): def connect_to_actual_backend(self, password):
self.server_socket = socket.create_connection((self.target_backend, 5432)) self.server_socket = socket.create_connection((self.target_backend, 5432))
packet = '%(length)s%(version)s' % { packet = b'%(length)s%(version)s' % {
'length': struct.pack('!I', 8), b'length': struct.pack('!I', 8),
'version': VERSION_SSL, b'version': VERSION_SSL,
} }
self.server_socket.sendall(packet) self.server_socket.sendall(packet)
data = read_n_bytes_from_socket(self.server_socket, 1) 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 = self.ssl_context.wrap_socket(self.server_socket)
self.server_socket.do_handshake() self.server_socket.do_handshake()
self.server_socket.sendall(self.startup_packet) self.server_socket.sendall(self.startup_packet)
@ -343,7 +343,7 @@ class ClientConnection(threading.Thread):
if auth_request.method == 'AUTH_REQ_MD5': if auth_request.method == 'AUTH_REQ_MD5':
# options is 4-byte salt # options is 4-byte salt
salt = auth_request.options 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) self.server_socket.sendall(response)
else: else:
_logger.debug('Unsupported backend auth method: %s' % auth_request.method) _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): def receive_auth_request_from_backend(self):
first_9_bytes = self.read_n_bytes_from_server(9) 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] packet_length = struct.unpack('!I', first_9_bytes[1:5])[0]
# Tag doesn't count on length, read the rest # Tag doesn't count on length, read the rest
the_rest = self.read_n_bytes_from_server(packet_length - 8) 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) nbytes = sock.recv_into(view, n)
view = view[nbytes:] # slicing views is cheap view = view[nbytes:] # slicing views is cheap
n -= nbytes n -= nbytes
return str(buf) return buf
def create_md5_auth_packet(username, password, salt): def create_md5_auth_packet(username, password, salt):
pw_and_username = password + username pw_and_username = password + username
pw_hash = hashlib.md5(pw_and_username).hexdigest() pw_hash = hashlib.md5(pw_and_username).hexdigest()
salted_hash = 'md5' + hashlib.md5(pw_hash + salt).hexdigest() salted_hash = 'md5' + hashlib.md5(pw_hash.encode('utf-8') + salt).hexdigest()
response = 'p%(length)s%(salted_hash)s\0' % { response = b'p%(length)s%(salted_hash)s\x00' % {
# 32 bytes of digest, four bytes length, 3 bytes for 'md5', one byte terminating null # 32 bytes of digest, four bytes length, 3 bytes for 'md5', one byte terminating null
'length': struct.pack('!I', 40), b'length': struct.pack('!I', 40),
'salted_hash': salted_hash, b'salted_hash': salted_hash.encode('utf-8'),
} }
return response return response
@ -409,22 +409,22 @@ def create_md5_auth_packet(username, password, salt):
def parse_options_from_startup_packet(data): def parse_options_from_startup_packet(data):
# format is <in32 length><in32 protocol>[<key>\0<value>\0]+\0 # format is <in32 length><in32 protocol>[<key>\0<value>\0]+\0
raw_key_value_pairs = data[8:] 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] 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 = {} 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): for i in range(0, len(key_value_pairs), 2):
key = key_value_pairs[i] key = key_value_pairs[i]
value = key_value_pairs[i + 1] value = key_value_pairs[i + 1]
options[key] = value options[key.decode('utf-8')] = value
return options return options
def parse_password_from_authentication_packet(data): def parse_password_from_authentication_packet(data):
assert data[-1] == '\0' assert data[-1] == 0
return data[5:-1] return data[5:-1]