jwt_postgresql_proxy.py

from gevent import monkey
monkey.patch_all()

from base64 import urlsafe_b64decode
from hashlib import md5
import itertools
import os
import re
import gevent
import json
import socket
import ssl
import struct
import time

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.serialization import load_pem_public_key


class ConnectionClosed(Exception):
    pass


class ProtocolError(Exception):
    pass


class AuthenticationError(Exception):
    pass


def main():
    env = normalise_environment(os.environ)
    UPSTREAM__HOST = env['UPSTREAM']['HOST']
    UPSTREAM__PORT = env['UPSTREAM']['PORT']
    UPSTREAM__PASSWORD = env['UPSTREAM']['PASSWORD']
    DOWNSTREAM__IP = env['DOWNSTREAM']['IP']
    DOWNSTREAM__PORT = env['DOWNSTREAM']['PORT']
    DOWNSTREAM__CERTFILE = env['DOWNSTREAM']['CERTFILE']
    DOWNSTREAM__KEYFILE = env['DOWNSTREAM']['KEYFILE']
    PUBLIC_KEYS = env['PUBLIC_KEYS']

    TLS_REQUEST = b'\x00\x00\x00\x08\x04\xd2\x16/'
    TLS_RESPONSE = b'S'

    STARTUP_MESSAGE_HEADER = struct.Struct('!LL')
    MESSAGE_HEADER = struct.Struct('!cL')
    INT = struct.Struct('!L')

    PROTOCOL_VERSION = 196608

    AUTHENTICATION_CLEARTEXT_PASSWORD = 3
    AUTHENTICATION_MD5_PASSWORD = 5
    AUTHENTICATION_OK = 0
    PASSWORD_RESPONSE = b'p'

    MAX_READ = 66560
    MAX_IN_MEMORY_MESSAGE_LENGTH = 66560

    ssl_context_downstream = ssl.SSLContext(ssl.PROTOCOL_TLS)
    ssl_context_downstream.load_cert_chain(
        certfile=DOWNSTREAM__CERTFILE, keyfile=DOWNSTREAM__KEYFILE)

    ssl_context_upstream = ssl.SSLContext(ssl.PROTOCOL_TLS)
    ssl_context_upstream.verify_mode = ssl.CERT_NONE

    public_keys = [
        load_pem_public_key(public_key.encode(), backend=default_backend())
        for public_key in PUBLIC_KEYS
    ]

    def b64_decode(b64_bytes):
        return urlsafe_b64decode(b64_bytes + (b'=' * ((4 - len(b64_bytes) % 4) % 4)))

    def handle_downstream(downstream_sock):
        # The high level logic of connection, authentication, and proxying, is all here

        downstream_sock_ssl = None
        upstream_sock = None
        upstream_sock_ssl = None

        try:
            # Initiate TLS
            downstream_sock_ssl = downstream_convert_to_ssl(downstream_sock)

            # Startup PostgreSQL downstream
            user, database = downstream_startup(downstream_sock_ssl)

            # Authenticate downstream user
            downstream_authenticate(downstream_sock_ssl, user)

            # Connect on TCP level upstream
            upstream_sock = upstream_connect()

            # Convert upstream to TLS
            upstream_sock_ssl = upstream_convert_to_ssl(upstream_sock)

            # Startup PostgreSQL upstream
            upstream_startup(upstream_sock_ssl, user, database)

            # Authenticate upstream user
            upstream_authenticate(upstream_sock_ssl, user)

            # Run proxy
            proxy_both_directions(downstream_sock_ssl, upstream_sock_ssl)

        except AuthenticationError:
            downstream_send_auth_error(downstream_sock_ssl or downstream_sock)

        except ProtocolError:
            downstream_send_error(downstream_sock_ssl or downstream_sock)

        finally:
            # Slightly faffy cleanup to deal the various cases where things could have stopped
            # at various points in the process

            if upstream_sock_ssl is not None:
                try:
                    upstream_sock = upstream_sock_ssl.unwrap()
                except (OSError, ssl.SSLError):
                    pass

            if upstream_sock is not None:
                try:
                    upstream_sock.shutdown(socket.SHUT_RDWR)
                except OSError:
                    # Could have shutdown already
                    pass
                finally:
                    upstream_sock.close()

            if downstream_sock_ssl is not None:
                try:
                    downstream_sock = downstream_sock_ssl.unwrap()
                except (OSError, ssl.SSLError):
                    pass

            try:
                downstream_sock.shutdown(socket.SHUT_RDWR)
            except OSError:
                # Could have shutdown already
                pass
            finally:
                downstream_sock.close()

    def downstream_convert_to_ssl(downstream_sock):
        chunk = recv_exactly(downstream_sock, len(TLS_REQUEST))
        if chunk != TLS_REQUEST:
            downstream_sock.sendall(MESSAGE_HEADER.pack(b'E', 4 + 1) + b'\x00')
            raise ProtocolError()
        downstream_sock.sendall(TLS_RESPONSE)
        downstream_sock_ssl = ssl_context_downstream.wrap_socket(downstream_sock, server_side=True)
        return downstream_sock_ssl

    def downstream_startup(downstream_sock_ssl):
        startup_message_len, protocol_version = STARTUP_MESSAGE_HEADER.unpack(
            recv_exactly(downstream_sock_ssl, STARTUP_MESSAGE_HEADER.size))
        if startup_message_len > MAX_IN_MEMORY_MESSAGE_LENGTH:
            raise ProtocolError('Startup message too large')

        if protocol_version != PROTOCOL_VERSION:
            downstream_sock_ssl.sendall(MESSAGE_HEADER.pack(b'E', 4 + 1) + b'\x00')
            raise ProtocolError('Unsupported downstream protocol version')

        startup_key_value_pairs = recv_exactly(
            downstream_sock_ssl, startup_message_len - STARTUP_MESSAGE_HEADER.size)

        pairs = dict(re.compile(b'([^\x00]+)\x00([^\x00]*)').findall(startup_key_value_pairs))

        return pairs[b'user'].decode(), pairs[b'database'].decode()

    def downstream_authenticate(downstream_sock_ssl, claimed_user):
        # Request password
        downstream_sock_ssl.sendall(MESSAGE_HEADER.pack(b'R', 4 + INT.size)
                                    + INT.pack(AUTHENTICATION_CLEARTEXT_PASSWORD))

        # Get password response
        tag, payload_length = MESSAGE_HEADER.unpack(
            recv_exactly(downstream_sock_ssl, MESSAGE_HEADER.size))
        if payload_length > MAX_IN_MEMORY_MESSAGE_LENGTH:
            raise ProtocolError('Password response message too large')
        if tag != PASSWORD_RESPONSE:
            raise ProtocolError('Expected password to request for password')
        password = (recv_exactly(downstream_sock_ssl, payload_length - 4))[:-1]

        # Verify signature
        header_b64, payload_b64, signature_b64 = password.split(b'.')
        for public_key in public_keys:
            try:
                # pylint: disable=no-value-for-parameter
                public_key.verify(b64_decode(signature_b64), header_b64 + b'.' + payload_b64)
            except InvalidSignature:
                continue
            else:
                break
        else:
            raise AuthenticationError()

        # Ensure the signed JWT `sub` is the same as the claimed database user
        payload = json.loads(b64_decode(payload_b64))
        if claimed_user != payload['sub']:
            raise AuthenticationError()

        now = time.time()
        if payload['exp'] <= now:
            raise AuthenticationError()

        # Tell downstream we are authenticated
        downstream_sock_ssl.sendall(MESSAGE_HEADER.pack(
            b'R', 4 + INT.size) + INT.pack(AUTHENTICATION_OK))

    def downstream_send_auth_error(downstream_sock):
        failed = \
            b'S' + b'FATAL\x00' + \
            b'M' + b'Authentication failed\x00' + \
            b'C' + b'28P01\x00' + \
            b'\x00'
        downstream_sock.sendall(MESSAGE_HEADER.pack(b'E', 4 + len(failed)) + failed)

    def downstream_send_error(downstream_sock):
        downstream_sock.sendall(MESSAGE_HEADER.pack(b'E', 4 + 1) + b'\x00')

    def upstream_connect():
        upstream_sock = socket.create_connection((UPSTREAM__HOST, UPSTREAM__PORT))
        upstream_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        return upstream_sock

    def upstream_convert_to_ssl(upstream_sock):
        upstream_sock.sendall(TLS_REQUEST)
        response = recv_exactly(upstream_sock, len(TLS_RESPONSE))
        if response != TLS_RESPONSE:
            raise ProtocolError()

        upstream_sock_ssl = ssl_context_upstream.wrap_socket(upstream_sock)
        return upstream_sock_ssl

    def upstream_startup(upstream_sock_ssl, user, database):
        pairs = \
            b'user\x00' + user.encode() + b'\x00' + \
            b'database\x00' + database.encode() + b'\x00' + \
            b'\x00'
        upstream_sock_ssl.sendall(
            STARTUP_MESSAGE_HEADER.pack(8 + len(pairs), PROTOCOL_VERSION) + pairs
        )

    def upstream_authenticate(upstream_sock_ssl, user):
        tag, message_len = MESSAGE_HEADER.unpack(
            recv_exactly(upstream_sock_ssl, MESSAGE_HEADER.size))
        message = recv_exactly(upstream_sock_ssl, message_len - 4)
        if tag != b'R':
            raise ProtocolError()

        # We only support MD5 password exchange or none
        auth_method, = INT.unpack(message[:4])

        if auth_method == AUTHENTICATION_OK:
            return

        if auth_method != AUTHENTICATION_MD5_PASSWORD:
            raise ProtocolError()

        salt = message[4:]
        hashed_password = b'md5' + \
            md5(md5(UPSTREAM__PASSWORD.encode() + user.encode()).hexdigest().encode()
                + salt).hexdigest().encode() + b'\x00'

        upstream_sock_ssl.sendall(MESSAGE_HEADER.pack(
            PASSWORD_RESPONSE, 4 + len(hashed_password)) + hashed_password)

        tag, message_len = MESSAGE_HEADER.unpack(
            recv_exactly(upstream_sock_ssl, MESSAGE_HEADER.size))
        message = recv_exactly(upstream_sock_ssl, message_len - 4)

        if tag != b'R':
            raise AuthenticationError()
        auth_result, = INT.unpack(message)
        if auth_result != AUTHENTICATION_OK:
            raise AuthenticationError()

    def proxy_both_directions(sock_a, sock_b):
        done = gevent.event.Event()

        def _proxy(source, target):
            try:
                while chunk := source.recv(MAX_READ):
                    target.sendall(chunk)
            finally:
                done.set()

        a_to_b_greenlet = gevent.spawn(_proxy, sock_a, sock_b)
        b_to_a_greenlet = gevent.spawn(_proxy, sock_b, sock_a)
        done.wait()

        a_to_b_greenlet.kill()
        b_to_a_greenlet.kill()
        a_to_b_greenlet.join()
        b_to_a_greenlet.join()

    def get_new_socket():
        sock = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM,
                             proto=socket.IPPROTO_TCP)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        return sock

    def recv_exactly(sock, amount):
        chunks = []
        while amount:
            chunk = sock.recv(min(amount, MAX_READ))
            if chunk:
                chunks.append(chunk)
                amount -= len(chunks[-1])
            else:
                raise ProtocolError()
        joined = b''.join(chunks)
        return joined

    sock = get_new_socket()
    sock.bind((DOWNSTREAM__IP, int(DOWNSTREAM__PORT)))
    sock.listen(socket.IPPROTO_TCP)

    while True:
        downstream_sock, _ = sock.accept()
        gevent.spawn(handle_downstream, downstream_sock)
        downstream_sock = None  # To make sure we don't have it hanging around


