//
// Created by WolverinDEV on 29/07/2020.
//

#include <log/LogUtils.h>
#include <misc/base64.h>
#include <misc/digest.h>
#include <src/InstanceHandler.h>
#include <misc/endianness.h>
#include "CryptSetupHandler.h"
#include "./VoiceClientConnection.h"

using namespace ts;
using namespace ts::connection;
using namespace ts::server::server::udp;

inline void generate_random(uint8_t *destination, size_t length) {
    while(length-- > 0)
        *(destination++) = (uint8_t) rand();
}

CryptSetupHandler::CryptSetupHandler(VoiceClientConnection *connection) : connection{connection} { }

CryptSetupHandler::CommandHandleResult CryptSetupHandler::handle_command(const std::string_view &payload) {
    std::variant<ts::command_result, CommandHandleResult>(CryptSetupHandler::*command_handler)(const ts::command_parser&) = nullptr;

    if(payload.starts_with("clientinitiv"))
        command_handler = &CryptSetupHandler::handleCommandClientInitIv;
    else if(payload.starts_with("clientek"))
        command_handler = &CryptSetupHandler::handleCommandClientEk;
    else if(payload.starts_with("clientinit"))
        command_handler = &CryptSetupHandler::handleCommandClientInit;

    if(!command_handler)
        return CommandHandleResult::PASS_THROUGH;

    this->last_command_ = std::chrono::system_clock::now();

    ts::command_parser parser{payload};
    try {
        std::unique_lock cmd_lock{this->command_lock};
        auto result = (this->*command_handler)(parser);

        CommandHandleResult handle_result;
        if(std::holds_alternative<CommandHandleResult>(result)) {
            handle_result = std::get<CommandHandleResult>(result);
        } else {
            auto cmd_result = std::move(std::get<ts::command_result>(result));

            ts::command_builder notify{"error"};
            cmd_result.build_error_response(notify, "id");

            if(parser.has_key("return_code"))
                notify.put_unchecked(0, "return_code", parser.value("return_code"));

            this->connection->send_command(notify.build(), false, nullptr);

            handle_result = cmd_result.has_error() ? CommandHandleResult::CLOSE_CONNECTION : CommandHandleResult::CONSUME_COMMAND;
            cmd_result.release_data();
        }
        return handle_result;
    } catch (std::exception& ex) {
        debugMessage(this->connection->virtual_server_id(), "{} Failed to handle connection command: {}. Closing connection.", this->connection->log_prefix(), ex.what());
        return CommandHandleResult::CLOSE_CONNECTION;
    }
}

