/*
  JS 2024-05-07 - doplnene priklady nastavenia premennej DEFAULT_GROUPS, doplnene info o ML-KEM
                - testovane s WinLibs 13.2.0 release 8
*/
/*
  JS 2024-04-11 testovane s OpenSSL 3.3.0, liboqs 0.10.0 a oqs-provider 0.6.0
*/
/*
  JS 2024-04-11 testovane s OpenSSL 3.3.0, liboqs 0.10.0 a oqs-provider 0.5.3
*/
/*
    JS 2024-03-04 - doplnenie funkcii pre uvolnenie nacitanych providerov
*/

/*
   JS 2024-02-24 - funkcia initSSLContext() rozsirena o nacitanie OQS-providera
		   pre plne funkcne PQ algoritmy na kazdej platforme
		 - doplnene priklady nastavenia premennej DEFAULT_GROUPS
		 - testovane s oqsprovider 0.5.3
*/
/*
   JS 2024-02-18 - testovanie PQ algoritmov s pouzitim oqs-providera
		 - uprava vypisu o pripojeni klienta
		 - pridane vypisy s informaciami o pouzivanych algoritmoch
		   pre KEX a certifikaty
		 - oprava kontroly navratovej hodnoty funkcii 
			SSL_CTX_use_certificate_file() a SSL_CTX_use_PrivateKey_file
*/

/*
  JS 2024-02-08 testovane s aktualnou najnovsou verziou OpenSSL 3.2.1
*/


/*
  MJ 2023-18-05 - Uprava funkcie printHeader()
                     -> Zmazanie SSL metody
                     -> Upraveny vypis ako spustit program
               - Zmazanie funkcii spojene s generovanim certifikatov
                    -> callbackGeneratingKey()
                    -> makekCert()
                - Uprava funkcie loadCertificates() aby nacitavala kluce klienta a servera zo suboru
                - Uprava funkcie initSSLContext() aby realizovala fixne iba TLS spojenie
                - Uprava hlavnej funkcie main() na zaklade predoslej upravy (nastavenie fixne TLS)
                    -> Zmazanie premennej ctxMethod
                    -> Uprava argumentov [argv] - zmazanie TLS metody, pridanie suborov s klucmi
                    -> zmena nazvoslovia premennych na ulozenie klucov


  MJ 2022-12-28 - Oprava Warningu vo funckii routine() vid tag MJ (vid. komentare MJ)

*/

/*

  MD 2021-03-29 testovane s aktualne najnovsou verziou OpenSSL v 1.1.1k
  MD 2018-11-07 Upravene pre linkovanie s OpenSSL v. 1.1.1
                (pouzitie novej funkcie TLS_server_method)
  MD 2018-11-02 Odstanene varovanie 
                dodane nacitanie CA certifikatu 
                dodany kod na realizaciu autentizacie klienta
                (vid //MD)
*/

/** SSL/TLS Server
 * SSL/TLS server demonstration. This source code is cross-plateforme Windows and Linux.
 * Compile under Linux with : g++ main.cpp -Wall -lssl -lcrypto -o main
 * Certificat and private key to protect transaction can be used from :
 * - External(s) file(s), created with command : openssl req -x509 -nodes -newkey rsa:2048 -keyout server.pem -out server.pem
 * - Internal uniq hardcoded certificat and private key, equal into each server instance
 * - Randomly generated certificat and private key, best solution to used dynamic keying material at each server lauching.
 */

// __unix__ is usually defined by compilers targeting Unix systems
#ifdef __unix__         
#   include <unistd.h>
#   include <sys/socket.h>
#   include <arpa/inet.h>
#   include <resolv.h>
#   define SOCKLEN_T socklen_t
#   define CLOSESOCKET close
    // _Win32 is usually defined by compilers targeting 32 or 64 bit Windows systems
#elif defined _WIN32    
    // MD, odstrani varovanie o potrebe prehodit poradie hlavickovych suborov
#   define WIN32_LEAN_AND_MEAN     
#   include <windows.h>
#   include <winsock2.h>
#   define SOCKLEN_T int
#   define CLOSESOCKET closesocket
#endif

#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <malloc.h>
#include <string.h>

