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

    SPDX-License-Identifier: GPL-2.0-or-later
*/
#include "encryptedpacketsocket.h"
#include <QtGlobal>
#include <cerrno>
#include <cstdlib>
#include <sys/types.h>
#include <util/log.h>
#include <util/sha1hash.h>
#ifndef Q_OS_WIN
#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h>
#else
#define IPTOS_THROUGHPUT 0x08
#endif
#include "rc4encryptor.h"
#include <net/socketmonitor.h>
#include <util/functions.h>

using namespace bt;
using namespace net;

namespace mse
{
Uint8 EncryptedPacketSocket::tos = IPTOS_THROUGHPUT;

EncryptedPacketSocket::EncryptedPacketSocket(int ip_version)
    : net::PacketSocket(true, ip_version)
    , monitored(false)
{
    sock->setBlocking(false);
    sock->setTOS(tos);
    reinserted_data = nullptr;
    reinserted_data_size = 0;
    reinserted_data_read = 0;
}

EncryptedPacketSocket::EncryptedPacketSocket(int fd, int ip_version)
    : net::PacketSocket(fd, ip_version)
    , monitored(false)
{
    sock->setBlocking(false);
    sock->setTOS(tos);
    reinserted_data = nullptr;
    reinserted_data_size = 0;
    reinserted_data_read = 0;
}

EncryptedPacketSocket::EncryptedPacketSocket(std::unique_ptr<net::SocketDevice> sd)
    : net::PacketSocket(std::move(sd))
    , monitored(false)
{
    sock->setBlocking(false);
    sock->setTOS(tos);
    reinserted_data = nullptr;
    reinserted_data_size = 0;
    reinserted_data_read = 0;
}

EncryptedPacketSocket::~EncryptedPacketSocket()
{
    if (monitored) {
        stopMonitoring();
    }

    delete[] reinserted_data;
}

void EncryptedPacketSocket::startMonitoring(net::SocketReader *rdr)
{
    setReader(rdr);
    SocketMonitor::instance().add(this);
    monitored = true;
    if (reinserted_data) {
        if (enc) {
            enc->decrypt(reinserted_data + reinserted_data_read, reinserted_data_size - reinserted_data_read);
        }

        rdr->onDataReady(reinserted_data + reinserted_data_read, reinserted_data_size - reinserted_data_read);
        delete[] reinserted_data;
        reinserted_data = nullptr;
        reinserted_data_size = 0;
    }
}

void EncryptedPacketSocket::stopMonitoring()
{
    SocketMonitor::instance().remove(this);
    monitored = false;
    rdr = nullptr;
}

Uint32 EncryptedPacketSocket::sendData(const Uint8 *data, Uint32 len)
{
    if (enc) {
        // we need to make sure all data is sent because of the encryption
        Uint32 ds = 0;
        const Uint8 *ed = enc->encrypt(data, len);
        while (sock->ok() && ds < len) {
            const Uint32 ret = sock->send(ed + ds, len - ds);
            ds += ret;
            if (ret == 0) {
                Out(SYS_CON | LOG_DEBUG) << "ret = 0" << endl;
            }
        }
        if (ds != len) {
            Out(SYS_CON | LOG_DEBUG) << "ds != len" << endl;
        }
        return ds;
    } else {
        const Uint32 ret = sock->send(data, len);
        if (ret != len) {
            Out(SYS_CON | LOG_DEBUG) << "ret != len" << endl;
        }
        return ret;
    }
}

Uint32 EncryptedPacketSocket::readData(Uint8 *buf, Uint32 len)
{
    Uint32 ret2 = 0;
    if (reinserted_data) {
        Uint32 tr = reinserted_data_size - reinserted_data_read;
        if (tr < len) {
            memcpy(buf, reinserted_data + reinserted_data_read, tr);
            delete[] reinserted_data;
            reinserted_data = nullptr;
            reinserted_data_size = reinserted_data_read = 0;
            ret2 = tr;
            if (enc) {
                enc->decrypt(buf, tr);
            }
        } else {
            tr = len;
            memcpy(buf, reinserted_data + reinserted_data_read, tr);
            reinserted_data_read += tr;
            if (enc) {
                enc->decrypt(buf, tr);
            }
            return tr;
        }
    }

    if (len == ret2) {
        return ret2;
    }

    const Uint32 ret = sock->recv(buf + ret2, len - ret2);
    if (ret + ret2 > 0 && enc) {
        enc->decrypt(buf, ret + ret2);
    }

    return ret;
}

Uint32 EncryptedPacketSocket::bytesAvailable() const
{
    const Uint32 ba = sock->bytesAvailable();
    if (reinserted_data_size - reinserted_data_read > 0) {
        return ba + (reinserted_data_size - reinserted_data_read);
    } else {
        return ba;
    }
}

void EncryptedPacketSocket::close()
{
    sock->close();
}

bool EncryptedPacketSocket::connectTo(const QString &ip, Uint16 port)
{
    // do a safety check
    if (ip.isNull() || ip.length() == 0) {
        return false;
    }

    return connectTo(net::Address(ip, port));
}

bool EncryptedPacketSocket::connectTo(const net::Address &addr)
{
    // we don't wanna block the current thread so set non blocking
    sock->setBlocking(false);
    sock->setTOS(tos);
    if (sock->connectTo(addr)) {
        return true;
    }

    return false;
}

void EncryptedPacketSocket::initCrypt(const bt::SHA1Hash &dkey, const bt::SHA1Hash &ekey)
{
    enc = std::make_unique<RC4Encryptor>(dkey, ekey);
}

void EncryptedPacketSocket::disableCrypt()
{
    enc.reset();
}

bool EncryptedPacketSocket::ok() const
{
    return sock->ok();
}

QString EncryptedPacketSocket::getRemoteIPAddress() const
{
    return sock->getPeerName().toString();
}

bt::Uint16 EncryptedPacketSocket::getRemotePort() const
{
    return sock->getPeerName().port();
}

net::Address EncryptedPacketSocket::getRemoteAddress() const
{
    return sock->getPeerName();
}

void EncryptedPacketSocket::setRC4Encryptor(std::unique_ptr<RC4Encryptor> e)
{
    enc = std::move(e);
}

void EncryptedPacketSocket::reinsert(const Uint8 *d, Uint32 size)
{
    //      Out() << "Reinsert : " << size << endl;
    Uint32 off = 0;
    if (reinserted_data) {
        off = reinserted_data_size;
        reinserted_data = (Uint8 *)realloc(reinserted_data, reinserted_data_size + size);
        reinserted_data_size += size;
    } else {
        reinserted_data = new Uint8[size];
        reinserted_data_size = size;
    }
    memcpy(reinserted_data + off, d, size);
}

bool EncryptedPacketSocket::connecting() const
{
    return sock->state() == net::SocketDevice::CONNECTING;
}

bool EncryptedPacketSocket::connectSuccessful() const
{
    return sock->connectSuccessful();
}

void EncryptedPacketSocket::setRemoteAddress(const net::Address &addr)
{
    sock->setRemoteAddress(addr);
}

void EncryptedPacketSocket::preProcess(Uint8 *data, Uint32 size)
{
    if (enc) {
        enc->encryptReplace(data, size);
    }
}

void EncryptedPacketSocket::postProcess(Uint8 *data, Uint32 size)
{
    if (enc) {
        enc->decrypt(data, size);
    }
}

}
