#include "Server.h"

#include <openssl/err.h>
#include <openssl/ssl.h>
#ifndef WIN32
#    include <string.h>
#endif
#include <iostream>
#include <Key.h>
#include <Logging.h>
#include <Text.h>

using namespace Network;

// Content of the Server class from Server.h
// Constructor
Server::Server()
    : ReferenceCounter()
{
    sock = 0;
    memset(&address, 0, sizeof(address)); // Set address
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = ADDR_ANY;
    clients = 0;
}

// Destructor
Server::~Server()
{
    disconnect();
}

// non-constant
bool Server::connect(unsigned short port, int queueLength) // Opens the socket
{
    sock = socket(AF_INET, SOCK_STREAM, 0); // Create socket
    address.sin_port = htons(port);        // set port
    if (sock < 0)
    {
        sock = 0;
        return 0;
    }
#ifdef WIN32
    char reuseSocket = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuseSocket, sizeof(char));
#else
    int reuseSocket = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuseSocket, sizeof(int));
#endif
    if (bind(sock, (struct sockaddr*)&address, sizeof(address))
        == -1) // open socket
    {
        disconnect();
        return 0; // Error
    }
    if (listen(sock, queueLength) == -1) // accept clients
    {
        disconnect();
        return 0; // Error
    }
    return 1;
}

SClient* Server::getClient() // accepts a client
{
    if (!sock) return 0;
    sockaddr_in client;
    int len = sizeof(address);
    fd_set set;
    int rv = 0;
    struct timeval timeout;
    while (rv == 0 && sock)
    {
        FD_ZERO(&set);      /* clear the set */
        FD_SET(sock, &set); /* add our file descriptor to the set */
        timeout.tv_sec = 10;
        timeout.tv_usec = 0;
        rv = select((int)sock + 1, &set, NULL, NULL, &timeout);
        if (rv == -1) return 0;
    }
    if (!sock) return 0;
#ifdef WIN32
    SOCKET cls = accept(sock, (sockaddr*)&client, &len); // receive client
    if (cls == INVALID_SOCKET)
    {
        disconnect();
        return 0;
    }
#else
    SOCKET cls = accept(
        sock, (sockaddr*)&client, (socklen_t*)&len); // receive client
    if (!cls)
    {
        if (errno == ECONNABORTED || errno == EBADF) disconnect();
        return 0;
    }
#endif
    client.sin_port = address.sin_port;
    clients++;
    return new SClient(client, cls); // return client handle class
}

int Server::getClients(bool reset) // returns the number of clients
{
    int ret = clients;
    if (reset) clients = 0;
    return ret;
}

bool Server::disconnect() // stops the server
{
    if (!sock) return 1;
    if (closesocket(sock) < 0) // close socket
        return 0;
    sock = 0;
    return 1;
}

// constant
unsigned short Server::getPort() const // returns the port
{
    return htons(address.sin_port);
}

bool Server::isConnected()
    const // returns 1 if the server is connected
{
    return sock != 0;
}

// Content of the SClient class from Server.h
// Constructor
SClient::SClient(sockaddr_in address, SOCKET sock)
    : ReferenceCounter()
{
    clientAddr = address;
    this->sock = sock;
    downStreamBytes = 0;
    upStreamBytes = 0;
    sendKey = 0;
    receiveKey = 0;
}

// Destructor
SClient::~SClient()
{
    disconnect();
    if (sendKey) sendKey->release();
    if (receiveKey) receiveKey->release();
}

// non-constant
void SClient::setReceiveTimeout(
    int miliseconds) // Sets a timeout for receiving data
{
#ifdef WIN32
    DWORD timeout = miliseconds;
    setsockopt(
        sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof timeout);
#else
    struct timeval tv;
    tv.tv_sec = miliseconds / 1000;
    tv.tv_usec = (miliseconds % 1000) * 1000;
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof tv);
#endif
}

void SClient::setSendKeyZ(Encryption::Key* key) // Sets the key for sending
{
    if (sendKey) sendKey->release();
    sendKey = key;
}

void SClient::setReceiveKeyZ(
    Encryption::Key* key) // Sets the key for receiving
{
    if (receiveKey) receiveKey->release();
    receiveKey = key;
}

void SClient::setSendKey(const char* key, int len) // Sets the key for sending
{
    if (!sendKey) sendKey = new Encryption::Key();
    sendKey->setKey(key, len);
}

void SClient::setReceiveKey(
    const char* key, int len) // Sets the key for receiving
{
    if (!receiveKey) receiveKey = new Encryption::Key();
    receiveKey->setKey(key, len);
}

