#include "ssl.h"


SSL_CTX* setup_ctx(CertList *certItem) {
	const SSL_METHOD *method;
	SSL_CTX *ctx;
	OpenSSL_add_all_algorithms();  /* load & register all cryptos, etc. */
	SSL_load_error_strings();   /* load all error messages */
	method = TLS_server_method();  /* create new server-method instance */


	ctx = SSL_CTX_new(method);   /* create new context from method */
	if ( ctx == NULL ) {
		ERR_print_errors_fp(stderr);
		abort();
	}

	//set the local certificate from CertFile
	if ( SSL_CTX_use_certificate(ctx, certItem->cert) <= 0 ){
		ERR_print_errors_fp(stderr);
		abort();
	}
	
	//set the private key from KeyFile (may be the same as CertFile)
	if ( SSL_CTX_use_PrivateKey(ctx, certItem->key) <= 0 ){
		ERR_print_errors_fp(stderr);
		abort();
	}
	//verify private key
	if ( !SSL_CTX_check_private_key(ctx) ) {
		fprintf(stderr, "Private key does not match the public certificate\n");
		abort();
	}

	return ctx;
}



// Generates a 2048-bit RSA key.
// Largely stolen from here: https://gist.github.com/nathan-osman/5041136
EVP_PKEY* generate_ca_key() {
	// Allocate memory for the EVP_PKEY structure.
	EVP_PKEY *pkey = EVP_PKEY_new();
	if(!pkey) {
		perror("Unable to create EVP_PKEY structure.\n");
		return NULL;
	}

	// Generate the RSA key and assign it to pkey.
	RSA *rsa = RSA_new();
	BIGNUM *bn = BN_new();
	BN_set_word(bn, RSA_F4);
	RSA_generate_key_ex(rsa, 2048, bn, NULL );
	if(!EVP_PKEY_assign_RSA(pkey, rsa)) {
		perror("Unable to generate 2048-bit RSA key.\n");
		EVP_PKEY_free(pkey);
		return NULL;
	}

    // The key has been generated, return it.
    return pkey;
}

// Generates a self-signed x509 certificate.
// Largely stolen from here: https://gist.github.com/nathan-osman/5041136
X509* generate_ca_cert(EVP_PKEY * pkey) {
	// Allocate memory for the X509 structure.
	X509 *x509 = X509_new();
	if(!x509) {
		perror( "Unable to create X509 structure.\n");
		return NULL;
	}

	
	ASN1_INTEGER_set(X509_get_serialNumber(x509), 1);

	// This certificate is valid from now until exactly one year from now.
	X509_gmtime_adj(X509_get_notBefore(x509), 0);
	X509_gmtime_adj(X509_get_notAfter(x509), 31536000L);

	// Set the public key for our certificate.
	X509_set_pubkey(x509, pkey);

	// We want to copy the subject name to the issuer name.
	X509_NAME *name = X509_get_subject_name(x509);

	// Set the country code and common name.
	X509_NAME_add_entry_by_txt(name, "C",  MBSTRING_ASC, (unsigned char *)"UK",        -1, -1, 0);
	X509_NAME_add_entry_by_txt(name, "O",  MBSTRING_ASC, (unsigned char *)"Yet Another Intercepting Proxy", -1, -1, 0);
	//I want this to be a fake domain because CAs can't sign certificates with
	//the same domain. I thought about using jn.hn but I want to test using
	//that.
	X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC, (unsigned char *)"yaip.yaip", -1, -1, 0);

	// Now set the issuer name.
	X509_set_issuer_name(x509, name);


	// TODO: mitmproxy adds some extensions. I don' tknow if I need those
	// https://github.com/mitmproxy/mitmproxy/blob/3cb89069b9f70c198cc89337bd5a9b1bbf3a69d0/mitmproxy/certs.py#L190-L204


	// Actually sign the certificate with our key.
	if(!X509_sign(x509, pkey, EVP_sha1())) {
		perror( "Error signing certificate.\n");
		X509_free(x509);
		return NULL;
	}



	return x509;
}

