/***
    Copyright (C) 2024 Jamie Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuCertificateStore.cpp
    Date: 2024-10-14
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuCertificateStore.hpp"
#include <mbedtls/x509.h>
#include <mbedtls/x509_crt.h>
#include <Source/Crypto/X509/AuCertificateChain.hpp>

namespace Aurora::Crypto::X509
{
    AuInt64 ConvertTime(const mbedtls_x509_time &time);
}

namespace Aurora::Crypto::CA
{
    // Digicerts "high assurance" and other special roots trusted by payment insitutions use 2048... FUCKING WHY?
    // These certs should be expiring in early 2030. Hopefully PayPal, Stripe, and others will get their act
    // together soon. Oh, and guess what, Amazon and tons of AWS customers use Amazon Root CA 1 with a 2048-bit PK as well.
    static const mbedtls_x509_crt_profile kAssumeRootIsSmart =
    {
        MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA1) |       // TODO(by 2030): delete me
        MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA256) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA384) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA512),
        0xFFFFFFF, /* Any PK alg    */
    #if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP256R1) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP384R1) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP521R1) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_BP256R1) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_BP384R1) |
        MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_BP512R1) |
        0,
    #else
        0,
    #endif
        2048, // TODO(by 2030): Raise to 4k
    };

    bool CertificateStore::CheckCertificate(const AuSPtr<X509::ICertificateChain> &pChain,
                                            const AuMemoryViewRead &derCertificate)
    {
        AU_LOCK_GUARD(this->rwLock->AsReadable());
        bool bFoundTrustedCA {};

        if (!pChain)
        {
            return false;
        }

        auto uCount = pChain->GetCertificateCount();
        for (AU_ITERATE_N(i, uCount))
        {
            auto pCurrentCert = AuStaticCast<AuCrypto::X509::CertificateChain>(pChain)->GetCertificateInternal(i);
            if (!pCurrentCert)
            {
                return false;
            }

            if (i)
            {
                auto pPrevCert = AuStaticCast<AuCrypto::X509::CertificateChain>(pChain)->GetCertificateInternal(i - 1);
                if (!pPrevCert)
                {
                    return false;
                }

                uint32_t flags {};
                if (mbedtls_x509_crt_verify_restartable(pCurrentCert, pPrevCert, NULL, &kAssumeRootIsSmart, NULL, &flags, NULL, NULL, NULL) < 0)
                {
                    SysPushErrorCrypto("Bad certificate chain");
                    return false;
                }

                if (flags != 0)
                {
                    SysPushErrorCrypto("Bad certificate chain");
                    return false;
                }
            }
 
            auto uCode = AuFnv1a32Runtime(pCurrentCert->pk_raw.p, pCurrentCert->pk_raw.len);

            auto uItrA = this->storage.find(uCode);
            if (uItrA == this->storage.end())
            {
                continue;
            }
 
            for (const auto &[that, iDate] : uItrA->second)
            {
                AuMemoryViewRead toCompare   { that };
                AuMemoryViewRead currentCert { pCurrentCert->pk_raw.p, pCurrentCert->pk_raw.len };

                if (toCompare.uLength != currentCert.uLength)
                {
                    continue;
                }

                if (AuTime::CurrentClockMS() > iDate)
                {
                    SysPushErrorCrypto("Expired Certificate Authority");
                    continue;
                }

                if (AuMemcmp(toCompare.Begin(),
                             currentCert.Begin(),
                             currentCert.uLength) == 0)
                {
                    bFoundTrustedCA = true;
                    break;
                }
            }

            if (bFoundTrustedCA)
            {
                break;
            }
        }

        return bFoundTrustedCA;
    }

    bool CertificateStore::AddCertificate(const AuMemoryViewRead &x509Certificate)
    {
        mbedtls_x509_crt crt;
        if (!x509Certificate)
        {
            SysPushErrorArg();
            return false;
        }

        mbedtls_x509_crt_init(&crt);

        if (mbedtls_x509_crt_parse(&crt,
                                   x509Certificate.Begin<unsigned char>(),
                                   x509Certificate.Size()) < 0)
        {
            SysPushErrorGeneric();
            return false;
        }

        {
            AuUInt32     uCode  { AuFnv1a32Runtime(crt.pk_raw.p, crt.pk_raw.len)    };
            AuByteBuffer buffer { AuMemoryViewRead { crt.pk_raw.p, crt.pk_raw.len } };
            
            if (!buffer)
            {
                SysPushErrorMemory();
                mbedtls_x509_crt_free(&crt);
                return false;
            }
         
            AU_LOCK_GUARD(this->rwLock->AsWritable());
            this->storage[uCode].push_back(AuMakePair(AuMove(buffer), X509::ConvertTime(crt.valid_to)));
        }

        mbedtls_x509_crt_free(&crt);
        return true;
    }

    bool CertificateStore::AddCertificateChain(X509::ICertificateChain *pChain)
    {
        if (!pChain)
        {
            SysPushErrorArg();
            return false;
        }

        AU_LOCK_GUARD(this->rwLock->AsWritable());

        auto uCount = pChain->GetCertificateCount();
        for (AU_ITERATE_N(i, uCount))
        {
            auto pCert = AuStaticCast<AuCrypto::X509::CertificateChain>(pChain)->GetCertificateInternal(i);
            if (!pCert)
            {
                SysPushErrorMemory();
                return false;
            }

            AuUInt32     uCode  { AuFnv1a32Runtime(pCert->pk_raw.p, pCert->pk_raw.len)    };
            AuByteBuffer buffer { AuMemoryViewRead { pCert->pk_raw.p, pCert->pk_raw.len } };
            
            if (!buffer)
            {
                SysPushErrorMemory();
                return false;
            }

            this->storage[uCode].push_back(AuMakePair(AuMove(buffer), X509::ConvertTime(pCert->valid_to)));
        }

        return true;
    }

    void CertificateStore::Serialize(Memory::ByteBuffer &buffer)
    {
        for (auto &[uID, storageList] : this->storage)
        {
            for (auto &[cert, uValidTo] : storageList)
            {
                buffer.Write<AuUInt32>(cert.size());
                buffer.Write<AuInt64>(uValidTo);
                buffer.Write(AuMemoryViewRead { cert });
            }
        }
    }

    bool CertificateStore::Deserialize(Memory::ByteBuffer &buffer)
    {
        AU_LOCK_GUARD(this->rwLock->AsWritable());

        while (buffer.RemainingBytes())
        {
            auto uLength = buffer.Read<AuUInt32>();
            if (!uLength)
            {
                return false;
            }

            auto iValidTo = buffer.Read<AuInt64>();
          
            auto view = buffer.GetLinearReadableForExactly(uLength);
            if (!view)
            {
                return false;
            }
            buffer.readPtr += uLength;

            AuUInt32     uCode  { AuFnv1a32Runtime(view.Begin(), view.Size()) };
            AuByteBuffer buffer { view                                        };
            
            if (!buffer)
            {
                SysPushErrorMemory();
                return false;
            }

            this->storage[uCode].push_back(AuMakePair(AuMove(buffer), iValidTo));
        }

        return true;
    }

    AUKN_SYM ICertificateStore *NewCertificateStoreNew()
    {
        return _new CertificateStore();
    }

    AUKN_SYM void NewCertificateStoreRelease(ICertificateStore *pHandle)
    {
        AuSafeDelete<CertificateStore *>(pHandle);
    }
}