#include <openssl/crypto.h>
#include <openssl/x509v3.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/bio.h>
#include <openssl/provider.h>

#ifdef _WIN32
WSADATA wsa;    // Winsock data
#endif

/* JS
Premenna na definovanie "exchange groups", algoritmov na vymenu klucov
Server moze podporovat viacero KEX/KEM algoritmov, v premennej musi byt kazdy
algoritmus oddeleny dvojbodkou, napr. "kyber512:bikel1"
Podporovane algoritmy je mozne ziskat prikazom "openssl list -kem-algorithms"
Prazdna premenna alebo naplnena neplatnymi protokolmi sposobi prerusenie aplikacie

Ak server a klient nenajdu spolocny KEX/KEM protokol, tak
SSL vyhodi chybu: "SSL routines:final_key_share:no suitable key share"

Priklad pouzitia:
- klasicke algoritmy: "x25519:x448:prime256v1:secp521r1:secp384r1:ffdhe2048:ffdhe3072"
- PQ algoritmy: "mlkem512:mlkem768:mlkem1024:kyber512:kyber768:kyber1024:bikel1:bikel3:bikel5:hqc128:hqc192:hqc256:frodo640aes:frodo640shake:frodo976aes:frodo976shake:frodo1344aes:frodo1344shake"
- hybrid algoritmy: "x25519_kyber768:x25519_frodo640aes:x25519_hqc128:x448_bikel3:x448_kyber768:p256_kyber768"

Podporovane algoritmy OQS-providerom: https://github.com/open-quantum-safe/oqs-provider/blob/0.6.0/ALGORITHMS.md 

POZOR - BIKE protokol nefunguje na Windows platforme (liboqs 0.10.0, oqs-provider 0.6.0)

*/
#define DEFAULT_GROUPS "mlkem512:kyber768:frodo976aes:kyber1024"

#define DEFAULT_PORT 443

/**
 * printUsage function who describe the utilisation of this script.
 * @param char* bin : the name of the current binary.
 */
void printHeader(char* bin){
    // JS update
    printf("[?] Usage : %s <port> <server_public_key> <server_private_key>]\n", bin);
    return;
}

/**
 * makeServerSocket function who create a traditionnal server socket, bind it and listen to it.
 * @param int port : the port to listen
 * @return int socket : the socket number created
 */
int makeServerSocket(int port){
    int                 sock;
    struct sockaddr_in  addr;
    #ifdef _WIN32
    WSAStartup(MAKEWORD(2,0),&wsa);
    #endif
    sock = socket(PF_INET, SOCK_STREAM, 0);
    memset(&addr, 0, sizeof(addr));
    addr.sin_family         = AF_INET;
    addr.sin_port           = htons(port);
    addr.sin_addr.s_addr    = INADDR_ANY;
    if(bind(sock, (struct sockaddr*)&addr, sizeof(addr)) != 0){
        perror("[-] Can't bind port on indicated port...");
        abort();
    }
    if(listen(sock, 10) != 0){
        perror("[-] Can't listening on indicated port...");
        abort();
    }
    printf("\n");
    printf("[+] Server listening on the %d port...\n", port);
    printf("[+] Waiting for connection\n");
    printf("\n");
    return sock;
}


/** JS update
 * initSSLContext function who initialize the SSL/TLS engine with right method/protocol
 * SSL/TLS engine provided by PQ algorithms functions of OQS provider
 * @return SSL_CTX *ctx : a pointer to the SSL context created
 */
SSL_CTX* initSSLContext(){
    const SSL_METHOD    *method;
    SSL_CTX             *ctx;

    // initialize the SSL library
    SSL_library_init();                     
    SSL_load_error_strings();			    
    OpenSSL_add_all_algorithms();		    

    // MJ only TLS connection
    method = TLS_server_method();
    //printf("[+] Use TLS server method.\n");

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

    return ctx;
}

/**
 * loadCertificates function who load private key and certificat from files.
 * 3 mecanisms available :
 * - loading certificate and private key from file(s)
 * - use embed hardcoded certificate and private key in the PEM format
 * - generate random and dynamic certificate and private key at each server's launch instance.
 * @param SSL_CTX* ctx : the SSL/TLS context
 * @param char *server_pem : filename of the PEM certificat
 * @param char *server_key : filename of the PEM private key
 */
