import ssl
import socket
import json
import argparse

exception_ciphers = {}

def enumerate_tls_ciphers(protocol_options, host, port, cert, cafile):
    root_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
    root_context.options |= protocol_options
    root_context.set_ciphers('ALL:COMPLEMENTOFALL:-PSK:-SRP')

    ciphers = {cipher['name'] for cipher in root_context.get_ciphers()}

    accepted_ciphers = []

    for cipher_name in ciphers:
        context = ssl.SSLContext(root_context.protocol)
        try:
            context.set_ciphers(cipher_name)
        except ssl.SSLError as error:
            exception_ciphers[cipher_name] = str(error)
        context.options = root_context.options
        context.load_verify_locations(cafile=cafile)
        context.load_cert_chain(certfile=cert)

        with socket.socket(socket.AF_INET) as sock:
            with context.wrap_socket(sock, server_hostname=host) as conn:
                try:
                    conn.connect((host, port))
                except Exception as e:
                    continue
                accepted_ciphers.append(cipher_name)

    return sorted(accepted_ciphers)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MongoDB TLS Cipher Suite Enumerator')
    parser.add_argument('--port', type=int, default=27017, help='Port to connect to')
    parser.add_argument('-o', '--outfile', type=str, default='ciphers.json',
                        help='file to write the output to')
    parser.add_argument('--host', type=str, default='localhost', help='host to connect to')
    parser.add_argument('--cafile', type=str, help='Path to CA certificate')
    parser.add_argument('--cert', type=str, help='Path to client certificate')
    args = parser.parse_args()

    # MacOS version of the toolchain does not have python linked with OpenSSL 1.1.1 yet, so we monkey patch this in here
    if not hasattr(ssl, 'OP_NO_TLSv1_3'):
        ssl.OP_NO_TLSv1_3 = 0

    exclude_ops = {
        ssl.OP_NO_SSLv2,
        ssl.OP_NO_SSLv3,
        ssl.OP_NO_TLSv1,
        ssl.OP_NO_TLSv1_1,
        ssl.OP_NO_TLSv1_2,
        ssl.OP_NO_TLSv1_3,
    }

    def exclude_except(op):
        option = 0
        for other_op in exclude_ops - {op}:
            option |= other_op
        return option

    suites = {
        'sslv2': exclude_except(ssl.OP_NO_SSLv2),
        'sslv3': exclude_except(ssl.OP_NO_SSLv3),
        'tls1': exclude_except(ssl.OP_NO_TLSv1),
        'tls1_1': exclude_except(ssl.OP_NO_TLSv1_1),
        'tls1_2': exclude_except(ssl.OP_NO_TLSv1_2),
    }

    results = {
        key: enumerate_tls_ciphers(protocol_options=proto, host=args.host, port=args.port,
                                   cafile=args.cafile, cert=args.cert)
        for key, proto in suites.items()
    }

    if exception_ciphers:
        print("System could not process the following ciphers")
        for cipher, error in exception_ciphers.items():
            print(cipher + '\tError: ' + error)

    with open(args.outfile, 'w+') as outfile:
        json.dump(results, outfile)