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

    File: Aes.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "../Crypto.hpp"
#include "Aes.hpp"
#include <tomcrypt.h>
#include <Source/AuCrypto.hpp>

namespace Aurora::Crypto::AES
{
    AUKN_SYM AuUInt GetSafeCipherPadding(const Memory::MemoryViewRead &parameters)
    {
        auto tptr = reinterpret_cast<const unsigned char *>(parameters.ptr);

        auto overhang = parameters.length & (128 - 1);
        auto primaryLength = parameters.length - overhang;

        if (overhang == 0)
        {
            auto lastbyte = tptr[primaryLength - 1];
            if (lastbyte <= 128)
            {
                return parameters.length + 128;
            }
            else
            {
                return parameters.length;
            }
        }
        else
        {
            return 128 - overhang;
        }
    }

    static bool EncryptCoolCodePadding(const void *plainText, AuUInt plainTextLength,
                                       const void *iv, void *outIv, AuUInt ivLength,
                                       const void *key, AuUInt keyLength,
                                       Memory::ByteBuffer &out,
                                       symmetric_CBC &cbc)
    {

        auto overhang = plainTextLength & (128 - 1) /* || fast_rng_bool */;
        auto padding = 128 - overhang;
        auto primaryLength = plainTextLength - overhang;

        int iRet;
        auto tptr = AuReinterpretCast<const unsigned char *>(plainText);

        unsigned char pad[128] =
        {
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE,
            0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00,
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE,
            0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00,
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE,
            0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00,
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE,
            0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00,
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE,
            0x00, 0x00, 0x00, 0x00, 0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x00,
            0xC0, 0x01, 0xC0, 0xDE, 0x00, 0x00, 0x00, 0x80
        };

        if (overhang == 0)
        {
            auto lastbyte = tptr[primaryLength - 1];
            if (lastbyte <= 128)
            {
                if (!out.GetOrAllocateLinearWriteable(primaryLength + 128))
                {
                    SysPushErrorMem();
                    return false;
                }

                // encrypt plain text
                iRet = cbc_encrypt(tptr, out.writePtr, primaryLength, &cbc);
                if (iRet != CRYPT_OK)
                {
                    SysPushErrorCrypto("{}", iRet);
                    return false;
                }

                // encrypt another 128 bytes of trash :( 50/50 chance 
                iRet = cbc_encrypt(pad, out.writePtr + primaryLength, 128, &cbc);
                if (iRet != CRYPT_OK)
                {
                    SysPushErrorCrypto("{}", iRet);
                    return false;
                }

                out.writePtr += primaryLength + 128;
            }
            else
            {
                if (!out.GetOrAllocateLinearWriteable(primaryLength))
                {
                    SysPushErrorMem();
                    return false;
                }

                iRet = cbc_encrypt(tptr, out.writePtr, primaryLength, &cbc);
                if (iRet != CRYPT_OK)
                {
                    SysPushErrorCrypto("{}", iRet);
                    return false;
                }

                out.writePtr += primaryLength;
            }
        }
        else
        {
            if (!out.GetOrAllocateLinearWriteable(primaryLength + 128))
            {
                SysPushErrorMem();
                return false;
            }

            iRet = cbc_encrypt(tptr, out.writePtr, primaryLength, &cbc);
            if (iRet != CRYPT_OK)
            {
                SysPushErrorCrypto("{}", iRet);
                return false;
            }

            AuMemcpy(pad, tptr + primaryLength, overhang);
            pad[127] = padding;

            iRet = cbc_encrypt(pad, out.writePtr + primaryLength, 128, &cbc);
            if (iRet != CRYPT_OK)
            {
                SysPushErrorCrypto("{}", iRet);
                return false;
            }

            out.writePtr += primaryLength + 128;
        }

        return true;
    }

