#include <netinet/tcp.h>
#include <sys/socket.h>
#include <arpa/inet.h>

// include/ path is here for a reason - to make sure we compile against boringssl (temporary solution)
#include <include/openssl/rand.h>
#include <include/openssl/ssl.h>
#include <include/openssl/x509v3.h>
#include <pthread.h>

#include "common.h"

// Enforce this protocol version
#define TLS_PROT_VERSION    TLS1_3_VERSION
static const int Curves[3] = {NID_CECPQ2, NID_CECPQ2b, NID_X25519};

static const struct CertDesc_t {
    const char* arg;
    const char* cert;
    const char* key;
    const char* cipher;
    const int* curves;
} Certs[] = {
    {
        .arg = "RSA_2048",
        .key = "etc/rsa_2048.key",
        .cert = "etc/rsa_2048.pem",
        .cipher = "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
        .curves = Curves,
    },
    {
        .arg = "RSA_3072",
        .key = "etc/rsa_3072.key",
        .cert = "etc/rsa_3072.pem",
        .cipher = "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
        .curves = Curves,
    },
    {
        .arg = "ECDSA_256",
        .key = "etc/ecdsa_256.key",
        .cert = "etc/ecdsa_256.pem",
        .cipher = "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
        .curves = Curves,
    }
};

const struct CertDesc_t* get_cert_desc(const char* c) {
    for(size_t i=0; i<ARRAY_SIZE(Certs); i++) {
        if (!memcmp(Certs[i].arg, c, strlen(Certs[i].arg))) {
            return &Certs[i];
        }
    }
    ERR("Critical error");
}

// Buffer used for read/write tests
unsigned char rw_buf[BUFFER_SIZE];

int pass_cb(char *buf, int size, int rwflag, void *u) {
    DBG("Using password " KEYFILE_PASS);
    // password for all keys is "test123"
    memcpy(buf, KEYFILE_PASS, strlen(KEYFILE_PASS));
    return strlen(KEYFILE_PASS);
}

static int accept_once(void) {

	int fd, client;
	int reuseval = 1;
    struct sockaddr_in a;
    size_t a_len = sizeof(a);

    // configure
    memset(&a, 0, a_len);
    a.sin_family = AF_INET;
    a.sin_port = htons(PORT);
    a.sin_addr.s_addr = INADDR_ANY;

	fd = socket(PF_INET, SOCK_STREAM, 0);

    if (setsockopt(fd,SOL_SOCKET, SO_REUSEADDR, &reuseval, sizeof(reuseval))) {
    	ERR("setsockopt");
    }

    if (bind(fd, (struct sockaddr *)&a, sizeof(a))) {
	    ERR("bind");
    }

    if (listen(fd, 1)) {
        ERR("listen");
    }


    client = accept(fd, (struct sockaddr *)&a, (socklen_t *)&a_len);
    if (client<0) {
        ERR("accept");
    }
    return client;
}

static SSL_CTX* setup_server_ctx(const char* cert_name) {
	SSL_CTX* ctx = SSL_CTX_new(TLS_method());
    assert(ctx);

    const struct CertDesc_t *c = get_cert_desc(cert_name);
    SSL_CTX_set_default_passwd_cb(ctx, &pass_cb);

	if(SSL_CTX_use_certificate_file(ctx, c->cert, SSL_FILETYPE_PEM) != 1)
		ERR("Error loading certificate from file");
	if(SSL_CTX_use_PrivateKey_file(ctx, c->key, SSL_FILETYPE_PEM) != 1)
		ERR("Error loading private key from file");

    if (SSL_CTX_set_strict_cipher_list(ctx, c->cipher) != 1) {
        ERR("Error setting cipher list");
    }

    if (!SSL_CTX_set1_curves(ctx, c->curves, 3)) {
        ERR("Enforcing curve");
    }

    if (SSL_CTX_set_min_proto_version(ctx, TLS_PROT_VERSION) != 1 ||
        SSL_CTX_set_max_proto_version(ctx, TLS_PROT_VERSION) != 1) {
        ERR("Enforcing protocol to TLSv1.2");

    }
    return ctx;
}

