/*
 * Copyright (C) 2023 Uniontech Software Technology Co., Ltd.
 *
 * Author:     gongheng<gongheng@uniontech.com>
 *
 * Maintainer: gongheng<gongheng@uniontech.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "tpmwork.h"
#include "vaultglobaldefine.h"
#include "vaultconfig.h"
#include "operatorcenter.h"

#include <QProcess>
#include <QFile>
#include <QTextStream>
#include <QDebug>
#include <QStandardPaths>
#include <QLibrary>

#define MAX_SIZE 3000
#define TPM_LIBRARY_NAME "libutpm2.so"

TPMWork::TPMWork(QObject *parent)
    : QObject(parent)
    , tpmLib(new QLibrary(TPM_LIBRARY_NAME))
{
    if (!tpmLib->load())
        qInfo() << "Vault warning: load tpm failed!";
}

TPMWork::~TPMWork()
{
    if (tpmLib) {
        tpmLib->unload();
        delete tpmLib;
        tpmLib = nullptr;
    }
}

bool TPMWork::checkTPMAvailable()
{
    if (!tpmLib->isLoaded())
        return false;

    // Confirm whether tpm is available by judging whether the random number interface is available
    QString output;
    return getRandom(2, output);
}

bool TPMWork::getRandom(int size, QString &output)
{
    if (!tpmLib->isLoaded())
        return false;

    if (size % 2 != 0 && size < 1) {
        qCritical() << "Vault error: random size must be even and not less than one!";
        return false;
    }

    typedef bool (*p_get_random)(uint16_t *size, uint8_t ranbytes[]);
    p_get_random utpm2_get_random = (p_get_random) tpmLib->resolve("utpm2_get_random");
    if (utpm2_get_random) {
        uint16_t len = size / 2;
        uint8_t *random = (uint8_t *)malloc(sizeof(uint8_t) * len);
        memset(random, 0, sizeof(uint8_t) * len);
        if (utpm2_get_random(&len, random)) {
            char *out = (char *)malloc(sizeof(uint8_t) * size + 1);
            memset(out, 0, sizeof(uint8_t) * size + 1);
            for (size_t i = 0; i < len; ++i) {
                sprintf(out + (i * 2), "%02x", random[i]);
            }
            output = QString(out);
            free(random);
            free(out);
            return true;
        }
    }
    return false;
}

bool TPMWork::isSupportAlgo(const QString &algoName)
{
    if (!tpmLib->isLoaded())
        return false;

    typedef bool (*p_check_algo)(const char *alg);
    p_check_algo utpm2_check_algo = (p_check_algo) tpmLib->resolve("utpm2_check_alg");
    if (utpm2_check_algo) {
        QByteArray arAlgoName = algoName.toUtf8();
        if (utpm2_check_algo(arAlgoName.data())) {
            return true;
        }
    }
    return false;
}

bool TPMWork::initTpm2(const QString hashAlgo, const QString keyAlgo, const QString keyPin)
{
    if (!tpmLib->isLoaded())
        return false;

    typedef bool (*p_init)(char *algdetail, char *galg, const char *auth, const char *dir);
    p_init utpm2_init = (p_init)tpmLib->resolve("utpm2_init");
    if (utpm2_init) {
        QByteArray arKeyAlgo = keyAlgo.toUtf8();
        QByteArray arHashAlgo = hashAlgo.toUtf8();
        QByteArray arKeyPin = keyPin.toUtf8();
        QByteArray arDir = QString(VAULT_BASE_PATH).toUtf8();
        if (utpm2_init(arKeyAlgo.data(), arHashAlgo.data(), arKeyPin.data(), arDir.data())) {
            return true;
        } else {
            qInfo() << "Vault error: utpm2_init return false!";
        }
    } else {
        qInfo() << "Vault error: resolve utpm2_init failed!";
    }
    return false;
}

void TPMWork::slotDoEncrypt(const QString hashAlgo, const QString keyAlgo, const QString keyPin, const QString password)
{
    if (!initTpm2(hashAlgo, keyAlgo, keyPin)) {
        emit sigEncryptFinished(false);
        return;
    }

    typedef int (*p_encrypt_decrypt)(const char *dir, bool isdecrypt, const char *auth, uint8_t inbytes[], uint8_t outbytes[], uint16_t *size);
    p_encrypt_decrypt utpm2_encrypt_decrypt = (p_encrypt_decrypt) tpmLib->resolve("utpm2_encrypt_decrypt");
    if (utpm2_encrypt_decrypt) {
        QByteArray arDir = QString(VAULT_BASE_PATH).toUtf8();
        QByteArray arKeyPin = keyPin.toUtf8();
        QByteArray arrPassword = password.toUtf8();
        uint16_t len = static_cast<uint16_t>(arrPassword.size());
        uint8_t *inbuffer = reinterpret_cast<uint8_t*>(arrPassword.data());
        uint8_t out_text[MAX_SIZE] = { 0 };
        if (utpm2_encrypt_decrypt(arDir.data(), false, arKeyPin.data(), inbuffer, out_text, &len)) {
            QFile file(TPM_ENCRYPT_FILE_PATH);
            if (file.open(QIODevice::WriteOnly | QIODevice::Text)) {
                file.write(reinterpret_cast<const char*>(out_text), len);
                file.close();
                emit sigEncryptFinished(true);
                return;
            } else {
                qInfo() << "Vault error: open encrypt file failed!";
            }
        } else {
            qInfo() << "Vault error: utpm2_encry_decrypt return false!";
        }
    } else {
        qInfo() << "Vault error: resolve utpm2_encry_decrypt failed!";
    }
    emit sigEncryptFinished(false);
}

void TPMWork::slotDoDecrypt(const QString &keyPin)
{
    if (!tpmLib->isLoaded())
        return;

    QString password { "" };
    typedef int (*p_encrypt_decrypt)(const char *dir, bool isdecrypt, const char *auth, uint8_t inbytes[], uint8_t outbytes[], uint16_t *size);
    p_encrypt_decrypt utpm2_encrypt_decrypt = (p_encrypt_decrypt) tpmLib->resolve("utpm2_encrypt_decrypt");
    if (utpm2_encrypt_decrypt) {
        QByteArray arDir = QString(VAULT_BASE_PATH).toUtf8();
        QByteArray arKeyPin = keyPin.toUtf8();
        QFile file(TPM_ENCRYPT_FILE_PATH);
        if (file.open(QIODevice::ReadOnly | QIODevice::Text)) {
            QByteArray ciphertext = file.readAll();
            file.close();
            uint16_t len = static_cast<uint16_t>(ciphertext.size());
            uint8_t out_text[MAX_SIZE] = { 0 };
            uint8_t *inbuffer = reinterpret_cast<uint8_t*>(ciphertext.data());
            if (utpm2_encrypt_decrypt(arDir.data(), true, arKeyPin.data(), inbuffer, out_text, &len)) {
                password = QString::fromUtf8(reinterpret_cast<const char*>(out_text), len);
                emit sigDecryptFinished(true, password);
                return;
            } else {
                qInfo() << "Vault error: utpm2_encry_decrypt return failed!";
            }
        } else {
            qInfo() << "Vault error: open encrypt file failed!";
        }
    } else {
        qInfo() << "Vault error: resolve utpm2_encry_decrypt failed!";
    }
    emit sigDecryptFinished(false, password);
}