    AUKN_SYM bool Encrypt(const AuMemoryViewRead &plainText,
                          const AuMemoryViewRead &inIv,
                          const AuMemoryViewWrite &outIv,
                          const AuMemoryViewRead &inKey,
                          AuByteBuffer &out,
                          bool auCoolCodePadding)
    {
        symmetric_CBC cbc;
        
    #if !defined(SHIP)
        if (!plainText)
        {
            return false;
        }
        
        if (!inIv)
        {
            return false;
        }

        if (!inKey)
        {
            return false;
        }
    #endif

        if (inIv.length != 16)
        {
            SysPushErrorCrypt("{}", inIv.length);
            return false;
        }

        if (outIv.length)
        {
            if (inIv.length != outIv.length)
            {
                SysPushErrorCrypt("{}", outIv.length);
                return false;
            }
        }

        if ((inKey.length != 16) &&
            (inKey.length != 24) &&
            (inKey.length != 32))
        {
            SysPushErrorCrypt("{}", inKey.length);
            return false;
        }

        auto iRet = cbc_start(::Crypto::gAesCipher,
                              AuReinterpretCast<const unsigned char *>(inIv.ptr),
                              AuReinterpretCast<const unsigned char *>(inKey.ptr),
                              inKey.length,
                              0,
                              &cbc);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        if (auCoolCodePadding)
        {
            if (!EncryptCoolCodePadding(AuReinterpretCast<const unsigned char *>(plainText.ptr),
                                        plainText.length,
                                        inIv.ptr,
                                        outIv.ptr,
                                        inIv.length,
                                        inKey.ptr,
                                        inKey.length,
                                        out,
                                        cbc))
            {
                return false;
            }
        }
        else
        {
            if ((plainText.length & (16 - 1)) != 0)
            {
                SysPushErrorArg("{}", plainText.length);
                return false;
            }

            if (!out.GetOrAllocateLinearWriteable(plainText.length))
            {
                SysPushErrorMem();
                return false;
            }

            iRet = cbc_encrypt(AuReinterpretCast<const unsigned char *>(plainText.ptr),
                               out.writePtr,
                               plainText.length,
                               &cbc);
            if (iRet != CRYPT_OK)
            {
                SysPushErrorCrypt("{}", iRet);
                return false;
            }

            out.writePtr += plainText.length;
        }

        if (outIv.ptr)
        {
            unsigned long ivLength = 16;

            iRet = cbc_getiv(AuReinterpretCast<unsigned char *>(outIv.ptr),
                             &ivLength,
                             &cbc);

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

        return true;
    }

    AUKN_SYM bool Decrypt(const AuMemoryViewRead &cipherText,
                          const AuMemoryViewRead &inIv,
                          const AuMemoryViewWrite &outIv,
                          const AuMemoryViewRead &inKey,
                          AuByteBuffer &plainText,
                          bool safe)
    {


        symmetric_CBC cbc;

    #if !defined(SHIP)
        if (!plainText)
        {
            return false;
        }

        if (!inIv)
        {
            return false;
        }

        if (!inKey)
        {
            return false;
        }
    #endif
        
        if (inIv.length != 16)
        {
            SysPushErrorCrypt("{}", inIv.length);
            return false;
        }

        if (outIv.length)
        {
            if (inIv.length != outIv.length)
            {
                SysPushErrorCrypt("{}", outIv.length);
                return false;
            }
        }

        if ((inKey.length != 16) &&
            (inKey.length != 24) &&
            (inKey.length != 32))
        {
            SysPushErrorCrypt("{}", inKey.length);
            return false;
        }

        auto iRet = cbc_start(::Crypto::gAesCipher,
                              AuReinterpretCast<const unsigned char *>(inIv.ptr),
                              AuReinterpretCast<const unsigned char *>(inKey.ptr),
                              inKey.length,
                              0,
                              &cbc);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        if (!plainText.GetOrAllocateLinearWriteable(cipherText.length))
        {
            SysPushErrorMem();
            return false;
        }

        iRet = cbc_decrypt(AuReinterpretCast<const unsigned char *>(cipherText.ptr),
                           plainText.writePtr,
                           cipherText.length,
                           &cbc);
        if (iRet != CRYPT_OK)
        {
            SysPushErrorCrypt("{}", iRet);
            return false;
        }

        if (safe)
        {
            auto magic = plainText[cipherText.length - 1];
            if (magic <= 128)
            {
                plainText.writePtr += cipherText.length - magic;
            }
            else
            {
                plainText.writePtr += cipherText.length;
            }
        }
        else
        {
            plainText.writePtr += cipherText.length;
        }

        if (outIv.ptr)
        {
            unsigned long ivLength = 16;

            iRet = cbc_getiv(reinterpret_cast<unsigned char *>(outIv.ptr), &ivLength, &cbc);
            if (iRet != CRYPT_OK)
            {
                SysPushErrorCrypt("{}", iRet);
                return false;
            }
        }

        return true;
    }
}