void loadCertificates(SSL_CTX* ctx, const char* server_pem, const char* server_key){

    if (SSL_CTX_use_certificate_file(ctx, server_pem, SSL_FILETYPE_PEM) != 1 ||
        // namiesto MD SSL_CTX_use_RSAprivateKey_file
        SSL_CTX_use_PrivateKey_file(ctx, server_key, SSL_FILETYPE_PEM) != 1)
    { 
        ERR_print_errors_fp(stderr);
        abort();
    }
    else
        printf("[*] Server's certificat and private key loaded from file.\n");

    // verify private key match the public key into the certificate
    if(!SSL_CTX_check_private_key(ctx)){
        fprintf(stderr, "[-] Private key does not match the public certificate...\n");
        abort();
    } else
        printf("[+] Server's private key match public certificate\n");

    // JS get server's signature algorithm name
    X509 *cert = SSL_CTX_get0_certificate(ctx);
    int pknid;
    if (X509_get_signature_info(cert, NULL, &pknid, NULL, NULL) != 1){
    	printf("Certificate signature algorithm: Unknown algorithm");
    }
    else{
	printf("Certificate signature algorithm: %s\n", OBJ_nid2sn(pknid));
    }

    // MD  kod na autentizaciu klienta
    // MJ Update [Define Macro]
    #ifdef AUTHENTICATION

    SSL_CTX_set_verify(ctx, SSL_VERIFY_PEER|SSL_VERIFY_FAIL_IF_NO_PEER_CERT, NULL);

    // load the trusted client CA certificate into context
    if (SSL_CTX_load_verify_locations(ctx, "myCA.pem", NULL) != 1)
    {
        fprintf(stderr, "[-] CA certificate not loaded...\n");
        abort();
    }
    #endif

    return;
}

/**
 * showCerts function who catch and print out certificate's data from the client.
 * @param SSL* ssl : the SSL/TLS connection
 */
void showCerts(SSL* ssl){
    X509 *cert;
    char *subject, *issuer;
    
    // get the client's certificate
    cert = SSL_get_peer_certificate(ssl);   
    if(cert != NULL){
        // get certificate's subject
        subject = X509_NAME_oneline(X509_get_subject_name(cert), 0, 0); 
        // get certificate's issuer
        issuer  = X509_NAME_oneline(X509_get_issuer_name(cert), 0, 0);  

        printf("[+] Client certificates :\n");
        printf("\tSubject: %s\n", subject);
        printf("\tIssuer: %s\n", issuer);

        // Free memory
        free(subject);                      
        free(issuer);					   
        X509_free(cert);		
    }
    else
        printf("[-] No client's certificates\n");
    return;
}

/**
 * routine function who treat the content of data received and reply to the client.
 * this function is threadable and his context sharedable.
 * @param SSL* ssl : the SSL/TLS connection
 */