def normalise_environment(key_values):
    separator = '__'

    def get_first_component(key):
        return key.split(separator)[0]

    def get_later_components(key):
        return separator.join(key.split(separator)[1:])

    without_more_components = {
        key: value
        for key, value in key_values.items()
        if not get_later_components(key)
    }

    with_more_components = {
        key: value
        for key, value in key_values.items()
        if get_later_components(key)
    }

    def grouped_by_first_component(items):
        def by_first_component(item):
            return get_first_component(item[0])

        return itertools.groupby(
            sorted(items, key=by_first_component),
            by_first_component,
        )

    def items_with_first_component(items, first_component):
        return {
            get_later_components(key): value
            for key, value in items
            if get_first_component(key) == first_component
        }

    nested_structured_dict = {
        **without_more_components, **{
            first_component: normalise_environment(
                items_with_first_component(items, first_component))
            for first_component, items in grouped_by_first_component(with_more_components.items())
        }}

    def all_keys_are_ints():
        def is_int(string):
            try:
                int(string)
                return True
            except ValueError:
                return False

        return all([is_int(key) for key, value in nested_structured_dict.items()])

    def list_sorted_by_int_key():
        return [
            value
            for key, value in sorted(
                nested_structured_dict.items(),
                key=lambda key_value: int(key_value[0])
            )
        ]

    return \
        list_sorted_by_int_key() if all_keys_are_ints() else \
        nested_structured_dict


if __name__ == '__main__':
    main()