CryptSetupHandler::CommandResult CryptSetupHandler::handleCommandClientInitIv(const ts::command_parser &cmd) {
    auto client = this->connection->getCurrentClient();
    assert(client);

    std::unique_lock state_lock{client->state_lock};
    if(client->connectionState() == ConnectionState::CONNECTED) { /* we've a reconnect */
        auto lastPingResponse = this->connection->ping_handler().last_ping_response();
        if(std::chrono::system_clock::now() - lastPingResponse < std::chrono::seconds(5)) {
            logMessage(this->connection->virtual_server_id(), "{} Client initialized session reconnect, but last ping response is not older then 5 seconds ({}). Ignoring attempt",
                       this->connection->log_prefix(),
                       duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - lastPingResponse).count()
            );
            return ts::command_result{error::ok};
        } else if(!config::voice::allow_session_reinitialize) {
            logMessage(this->connection->virtual_server_id(), "{} Client initialized session reconnect and last ping response is older then 5 seconds ({}). Dropping attempt because its not allowed due to config settings",
                       this->connection->log_prefix(),
                       duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - lastPingResponse).count()
            );
            return ts::command_result{error::ok};
        }
        logMessage(this->connection->virtual_server_id(), "{} Client initialized reconnect and last ping response is older then 5 seconds ({}). Allowing attempt",
                   this->connection->log_prefix(),
                   duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - lastPingResponse).count()
        );

        state_lock.unlock();

        {
            std::unique_lock server_channel_lock(client->server->get_channel_tree_lock()); /* we cant get moved if this is locked! */
            if(client->currentChannel)
                client->server->client_move(client->ref(), nullptr, nullptr, config::messages::timeout::connection_reinitialized, ViewReasonId::VREASON_TIMEOUT, false, server_channel_lock);
        }

        client->finalDisconnect();
        state_lock.lock();
    } else if(client->state >= ConnectionState::DISCONNECTING) {
        state_lock.unlock();
        std::shared_lock disconnect_finish{client->finalDisconnectLock}; /* await until the last disconnect has been processed */
        state_lock.lock();
        client->state = ConnectionState::INIT_HIGH;
    } else if(client->state == ConnectionState::INIT_HIGH) {
        logTrace(client->getServerId(), "{} Received a duplicated initiv. It seems like our initivexpand2 hasn't yet reached the client. The acknowledge handler should handle this issue for us.", CLIENT_STR_LOG_PREFIX_(client));
        return command_result{error::ok};
    } else {
        client->state = ConnectionState::INIT_HIGH;
    }
    state_lock.unlock();

    this->connection->reset();
    this->connection->packet_decoder().register_initiv_packet();
    this->connection->packet_statistics().reset_offsets();

    bool use_teaspeak = cmd.has_switch("teaspeak");
    if(use_teaspeak ? !config::server::clients::teaspeak : !config::server::clients::teamspeak)
        return command_result{error::client_type_is_not_allowed};

    if(use_teaspeak) {
        debugMessage(this->connection->virtual_server_id(), "{} Client using TeaSpeak.", this->connection->log_prefix());
        client->properties()[property::CLIENT_TYPE_EXACT] = ClientType::CLIENT_TEASPEAK;
    }

    /* normal TeamSpeak handling */
    this->seed_client = base64::decode(cmd.value("alpha"));
    if(this->seed_client.length() != 10)
        return ts::command_result{error::parameter_invalid, "alpha"};

    std::string clientOmega = base64::decode(cmd.value("omega")); //The identity public key
    std::string ip = cmd.value("ip");
    bool ot = cmd.has_key("ot") ? cmd.value_as<bool>("ot") : false;

    {
        this->remote_key = std::shared_ptr<ecc_key>(new ecc_key{}, [](ecc_key* key){
            if(!key) return;
            ecc_free(key);
            delete key;
        });

        auto state = ecc_import((const unsigned char *) clientOmega.data(), clientOmega.length(), &*this->remote_key);
        if(state != CRYPT_OK) {
            this->remote_key = nullptr;
            return ts::command_result{error::client_could_not_validate_identity};
        }

        client->properties()[property::CLIENT_UNIQUE_IDENTIFIER] = base64::encode(digest::sha1(cmd.value("omega")));
    }

    this->new_protocol = !use_teaspeak && ot && config::experimental_31 && (this->client_protocol_time_ >= 173265950ULL || this->client_protocol_time_ == (uint32_t) 5680278000ULL);

    {
        size_t server_seed_length = this->new_protocol ? 54 : 10;

        char beta[server_seed_length];
        generate_random((uint8_t *) beta, server_seed_length);

        this->seed_server = std::string{beta, server_seed_length};
    }

    if(this->new_protocol) {
        //Pre setup
        //Generate chain
        debugMessage(this->connection->virtual_server_id(), "{} Got client 3.1 protocol with build timestamp {}", this->connection->log_prefix(), this->client_protocol_time_);

        this->chain_data = serverInstance->getTeamSpeakLicense()->license();
        this->chain_data->chain->addEphemeralEntry();

        auto crypto_chain = this->chain_data->chain->exportChain();
        auto crypto_chain_hash = digest::sha256(crypto_chain);

        size_t sign_buffer_size{128};
        char sign_buffer[sign_buffer_size];

        prng_state prng_state{};
        memset(&prng_state, 0, sizeof(prng_state));

        auto sign_result = ecc_sign_hash(
                (u_char*) crypto_chain_hash.data(), crypto_chain_hash.length(),
                (u_char*) sign_buffer, &sign_buffer_size,
                &prng_state, find_prng("sprng"),
                client->getServer()->serverKey());

        if(sign_result != CRYPT_OK) {
            return ts::command_result{error::vs_critical, "failed to sign crypto chain"};
        }

        ts::command_builder answer{"initivexpand2"};
        answer.put_unchecked(0, "time", std::chrono::duration_cast<std::chrono::seconds>(std::chrono::system_clock::now().time_since_epoch()).count());
        answer.put_unchecked(0, "l", base64::encode(crypto_chain));
        answer.put_unchecked(0, "beta", base64::encode(this->seed_server));
        answer.put_unchecked(0, "omega", client->getServer()->publicServerKey());
        answer.put_unchecked(0, "proof", base64::encode((const char*) sign_buffer, sign_buffer_size));
        answer.put_unchecked(0, "tvd", "");
        answer.put_unchecked(0, "root", base64::encode((char*) this->chain_data->public_key, 32));
        answer.put_unchecked(0, "ot", "1");

        this->connection->send_command(answer.build(), false, nullptr);
        client->handshake.state = SpeakingClient::HandshakeState::SUCCEEDED; /* we're doing the verify via TeamSpeak */

        return CommandHandleResult::CONSUME_COMMAND; /* we don't want to send an error id=0 msg=ok */
    } else {
        debugMessage(this->connection->virtual_server_id(), "{} Got non client 3.1 protocol with build timestamp {}", this->connection->log_prefix(), this->client_protocol_time_, this->client_protocol_time_);

        auto server_public_key = client->getServer()->publicServerKey();
        if(server_public_key.empty()) {
            return ts::command_result{error::vs_critical, "failed to export server public key"};
        }

        {
            ts::command_builder answer{"initivexpand"};
            answer.put_unchecked(0, "alpha", base64::encode(this->seed_client));
            answer.put_unchecked(0, "beta", base64::encode(this->seed_server));
            answer.put_unchecked(0, "omega", server_public_key);

            if(use_teaspeak) {
                answer.put_unchecked(0, "teaspeak", "1");
                client->handshake.state = SpeakingClient::HandshakeState::BEGIN; /* we need to start the handshake */
            } else {
                client->handshake.state = SpeakingClient::HandshakeState::SUCCEEDED; /* we're using the provided identity as identity */
            }

            this->connection->send_command(answer.build(), false, nullptr);
            this->connection->packet_encoder().encrypt_pending_packets();
        }

        std::string error;
        if(!this->connection->getCryptHandler()->setupSharedSecret(this->seed_client, this->seed_server, &*this->remote_key, client->getServer()->serverKey(), error)){
            logError(this->connection->virtual_server_id(), "{} Failed to calculate shared secret {}. Dropping client.",
                    this->connection->log_prefix(), error);
            return ts::command_result{error::vs_critical};
        }

        return CommandHandleResult::CONSUME_COMMAND; /* we don't want to send an error id=0 msg=ok */
    }
}