bool SClient::send(const char* message, int len) // sends to client
{
    if (!sock) return 0;
    int ll = 0;
    while (len > 0)
    {
#ifdef WIN32
        int l = send(sock, message + ll, len, 0);
#else
        int l = (int)send(sock, message + ll, len, MSG_NOSIGNAL);
#endif
        if (l <= 0)
        {
#ifdef WIN32
#    ifdef _DEBUG
            Framework::Logging::warning()
                << "send: " << l << " Error: " << WSAGetLastError();
#    endif
#endif
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    upStreamBytes += ll;
    return 1;
}

bool SClient::getMessage(char* message, int len) // receives message from client
{
    if (!sock) return 0;
    int ll = 0;
    while (len > 0)
    {
        int l = (int)recv(sock, message + ll, len, MSG_WAITALL);
        if (l <= 0)
        {
#ifdef WIN32
#    ifdef _DEBUG
            Framework::Logging::warning()
                << "recv: " << l << " Error: " << WSAGetLastError();
#    endif
#endif
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    downStreamBytes += ll;
    return 1;
}

bool SClient::sendEncrypted(const char* message, int len) // sends to server
{
    if (!sendKey) return send(message, len);
    Encryption::Bytes* n = new Encryption::Bytes(message, len);
    sendKey->codieren(dynamic_cast<Encryption::Bytes*>(n->getThis()));
    int ll = 0;
    while (len > 0)
    {
#ifdef WIN32
        int l = send(sock, n->getBytes() + ll, len, 0);
#else
        int l = (int)send(sock, n->getBytes() + ll, len, MSG_NOSIGNAL);
#endif
        if (l <= 0)
        {
#ifdef WIN32
#    ifdef _DEBUG
            Framework::Logging::warning()
                << "send: " << l << " Error: " << WSAGetLastError();
#    endif
#endif
            n->release();
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    upStreamBytes += ll;
    n->release();
    return 1;
}

bool SClient::getMessageEncrypted(char* message, int len) // receives message
{
    if (!receiveKey) return getMessage(message, len);
    int ll = 0;
    while (len > 0)
    {
        int l = (int)recv(sock, message + ll, len, MSG_WAITALL);
        if (l <= 0)
        {
#ifdef WIN32
#    ifdef _DEBUG
            Framework::Logging::warning()
                << "recv: " << l << " Error: " << WSAGetLastError();
#    endif
#endif
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    Encryption::Bytes* n = new Encryption::Bytes();
    n->setBytesZ(message, ll);
    receiveKey->decodieren(n);
    downStreamBytes += ll;
    return 1;
}

int SClient::getDownloadBytes(
    bool reset) // returns the number of received bytes
{
    int ret = downStreamBytes;
    if (reset) downStreamBytes = 0;
    return ret;
}

int SClient::getUploadBytes(
    bool reset) // returns the number of sent bytes
{
    int ret = upStreamBytes;
    if (reset) upStreamBytes = 0;
    return ret;
}

bool SClient::disconnect() // disconnects from client
{
    if (!sock) return 0;
    if (closesocket(sock) < 0) // disconnect
        return 0;
    sock = 0;
    return 1;
}

// constant
bool SClient::hasMessage(
    int time) const // Waits for a message for a given time
{
    fd_set set;
    FD_ZERO(&set);
    FD_SET(sock, &set);
    timeval time = {time / 1000, time};
    int result = select(0, &set, 0, 0, &time);
    if (result < 0)
    {
#ifdef WIN32
#    ifdef _DEBUG
        Framework::Logging::warning()
            << "select: " << result << " Error: " << WSAGetLastError();
#    endif
#endif
    }
    return result > 0;
}

unsigned short SClient::getPort() const // returns the port
{
    return htons(clientAddr.sin_port);
}

const char* SClient::getIp() const // returns the client's IP
{
    return inet_ntoa(clientAddr.sin_addr);
}

int pem_passwd_cb(char* buf, int size, int rwflag, void* userdata)
{
    const char* pw = ((Text*)userdata)->getText();
    memcpy(buf, pw, MIN((unsigned int)size, strlen(pw) + 1));
    return (int)strlen(buf);
}

bool SSLErrorCheck(int result, SSL* ssl, const char* action)
{
    if (result <= 0)
    {
        Framework::Logging::error()
            << "'" << action
            << "' returned error code: " << SSL_get_error(ssl, result);
        return 0;
    }
    return 1;
}

bool SSLErrorCheck(__int64 result, const char* action)
{
    if (result <= 0)
    {
        unsigned long error = ERR_get_error();

        Framework::Logging::error() << "'" << action << "' returned " << result
                                    << " error code: " << error << "("
                                    << ERR_reason_error_string(error) << ")";
        return 0;
    }
    return 1;
}

bool SClient::waitForNextMessage()
    const // waits until there is something to receive
{
    fd_set set;
    int rv = 0;
    struct timeval timeout;
    while (rv == 0 && sock)
    {
        FD_ZERO(&set);      /* clear the set */
        FD_SET(sock, &set); /* add our file descriptor to the set */
        timeout.tv_sec = 10;
        timeout.tv_usec = 0;
        rv = select((int)sock + 1, &set, NULL, NULL, &timeout);
        if (rv == -1)
        {
#ifdef WIN32
#    ifdef _DEBUG
            Framework::Logging::warning()
                << "select: " << rv << " Error: " << WSAGetLastError();
#    endif
#endif
            return 0;
        }
    }
    if (!sock) return 0;
    char c;
#ifdef WIN32
    int l = (int)recv(sock, &c, 1, MSG_PEEK);
#else
    int l = (int)recv(sock, &c, 1, MSG_WAITALL | MSG_PEEK);
#endif
    if (l <= 0)
    {
#ifdef WIN32
#    ifdef _DEBUG

        Framework::Logging::warning()
            << "recv: " << l << " Error: " << WSAGetLastError();
#    endif
#endif
        return 0; // Error
    }
    return 1;
}

// Content of the SSLServer class
// Constructor
SSLServer::SSLServer()
    : ReferenceCounter()
{
    s = 0;
    const SSL_METHOD* method = TLS_server_method();
    ctx = SSL_CTX_new(method);
    SSLErrorCheck(SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION),
        "SSL_CTX_set_min_proto_version");
    SSLErrorCheck(SSL_CTX_set_max_proto_version(ctx, TLS1_3_VERSION),
        "SSL_CTX_set_max_proto_version");
    SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, 0);
    SSL_CTX_set_default_passwd_cb(ctx, pem_passwd_cb);
    password = new Text();
    SSL_CTX_set_default_passwd_cb_userdata(ctx, password);
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    clients = 0;
}

// Destructor
SSLServer::~SSLServer()
{
    disconnect();
    SSL_CTX_free(ctx);
    password->release();
#ifdef WIN32
    OPENSSL_thread_stop();
#endif
}

// non-constant
// Sets the path to the file where the certificate is stored
bool SSLServer::setCertificateFile(const char* file)
{
    return SSLErrorCheck(
        SSL_CTX_use_certificate_file(ctx, file, SSL_FILETYPE_PEM),
        "SSL_CTX_use_certificate_file");
}

// Sets the path to the file where the private key is stored
bool SSLServer::setPrivateKeyFile(const char* file)
{
    return SSLErrorCheck(
        SSL_CTX_use_PrivateKey_file(ctx, file, SSL_FILETYPE_PEM),
        "SSL_CTX_use_PrivateKey_file");
}

// sets the password of the private key (must be called before
// setPrivateKeyFile)
void SSLServer::setPrivateKeyPassword(const char* password)
{
    this->password->setText(password);
}

// Opens the socket
bool SSLServer::connect(unsigned short port, int queueLength)
{
    addr.sin_port = htons(port);
    s = socket(AF_INET, SOCK_STREAM, 0);
    if (s < 0)
    {
        s = 0;
        return 0;
    }
#ifdef WIN32
    char reuseSocket = 1;
    setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &reuseSocket, sizeof(char));
#else
    int reuseSocket = 1;
    setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &reuseSocket, sizeof(int));
#endif
    if (bind(s, (struct sockaddr*)&addr, sizeof(addr)) < 0)
    {
        disconnect();
        return 0;
    }
    if (listen(s, queueLength) < 0)
    {
        disconnect();
        return 0;
    }
    return 1;
}

// accepts a client
SSLSClient* SSLServer::getClient()
{
    if (!s) return 0;
    int len = sizeof(addr);
    struct sockaddr_in addr;
    fd_set set;
    int rv = 0;
    struct timeval timeout;
    while (rv == 0 && s)
    {
        FD_ZERO(&set);   /* clear the set */
        FD_SET(s, &set); /* add our file descriptor to the set */
        timeout.tv_sec = 10;
        timeout.tv_usec = 0;
        rv = select((int)s + 1, &set, NULL, NULL, &timeout);
        if (rv == -1) return 0;
    }
    if (!s) return 0;
#ifdef WIN32
    SOCKET client = accept(s, (struct sockaddr*)&addr, &len);
    if (client == INVALID_SOCKET)
    {
        disconnect();
        return 0;
    }
#else
    SOCKET client
        = accept(s, (sockaddr*)&addr, (socklen_t*)&len); // receive client
    if (!client)
    {
        if (errno == ECONNABORTED || errno == EBADF) disconnect();
        return 0;
    }
#endif
    addr.sin_port = this->addr.sin_port;
    SSL* ssl = SSL_new(ctx);
    if (ssl == 0 && !SSLErrorCheck(0, "SSL_new"))
    {
        closesocket(client);
        return 0;
    }
    if (!SSLErrorCheck(SSL_set_fd(ssl, (int)client), ssl, "SSL_set_fd"))
    {
        SSL_free(ssl);
        closesocket(client);
        return 0;
    }
    if (!SSLErrorCheck(SSL_accept(ssl), ssl, "SSL_accept"))
    {
        SSL_free(ssl);
        closesocket(client);
        return 0;
    }
    clients++;
    return new SSLSClient(addr, ssl, client);
}

// returns the number of clients
int SSLServer::getClients(bool reset)
{
    int ret = clients;
    if (reset) clients = 0;
    return ret;
}

// stops the server
bool SSLServer::disconnect()
{
    if (!s) return 1;
    if (closesocket(s) < 0) // close socket
        return 0;
    s = 0;
    return 1;
}

// constant
// returns the port
unsigned short SSLServer::getPort() const
{
    return htons(addr.sin_port);
}

// returns 1 if the server is connected
bool SSLServer::isConnected() const
{
    return s != 0;
}

// Content of the SSLSClient class
// Constructor
SSLSClient::SSLSClient(sockaddr_in client, SSL* ssl, SOCKET s)
    : ReferenceCounter()
{
    this->s = s;
    clientAddr = client;
    this->ssl = ssl;
    downStreamBytes = 0;
    upStreamBytes = 0;
}

// Destructor
SSLSClient::~SSLSClient()
{
    disconnect();
#ifdef WIN32
    OPENSSL_thread_stop();
#endif
}

// non-constant
void SSLSClient::setReceiveTimeout(
    int miliseconds) // Sets a timeout for receiving data
{
#ifdef WIN32
    DWORD timeout = miliseconds;
    setsockopt(
        s, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof timeout);
#else
    struct timeval tv;
    tv.tv_sec = miliseconds / 1000;
    tv.tv_usec = (miliseconds % 1000) * 1000;
    setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof tv);
#endif
}

bool SSLSClient::send(const char* message, int len) // sends to client
{
    if (!ssl) return 0;
    int ll = 0;
    while (len > 0)
    {
        int l = SSL_write(ssl, message + ll, len);
        if (l <= 0)
        {
#ifdef _DEBUG
            Framework::Logging::warning()
                << "SSL_write: " << l << " Error: " << SSL_get_error(ssl, l);
#endif
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    upStreamBytes += ll;
    return 1;
}

bool SSLSClient::getMessage(char* message, int len) // receives message from client
{
    if (!ssl) return 0;
    int ll = 0;
    while (len > 0)
    {
        int l = (int)SSL_read(ssl, message + ll, len);
        if (l <= 0)
        {
#ifdef _DEBUG
            Framework::Logging::warning()
                << "SSL_read: " << l << " Error: " << SSL_get_error(ssl, l);
#endif
            return 0; // Error
        }
        len -= l;
        ll += l;
    }
    downStreamBytes += ll;
    return 1;
}

int SSLSClient::getDownloadBytes(
    bool reset) // returns the number of received bytes
{
    int ret = downStreamBytes;
    if (reset) downStreamBytes = 0;
    return ret;
}

int SSLSClient::getUploadBytes(
    bool reset) // returns the number of sent bytes
{
    int ret = upStreamBytes;
    if (reset) upStreamBytes = 0;
    return ret;
}

bool SSLSClient::disconnect() // disconnects from client
{
    if (!ssl) return 0;
    SSL_free(ssl);
    if (closesocket(s) < 0) // disconnect
        return 0;
    ssl = 0;
    s = 0;
    return 1;
}

// constant
bool SSLSClient::hasMessage(
    int time) const // Waits for a message for a given time
{
    fd_set set;
    FD_ZERO(&set);
    FD_SET(SSL_get_rfd(ssl), &set);
    timeval time = {time / 1000, time};
    return SSL_pending(ssl) > 0 || select(0, &set, 0, 0, &time) == 1;
}

unsigned short SSLSClient::getPort() const // returns the port
{
    return htons(clientAddr.sin_port);
}

const char* SSLSClient::getIp() const // returns the client's IP
{
    return inet_ntoa(clientAddr.sin_addr);
}