X509* generate_site_cert( EVP_PKEY *pkey, X509 *caCert, char *host ){
	X509 *hostCert = X509_new();
	if(!hostCert) {
		perror( "Unable to create X509 structure.\n");
		return NULL;
	}

	
	ASN1_INTEGER_set(X509_get_serialNumber(hostCert), 1);

	// This certificate is valid from now until exactly one year from now.
	X509_gmtime_adj(X509_get_notBefore(hostCert), 0);
	X509_gmtime_adj(X509_get_notAfter(hostCert), 31536000L);

	// Set the public key for our certificate.
	X509_set_pubkey(hostCert, pkey);

	// We want to copy the subject name of the ca to the issuer of the new
	// certificate
	X509_NAME *ca_sn = X509_get_subject_name(caCert);
	X509_NAME *host_in = X509_NAME_dup( ca_sn );
	X509_set_issuer_name(hostCert, host_in);


	X509_NAME *host_sn = X509_get_subject_name(hostCert);

	// Set the country code and common name.
	X509_NAME_add_entry_by_txt(host_sn, "C",  MBSTRING_ASC, (unsigned char *)"UK",        -1, -1, 0);
	X509_NAME_add_entry_by_txt(host_sn, "O",  MBSTRING_ASC, (unsigned char *)"Yet Another Intercepting Proxy", -1, -1, 0);
	X509_NAME_add_entry_by_txt(host_sn, "CN", MBSTRING_ASC, (unsigned char *)strdup(host), -1, -1, 0);

	if(!X509_sign(hostCert, pkey, EVP_sha1())) {
		perror( "Error signing certificate.\n");
		X509_free(hostCert);
		return NULL;
	}

	return hostCert;
}

bool create_and_save_key(char keyfile[]) {
	// Generate the key
	EVP_PKEY *pkey = generate_ca_key();
	if ( pkey == NULL ){
		fprintf(stderr, "Unable to open generate key\n");
		return false;
	}

	// Open the PEM file for writing the key to disk.
	FILE * pkey_file = fopen(keyfile, "wb");
	if(!pkey_file) {
		fprintf(stderr, "Unable to open \"%s\" for writing.\n", keyfile);
		return false;
	}
    // Write the key to disk.
    bool ret = PEM_write_PrivateKey(pkey_file, pkey, NULL, NULL, 0, NULL, NULL);
    fclose(pkey_file);

    if(!ret) {
        perror("Unable to write private key to disk.\n");
        return false;
    }
	return true;
}

bool create_and_save_cert(char certfile[], EVP_PKEY *pkey) {
    // Open the PEM file for writing the certificate to disk.
    FILE * x509_file = fopen(certfile, "wb");
    if(!x509_file) {
		fprintf(stderr, "Unable to open \"%s\" for writing.\n", certfile);
        return false;
    }

	X509 *x509 = generate_ca_cert(pkey);
	if ( x509 == NULL ){
		fprintf(stderr, "Unable to open generate cert\n");
		return false;
	}

    // Write the certificate to disk.
    int ret = PEM_write_X509(x509_file, x509);
    fclose(x509_file);

    if(!ret) {
        perror("Unable to write certificate to disk.\n");
        return false;
    }
	return true;
}

EVP_PKEY* read_private_key(char keyfile[]){
	FILE *fp;
	EVP_PKEY *pkey;

	if (! (fp = fopen(keyfile, "r"))){
		perror("Error cant open certificate private key file.\n");
		return NULL;
	}

	pkey = PEM_read_PrivateKey( fp, NULL, NULL, NULL );

	if ( pkey == NULL ){
		perror("Error cant read certificate private key file.\n");
	}
	fclose(fp);

	return pkey;
}

X509* read_cert(char certfile[]){
	FILE *fp;
	X509 *cert;

	if (! (fp = fopen(certfile, "r"))){
		perror("Error cant open certificate private key file.\n");
		return NULL;
	}

	cert = PEM_read_X509( fp, NULL, NULL, NULL );

	if ( cert == NULL ){
		perror("Error cant read certificate private key file.\n");
	}
	fclose(fp);

	return cert;
}

CertList *newCertListItem(char *host, EVP_PKEY *key, X509 *cert){
	CertList *item = malloc(sizeof( CertList ));
	item->host = host;
	item->key = key;
	item->cert = cert;
	item->next = NULL;
	return item;
}

CertList *getLastCertListItem(CertList *item){
	if ( item == NULL ) return NULL;
	while ( item->next != NULL ) item=item->next;
	return item;
}

unsigned int countCertListItems(CertList *item){
	unsigned int count = 0;
	while ( item != NULL ){
		count++;
		item = item->next;
	}
	return count;
}

CertList *findCertListItem( CertList *item, char *host ){
	while ( item != NULL && strcmp( host, item->host ) != 0 ) item = item->next;
	return item;
}