CryptSetupHandler::CommandResult CryptSetupHandler::handleCommandClientEk(const ts::command_parser &cmd) {
    debugMessage(this->connection->virtual_server_id(), "{} Got client ek!", this->connection->log_prefix());

    auto client_key = base64::decode(cmd.value("ek"));
    auto private_key = this->chain_data->chain->generatePrivateKey(this->chain_data->root_key, this->chain_data->root_index);

    this->connection->getCryptHandler()->setupSharedSecretNew(this->seed_client, this->seed_server, (char*) private_key.data(), client_key.data());
    this->connection->packet_encoder().acknowledge_manager().reset();

    {
        char buffer[2];
        le2be16(1, buffer);

        auto pflags = protocol::PacketFlag::NewProtocol;
        this->connection->send_packet(protocol::PacketType::ACK, (protocol::PacketFlag::PacketFlag) pflags, buffer, 2);
        //Send the encrypted acknowledge (most the times the second packet; If not we're going into the resend loop)
        //We cant use the send_packet_acknowledge function since it sends the acknowledge unencrypted
    }

    return CommandHandleResult::CONSUME_COMMAND; /* we don't want to send an error id=0 msg=ok */
}

CryptSetupHandler::CommandResult CryptSetupHandler::handleCommandClientInit(const ts::command_parser &) {
    /* the client must have received everything */
    this->connection->packet_encoder().acknowledge_manager().reset();
    this->seed_client.clear();
    this->seed_server.clear();
    this->chain_data = nullptr;

    return CommandHandleResult::PASS_THROUGH;
}