static void start_writing(SSL *ssl) {
    int nread,err;

    DBG("Start writing");

    fill_buffer_from_file(rw_buf, sizeof(rw_buf));
    for (size_t i=0; i<RW_ITERATIONS; i++) {
        for(nread=0; nread<sizeof(rw_buf); nread+=err) {
            err = SSL_write(ssl, rw_buf+nread, sizeof(rw_buf) - nread);
            if(err <= 0) {
                ERR("SSL_write");
            }
        }
    }
}

static int do_server_loop(SSL* ssl) {
    int nread,err;

    DBG("Check if first packet is ReadRequest");
    for(nread=0; nread<strlen(READ_REQUEST); nread+=err) {
        err = SSL_read(ssl, rw_buf+nread, sizeof(rw_buf) - nread);
        if (err <=0) {
            break;
        }
    }

    DBG("Check if server loop should be finished");
    // Stop here if shutdown requested
    if (SSL_get_shutdown(ssl) & SSL_RECEIVED_SHUTDOWN) {
        return 1;
    }

    // Check if client wants to read
    if (nread && !memcmp(READ_REQUEST, rw_buf, nread)) {
        start_writing(ssl);
        return 2; // server initiated closure
    }

    // Keep reading till connection closed
    DBG("Continue reading");
	do {
		for(nread=0; nread<sizeof(rw_buf); nread+=err) {
            err = SSL_read(ssl, rw_buf+nread, sizeof(rw_buf) - nread);
			if(err <= 0) {
				break;
            }
		}
	} while(err>0);
	return (SSL_get_shutdown(ssl) & SSL_RECEIVED_SHUTDOWN) ? 1 : 0;
}

#ifdef DEBUG
static const char* DBG_get_curve(SSL* ssl) {
    const int id = SSL_get_curve_id(ssl);
    if (!id) return "N/A";
    return SSL_get_curve_name(id);
}
#endif

// Interface for accepting
int main(int argc, char *argv[])
{
	SSL     *ssl;
	SSL_CTX *ctx;
	int 	fd, ret;

    if (argc < 2) {
        goto usage;
    }

    // Start server
    init();

    ctx = setup_server_ctx(argv[1]);
    SSL_CTX_enable_pq_experiment_signal(ctx);
	fd = accept_once();

	for(;;) {

		DBG("SSL Accepting");
		if(!(ssl = SSL_new(ctx)))	{
			DBG("Error creating SSL context");
		}

		SSL_set_fd(ssl, fd);
        ret = SSL_accept(ssl);
		if (ret<=0)	{
            ret = SSL_get_error(ssl, ret);

            if (ret == SSL_ERROR_SYSCALL) {
                DBG("Connection closed");
                goto end;
            } else {
                ERR("Critical error occured %d", ret);
            }
		}

        DBG("Protocol: %s Cipher %s Curve %s", SSL_get_version(ssl), SSL_get_cipher_name(ssl), DBG_get_curve(ssl));

		DBG("server loop");
        ret = do_server_loop(ssl);
		if(ret) {

			DBG("Shutdown");
			SSL_set_read_ahead(ssl, 0);
		    if (ret==2) {
                // Server is closing
                SSL_shutdown(ssl);
            }
            assert(SSL_shutdown(ssl)==1);
		}
		else {
			DBG("Shutdown - error case");
			goto end;
		}
        SSL_clear(ssl);
	}
    SSL_CTX_free(ctx);

end:
    cleanup();
	return 0;
usage:
    ERR("\n\nUsage: %s server certificate \nOptions for 'server certificate':\n\t"
        "RSA_2048\n\t"
        "RSA_3072\n\t"
        "ECDSA_256\n\t",
        argv[0]);
    goto end;
}