void routine(SSL* ssl){
    /* MJ char        buf[1024], reply[1024];
        -> buf must be smaller than reply array because arrays are copied in sprintf function, 
           which have some additional string
        -> buf_size + (string in sprintf)_size > reply_size = Can do overflow
        -> 1024 bytes (buf) + "Enchante %s, je suis ServerName.\n" + some string(s) > 1024 bytes 
          (reply)
        -> solution is change size of buf array
    */
    char buf[1024/2], reply[1024]; 
    
    int         sock, bytes;
    const char* echo = "Enchante %s, je suis ServerName.\n";

    // accept SSL/TLS connection
    if(SSL_accept(ssl) == -1)				    
	ERR_print_errors_fp(stderr);
    else{
        printf("[+] Cipher used : %s\n", SSL_get_cipher(ssl));

	// JS get chosen (negotiated) key exchange/encapsulation algorithm name
	printf("Used group (KEM): %s\n", SSL_group_to_name(ssl, SSL_get_negotiated_group(ssl)));

	// JS get client's certificate algorithm name
	int nid;
	SSL_get_peer_signature_type_nid(ssl, &nid);
	printf("Peer signature name: %s\n", OBJ_nid2sn(nid));


        // Show certificats data
        showCerts(ssl);

        // read data from client request								
        bytes = SSL_read(ssl, buf, sizeof(buf));	
        if(bytes > 0){
            buf[bytes] = 0;
            printf("[+] Client data received : %s\n", buf);

            // construct response
            sprintf(reply, echo, buf);
            // send response			    
            SSL_write(ssl, reply, strlen(reply));	
        } else {
            switch(SSL_get_error(ssl, bytes)){
                case SSL_ERROR_ZERO_RETURN :
                    printf("SSL_ERROR_ZERO_RETURN : ");
                    break;
                case SSL_ERROR_NONE :
                    printf("SSL_ERROR_NONE : ");
                    break;
                case SSL_ERROR_SSL:
                    printf("SSL_ERROR_SSL : ");
                    break;
            }
            ERR_print_errors_fp(stderr);
        }

    }

    // get traditionnal socket connection from SSL connection
    sock = SSL_get_fd(ssl);							
    // release SSL connection state
    SSL_shutdown(ssl);
    SSL_free(ssl);								    
    // close socket
    CLOSESOCKET(sock);                              
}

/**
 * main function who coordinate the socket and SSL connection creation, then receive and emit 
   data to and from the client.
 */
int main(int argc, char **argv){
    int sock, port;
    SSL_CTX *ctx;
    const char  *server_pem, *server_key;

    if(argc != 4){
        printHeader(argv[0]);
        exit(0);
    }

    port =  (atoi(argv[1]) > 0 && atoi(argv[1]) < 65535) ? atoi(argv[1]) : DEFAULT_PORT;

    // JS
    // Load default and OQS providers
    // Default provider must be loaded before OQS provider
    OSSL_PROVIDER* provider;
    provider = OSSL_PROVIDER_load(NULL, "default");
    if (provider == NULL) {
    	printf("Failed to load Default provider\n");
	    exit(0);
    }

    OSSL_PROVIDER* custom_provider = OSSL_PROVIDER_load(NULL, "oqsprovider");
    if (custom_provider == NULL){
	    printf("Failed to load OQS-provider\n");
        OSSL_PROVIDER_unload(provider);
	    exit(0);
    }

    // load SSL library and dependances
    ctx =  initSSLContext();     
    
    // MJ Update argv parameters
    server_pem = argv[2];
    server_key = argv[3];

    // load certificats and keys
    loadCertificates(ctx, server_pem, server_key);       
    
    // make a classic server socket
    sock = makeServerSocket(port);           

    while(1){
        struct sockaddr_in  addr;
        SSL         *ssl;
        SOCKLEN_T   len = sizeof(addr);

        // accept connection of client
        int         client = accept(sock, (struct sockaddr*)&addr, &len);   
        printf("[+] Connection [%s:%d]\n", inet_ntoa(addr.sin_addr), ntohs(addr.sin_port));
	
        // get new SSL state with context
        ssl = SSL_new(ctx);
	    if (ssl == NULL){
            ERR_print_errors_fp(stderr);
            break;
        }

        // JS set key exchange/encapsulation protocols supported by server
        // Without this function, client will use default X25519 protocol
        // First protocol supported by both server and client will be used for KEX/KEM
        // Empty DEFAULT_GROUPS or list of unsupported protocol by server will fail communication    
        if (SSL_set1_groups_list(ssl, DEFAULT_GROUPS) != 1){
            printf("KEX/KEM algorithms undefined - check DEFAULT_GROUPS variable\n");
            ERR_print_errors_fp(stderr);
            break;
        }

        // set traditionnal socket to SSL
        SSL_set_fd(ssl, client);		
        
        // apply routine to the socket's content			
        routine(ssl);
        // JS remove break if you want server running in loop
	    break;							    
    }


    // close socket
    CLOSESOCKET(sock);                              
    #ifdef _WIN32
    WSACleanup();
    #endif

    // release SSL's context
    SSL_CTX_free(ctx);
    // JS Unload both providers
	OSSL_PROVIDER_unload(provider);
	OSSL_PROVIDER_unload(custom_provider);								
    return 0;
}