#include "pch.h"
#include "Crypt.h"
#include "Patches.h"
#include "Util.h"

// Carrot: Notes;
// - The public key should be taken from the attached executable. Find "BgIAAACk" and get 200 characters
// - I didn't test this
// - All the std::runtime_errors here should use RBX::runtime_error but I didn't want to hook it. YOu do it pizza
// - Idk how iffy the base64 decoding is lol

// Base/rbx/Crypt.cpp

Crypt::Crypt()
{
    static const char* keyB64 = "BgIAAACkAABSU0ExAAQAAAEAAQBBC9x87ewme1daAWZGJpK3rTepuMbivjpgTyNwu+f7gxpyreKUzQDSpqj+DhT7ni432DY0H4+Zv6MazZFrjeBnS6YnD02fuGkmLb0mg0yO784esDkImtDpqNrmoqPwAe94jr8K8vsxaOeF5b7OfA+lWUcag6Jh0Yb+mZxhFbmvrg==";
    char pbKeyBlob[256];
    int dwBlobLen = 256;

    // http://support.microsoft.com/kb/238187
    if (!CryptAcquireContext(&context, NULL, MS_DEF_PROV, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT))
    {
        if (::GetLastError() == NTE_BAD_KEYSET)
        {
            if (!CryptAcquireContext(&context, NULL, MS_DEF_PROV, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT | CRYPT_NEWKEYSET))
            {
                throw std::runtime_error("Error " + std::to_string(GetLastError()) + " during CryptAcquireContext 2\n");
            }
        }
        else
        {
            throw std::runtime_error("Error " + std::to_string(GetLastError()) + " during CryptAcquireContext\n");
        }
    }

    ATL::Base64Decode(keyB64, strlen(keyB64), (BYTE*)pbKeyBlob, &dwBlobLen);

    if (!CryptImportKey(context, (BYTE*)pbKeyBlob, dwBlobLen, 0, 0, &key))
    {
        throw std::runtime_error("Error " + std::to_string(GetLastError()) + " during CryptImportKey");
    }
}

Crypt::~Crypt()
{
    CryptDestroyKey(key);
    CryptReleaseContext(context, 0);
}

void Crypt::verifySignatureBase64(std::string message, std::string signatureBase64)
{
    // Check for a reasonable signature length before verifying
    if (signatureBase64.length() > 4096)
    {
        throw std::runtime_error("");
    }

    HCRYPTHASH hash;

    if (!CryptCreateHash(context, CALG_SHA_256, NULL, 0, &hash))
    {
        throw std::runtime_error("");
    }

    try
    {
        if (!CryptHashData(hash, (BYTE*)message.c_str(), message.size(), 0))
        {
            throw std::runtime_error("");
        }

        // Roblox does stupid stuff. We can just use C++17 features :-)

        std::vector<BYTE> signature = Util::base64Decode(signatureBase64);

        /*
            The native cryptography API uses little-endian byte order
            while the .NET Framework API uses big-endian byte order.
            If you are verifying a signature generated by using a .NET Framework
            API (or similar), you must swap the order of signature bytes before
            calling the CryptVerifySignature function to verify the signature.
        */

        std::reverse(signature.begin(), signature.end());

        BYTE* signatureData = new BYTE[signature.size()];
        std::copy(signature.begin(), signature.end(), signatureData);

        if (!CryptVerifySignature(hash, signatureData, signature.size(), key, NULL, 0))
        {
            throw std::runtime_error("");
        }
    }
    catch (...)
    {
        ::CryptDestroyHash(hash);
        throw std::runtime_error("");
    }

    ::CryptDestroyHash(hash);
}

// End Base/rbx/Crypt.cpp

Crypt__verifySignatureBase64_t Crypt__verifySignatureBase64 = (Crypt__verifySignatureBase64_t)ADDRESS_CRYPT__VERIFYSIGNATUREBASE64;

void __fastcall Crypt__verifySignatureBase64_hook(HCRYPTPROV* _this, void*, int a2, BYTE* pbData, int a4, int a5, int a6, DWORD dwDataLen, int a8, int a9, int a10, int a11, int a12, int a13, int a14, int a15)
{
    std::string message;
    std::string signatureBase64;

    // Get message
    /*
        v18 = pbData;
        if ((unsigned int)a8 < 0x10)
        {
            v18 = (const BYTE*)&pbData;
        }
    */

    const BYTE* v18 = pbData;
    if ((unsigned int)a8 < 0x10)
    {
        v18 = (const BYTE*)&pbData;
    }

    message = std::string(reinterpret_cast<char const*>(pbData), dwDataLen);

    // Get signatureBase64
    /*
        v21 = (int *)a10;
        if ((unsigned int)a15 < 0x10)
        {
            v21 = &a10;
        }

        sub_79EA70(v21, a14, &v30, &dwSigLen); // ATL::Base64Decode
    */

    int* v21 = (int*)a10;
    if ((unsigned int)a15 < 0x10)
    {
        v21 = &a10;
    }

    signatureBase64 = std::string(reinterpret_cast<const char*>(v21), a14);

    Crypt().verifySignatureBase64(message, signatureBase64);
}