/*
    SPDX-FileCopyrightText: 2005 Joris Guisson <joris.guisson@gmail.com>

    SPDX-License-Identifier: GPL-2.0-or-later
*/
#include "encryptedauthenticate.h"

#include <QRandomGenerator>
#include <algorithm>

#include "encryptedpacketsocket.h"
#include "functions.h"
#include "rc4encryptor.h"
#include <net/socks.h>
#include <torrent/globals.h>
#include <torrent/server.h>
#include <util/functions.h>
#include <util/log.h>

using namespace bt;

namespace mse
{
EncryptedAuthenticate::EncryptedAuthenticate(const net::Address &addr,
                                             TransportProtocol proto,
                                             const SHA1Hash &info_hash,
                                             const PeerID &peer_id,
                                             PeerConnector::WPtr pcon)
    : Authenticate(addr, proto, info_hash, peer_id, pcon)
{
    mse::GeneratePublicPrivateKey(xa, ya);
    state = NOT_CONNECTED;
    buf_size = 0;
    vc_off = 0;
    dec_bytes = 0;
    crypto_select = 0;
    pad_D_len = 0;
    end_of_crypto_handshake = 0;
    // Out(SYS_CON|LOG_DEBUG) << "EncryptedAuthenticate : " << ip << ":" << port << endl;
}

EncryptedAuthenticate::~EncryptedAuthenticate()
{
}

void EncryptedAuthenticate::connected()
{
    // we are connected so send ya and some padding
    Uint8 tmp[608];
    ya.toBuffer(tmp, 96);
    sock->sendData(tmp, 96 + QRandomGenerator::global()->bounded(512));
    state = SENT_YA;
}

/*
1 A->B: Diffie Hellman Ya, PadA
2 B->A: Diffie Hellman Yb, PadB
3 A->B: HASH('req1', S), HASH('req2', SKEY) xor HASH('req3', S), ENCRYPT(VC, crypto_provide, len(PadC), PadC, len(IA)), ENCRYPT(IA)
4 B->A: ENCRYPT(VC, crypto_select, len(padD), padD), ENCRYPT2(Payload Stream)
5 A->B: ENCRYPT2(Payload Stream)
*/

void EncryptedAuthenticate::handleYB()
{
    // if you can't sent 96 bytes you are not worth the effort
    if (buf_size < 96) {
        Out(SYS_CON | LOG_DEBUG) << "Not enough data received, encrypted authentication failed" << endl;
        onFinish(false);
        return;
    }

    // read Yb
    yb = BigInt::fromBuffer(buf, 96);

    // calculate s
    s = mse::DHSecret(xa, yb);

    state = GOT_YB;
    // now we must send line 3
    Uint8 tmp_buf[120]; // temporary buffer
    bt::SHA1Hash h1, h2; // temporary hash

    // generate and send the first hash
    memcpy(tmp_buf, "req1", 4);
    s.toBuffer(tmp_buf + 4, 96);
    h1 = SHA1Hash::generate(tmp_buf, 100);
    sock->sendData(h1.getData(), 20);

    // generate second and third hash and xor them
    memcpy(tmp_buf, "req2", 4);
    memcpy(tmp_buf + 4, info_hash.getData(), 20);
    h1 = SHA1Hash::generate(tmp_buf, 24);

    memcpy(tmp_buf, "req3", 4);
    s.toBuffer(tmp_buf + 4, 96);
    h2 = SHA1Hash::generate(tmp_buf, 100);
    sock->sendData((h1 ^ h2).getData(), 20);

    // now we enter encrypted mode the keys are :
    // HASH('keyA', S, SKEY) for the encryption key
    // HASH('keyB', S, SKEY) for the decryption key
    enc = mse::EncryptionKey(true, s, info_hash);
    dec = mse::EncryptionKey(false, s, info_hash);

    our_rc4 = std::make_unique<RC4Encryptor>(dec, enc);

    // now we must send ENCRYPT(VC, crypto_provide, len(PadC), PadC, len(IA))
    memset(tmp_buf, 0, 16); // VC are 8 0x00's
    if (bt::ServerInterface::unencryptedConnectionsAllowed())
        tmp_buf[11] = 0x03; // we support both plain text and rc4
    else
        tmp_buf[11] = 0x02;
    WriteUint16(tmp_buf, 12, 0x0000); // no padC
    WriteUint16(tmp_buf, 14, 68); // length of IA, which will be the bittorrent handshake
    // send IA which is the handshake
    makeHandshake(tmp_buf + 16, info_hash, our_peer_id);
    sock->sendData(our_rc4->encrypt(tmp_buf, 84), 84);

    // search for the encrypted VC in the data
    findVC();
}

void EncryptedAuthenticate::findVC()
{
    Uint8 vc[8] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

    RC4Encryptor rc4(enc, dec);
    memcpy(vc, rc4.encrypt(vc, 8), 8);

    Uint32 max_i = buf_size - 8;
    for (Uint32 i = 96; i < max_i; i++) {
        if (vc[0] == buf[i] && memcmp(buf + i, vc, 8) == 0) {
            state = FOUND_VC;
            vc_off = i;
            handleCryptoSelect();
            return;
        }
    }

    // we haven't found it in the first 616 bytes (96 + max 512 padding + 8 bytes VC)
    if (buf_size >= 616) {
        onFinish(false);
    }
}

void EncryptedAuthenticate::handleCryptoSelect()
{
    // not enough data available so lets come back later
    if (vc_off + 14 >= buf_size)
        return;

    // now decrypt the first 14 bytes
    our_rc4->decrypt(buf + vc_off, 14);
    // check the VC
    for (Uint32 i = vc_off; i < vc_off + 8; i++) {
        if (buf[i]) {
            Out(SYS_CON | LOG_DEBUG) << "Invalid VC " << endl;
            onFinish(false);
            return;
        }
    }

    crypto_select = ReadUint32(buf, vc_off + 8);
    pad_D_len = ReadUint16(buf, vc_off + 12);
    if (pad_D_len > 512) {
        Out(SYS_CON | LOG_DEBUG) << "Invalid pad D length" << endl;
        onFinish(false);
        return;
    }

    end_of_crypto_handshake = vc_off + 14 + pad_D_len;
    if (!(vc_off + 14 + pad_D_len < buf_size)) {
        // padD is not complete, wait for that
        state = WAIT_FOR_PAD_D;
        return;
    }

    handlePadD();
}

void EncryptedAuthenticate::handlePadD()
{
    // decrypt the padding
    our_rc4->decrypt(buf + (vc_off + 14), pad_D_len);

    if (crypto_select & 0x00000001) { // plain_text selected
        our_rc4.reset();
    } else if (crypto_select & 0x00000002) { // now it must be rc4 if not exit
        sock->setRC4Encryptor(std::move(our_rc4));
    } else { // we don't support anything else so error out
        onFinish(false);
        return;
    }

    // noz we wait for the normal handshake
    state = NORMAL_HANDSHAKE;
    // if we have read more then the crypto handshake, reinsert it
    if (buf_size > vc_off + 14 + pad_D_len) {
        Uint32 off = vc_off + 14 + pad_D_len;
        sock->reinsert(buf + off, buf_size - off);
        Authenticate::onReadyRead();
    }
}

void EncryptedAuthenticate::onReadyRead()
{
    if (finished)
        return;

    if (socks) {
        switch (socks->onReadyToRead()) {
        case net::Socks::FAILED:
            Out(SYS_CON | LOG_NOTICE) << "Failed to connect to host via socks server " << endl;
            onFinish(false);
            break;
        case net::Socks::CONNECTED:
            // connection established, so get rid of socks shit
            delete socks;
            socks = nullptr;
            connected();
            if (sock->bytesAvailable() > 0)
                onReadyRead();
            break;
        default:
            break;
        }
        return;
    }

    Uint32 ba = sock->bytesAvailable();
    if (ba == 0) {
        onFinish(false);
        return;
    }

    if (state != NORMAL_HANDSHAKE) {
        if (buf_size + ba > MAX_EA_BUF_SIZE)
            ba = MAX_EA_BUF_SIZE - buf_size;

        // do not read past the end of padD
        if (pad_D_len > 0 && buf_size + ba > vc_off + 14 + pad_D_len)
            ba = (vc_off + 14 + pad_D_len) - buf_size;
        // read data
        buf_size += sock->readData(buf + buf_size, ba);
    }

    switch (state) {
    case SENT_YA:
        if (ba > 608) {
            onFinish(false);
        } else {
            handleYB();
        }
        break;
    case GOT_YB:
        findVC();
        break;
    case FOUND_VC:
        handleCryptoSelect();
        break;
    case WAIT_FOR_PAD_D:
        handlePadD();
        break;
    case NORMAL_HANDSHAKE:
        // let AuthenticateBase deal with the data
        AuthenticateBase::onReadyRead();
        break;
    default:
        break;
    };
}

}

#include "moc_encryptedauthenticate.cpp"
