/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: RSAPrivate.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "../Crypto.hpp"
#include "RSA.hpp"
#include "RSAPrivate.hpp"
#include "RSAPublic.hpp"

namespace Aurora::Crypto::RSA
{
    PrivateRSA::PrivateRSA(rsa_key &key) : key_(key)
    {

    }

    PrivateRSA::~PrivateRSA()
    {
        ::rsa_free(&this->key_);
    }

    bool PrivateRSA::Sign(const Memory::MemoryViewRead & payload,
                          AuHashing::EHashType method,
                          EPaddingType type,
                          Memory::ByteBuffer &out)
    {
        prng_state yarrow_prng;
        const int salt = 0;
        
        if (!payload.HasMemory())
        {
            SysPushErrorParam();
            return {};
        }

        int padding = ::Crypto::PaddingToType(type);
        if (padding == 0xFF)
        {
            SysPushErrorCrypt("invalid pad {}", type);
            return false;
        }

        int hash = ::Crypto::HashMethodToId(method);
        if (hash == 0xFF)
        {
            SysPushErrorCrypt("invalid hash {}", method);
            return false;
        }

        auto writeView = out.GetOrAllocateLinearWriteable(1024);
        if (!writeView)
        {
            SysPushErrorMem();
            return {};
        }

        Memory::ByteBuffer hashVec;

        if (!AuTryResize(hashVec, 128))
        {
            SysPushErrorMem();
            return false;
        }

        unsigned long hashSize = hashVec.size();
        auto iRet = ::hash_memory(hash,
                                  AuReinterpretCast<const unsigned char *>(payload.ptr), payload.length,
                                  AuReinterpretCast<unsigned char *>(hashVec.data()), &hashSize);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        iRet = yarrow_start(&yarrow_prng);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        unsigned long len = 1024;
        iRet = ::rsa_sign_hash_ex(AuReinterpretCast<const unsigned char *>(hashVec.data()), hashSize,
                                  out.writePtr, &len,
                                  padding,
                                  &yarrow_prng,
                                  ::Crypto::gPrngYarrow,
                                  hash,
                                  salt,
                                  &this->key_);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        out.writePtr += len;
        return true;
    }

    bool PrivateRSA::Decrypt(const AuMemoryViewRead &payload,
                             EPaddingType type, 
                             AuByteBuffer &out)
    {
        if (!payload.HasMemory())
        {
            SysPushErrorParam();
            return {};
        }

        int padding = ::Crypto::PaddingToType(type);
        if (padding == 0xFF)
        {
            SysPushErrorCrypt("invalid pad {}", type);
            return false;
        }

        const int prng_idx = padding == LTC_PKCS_1_PSS ? ::Crypto::gPrngYarrow : 0;
        if (prng_idx < 0)
        {
            SysPushErrorCrypt("{}", prng_idx);
            return false;
        }

        auto writeView = out.GetOrAllocateLinearWriteable(payload.length);
        if (!writeView)
        {
            SysPushErrorMem();
            return {};
        }

        unsigned long len = payload.length;

        int stat = 0;
        auto ret = ::rsa_decrypt_key_ex(AuReinterpretCast<const unsigned char *>(payload.ptr),
                                        payload.length,
                                        out.writePtr, &len,
                                        NULL, 0,
                                        0, // hash? excuse me?
                                        padding,
                                        &stat,
                                        &this->key_);
        if (ret != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", ret);
            return false;
        }

        out.writePtr += len;
        return stat == 1;
    }

    AuSPtr<IRSAPublic> PrivateRSA::ToPublic()
    {
        rsa_key key {};
        key.type = PK_PUBLIC;

    #define COPY_KEY_PART(x)                                    \
        if (ltc_mp.init_copy(&key.x, this->key_.x) != CRYPT_OK) \
        {                                                       \
            return {};                                          \
        }

        COPY_KEY_PART(N)
        COPY_KEY_PART(e)
       
        return AuMakeShared<PublicRSA>(key);
    }

    bool PrivateRSA::ToKey(const RSAMeta &meta, AuByteBuffer &out)
    {
        return ExportRSAKey(this->key_, meta.type, meta.encoding, out);
    }

    AUKN_SYM IRSAPrivate *OpenRSAPrivateNew(const RSAKey &key)
    {
        rsa_key in {};

        if (key.meta.type == EKeyType::eKeyPublic)
        {
            SysPushErrorArg("Attempted to import a public key as a private key.");
            return nullptr;
        }

        if (!ImportRSAKey(in, key))
        {
            return nullptr;
        }

        auto pPrivate = _new PrivateRSA(in);
        if (!pPrivate)
        {
            ::rsa_free(&in);
            return nullptr;
        }

        return pPrivate;
    }

    AUKN_SYM void OpenRSAPrivateRelease(IRSAPrivate *pPrivate)
    {
        AuSafeDelete<PrivateRSA *>(pPrivate);
    }

    AUKN_SYM IRSAPrivate *NewRSAKeyNew(AuUInt16 keySize)
    {
        rsa_key key {};
        const int prng_idx = register_prng(&sprng_desc);
        prng_state yarrow_prng;

        auto iRet = ::rsa_make_key(&yarrow_prng, prng_idx, keySize / 8, 65537, &key);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return nullptr;
        }

        auto ret = _new PrivateRSA(key);
        if (!ret)
        {
            ::rsa_free(&key);
            return nullptr;
        }

        return ret;
    }

    AUKN_SYM void NewRSAKeyRelease(IRSAPrivate *re)
    {
        return OpenRSAPrivateRelease(re);
    }
}