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

    File: GenerateCertificate.cpp
    Date: 2022-11-18
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "../Crypto.hpp"
#include "x509.hpp"
#include <Source/IO/TLS/TLS.hpp>
#include <Source/IO/TLS/TLSPrivateKeyPair.hpp>
#include <mbedtls/asn1write.h>
#include <mbedtls/oid.h>

namespace Aurora::Crypto::X509
{
    AUKN_SYM AuResult<Memory::ByteBuffer> GenerateCertificate(const CertRequest &request)
    {
        int iRet {};
        mbedtls_x509write_cert crt;
        mbedtls_mpi serial;
        AuByteBuffer bufPb;
        mbedtls_pk_context ctxPb {};
        mbedtls_md_type_t hash {};
        AuByteBuffer bufPv;
        mbedtls_pk_context ctxPv {};
        Memory::ByteBuffer buffer;

        if (request.pRSAKey && request.pECCKey)
        {
            SysPushErrorArg("A certificate may only contain one key");
            return {};
        }

        if (!request.pSigningChain)
        {
            if (!(bool(request.pRSAKey) == bool(request.pSelfSigningRSAKey) &&
                  bool(request.pECCKey) == bool(request.pSelfSigningECCKey)))
            {
                SysPushErrorArg("Missing or conflicting private/public key types");
                return {};
            }
        }
        else
        {
            if (request.pSelfSigningRSAKey || request.pSelfSigningECCKey)
            {
                SysPushErrorArg("Missing or conflicting private/public key types");
                return {};
            }
        }

        if (request.uVersion > 3)
        {
            SysPushErrorArg("Invalid version");
            return {};
        }

        ::mbedtls_x509write_crt_init(&crt);
        ::mbedtls_mpi_init(&serial);
        ::mbedtls_x509write_crt_set_version(&crt, request.uVersion - 1);

        {
            if (request.pRSAKey)
            {
                if (!request.pRSAKey->ToKey(RSA::ERSAKeyType::eRsaKey, bufPb))
                {
                    SysPushErrorNested();
                    goto out;
                }
            }

            if (request.pECCKey)
            {
                if (!request.pECCKey->AsPublicECC(bufPb))
                {
                    SysPushErrorNested();
                    goto out;
                }
            }

            if ((iRet = ::mbedtls_pk_parse_public_key(&ctxPb, bufPb.readPtr, bufPb.RemainingBytes())) != 0)
            {
                SysPushErrorCrypto("Couldn't import public key: {}", iRet);
                goto out;
            }

            ::mbedtls_x509write_crt_set_subject_key(&crt, &ctxPb);
        }

        {

            if (request.pSigningChain)
            {
                ctxPv = AuStaticCast<IO::TLS::TLSPrivateKeyPairImpl>(request.pSigningChain)->GetInternal();
            }
            else
            {

                if (request.pSelfSigningRSAKey)
                {
                    RSA::RSAMeta meta;
                    meta.encoding = RSA::ERSAKeyType::eKey;
                    meta.type = AuCrypto::EKeyType::eKeyPrivate;
                    if (!request.pSelfSigningRSAKey->ToKey(meta, bufPv))
                    {
                        SysPushErrorNested();
                        goto out;
                    }
                }

                if (request.pSelfSigningECCKey)
                {
                    if (!request.pSelfSigningECCKey->AsPrivateECC(bufPv))
                    {
                        SysPushErrorNested();
                        goto out;
                    }
                }
                
                if ((iRet = ::mbedtls_pk_parse_key(&ctxPv,
                                                  bufPv.base, bufPv.RemainingBytes(),
                                                  nullptr, 0,
                                                  mbedtls_ctr_drbg_random,
                                                  &IO::TLS::gCtrDrbg)) != 0)
                {
                    SysPushErrorCrypt("Couldn't parse private issuer key: {}", iRet);
                    goto out;
                }

                ::mbedtls_x509write_crt_set_issuer_key(&crt, &ctxPv);
            }
        }
        
    #define _WRITE_ISSUER(oid, expression)                                                        \
        if (expression.size())                                                                    \
        {                                                                                         \
            mbedtls_asn1_named_data* pCur;                                                        \
            if (!(pCur =::mbedtls_asn1_store_named_data(&crt.private_subject,                     \
                                                        oid,                                      \
                                                        strlen(oid),                              \
                                                        (const unsigned char *)expression.c_str(),\
                                                        expression.size())))                      \
            {                                                                                     \
                goto out;                                                                        \
            }                                                                                     \
            pCur->val.tag = MBEDTLS_ASN1_UTF8_STRING;                                             \
        }

        _WRITE_ISSUER(MBEDTLS_OID_AT_CN, request.name.commonName);
        _WRITE_ISSUER(MBEDTLS_OID_AT_COUNTRY, request.name.countryCode);
        _WRITE_ISSUER(MBEDTLS_OID_AT_ORGANIZATION, request.name.organization);
        _WRITE_ISSUER(MBEDTLS_OID_AT_ORG_UNIT, request.name.department);
        _WRITE_ISSUER(MBEDTLS_OID_AT_STATE, request.name.state);
        _WRITE_ISSUER(MBEDTLS_OID_PKCS9_EMAIL, request.name.email);
        _WRITE_ISSUER(MBEDTLS_OID_AT_TITLE, request.name.title);
        _WRITE_ISSUER(MBEDTLS_OID_AT_GIVEN_NAME, request.name.name);
        _WRITE_ISSUER(MBEDTLS_OID_AT_POSTAL_CODE, request.name.postcode);
        _WRITE_ISSUER(MBEDTLS_OID_AT_POSTAL_ADDRESS, request.name.address);
        _WRITE_ISSUER(MBEDTLS_OID_AT_LOCALITY, request.name.locality);
    #undef _WRITE_ISSUER

        if (request.pSigningChain)
        {
            crt.private_issuer = &AuStaticCast<IO::TLS::CertificateChain>(request.pSigningChain->GetChain())->pCertificate->issuer;
        }
        else
        {
            crt.private_issuer = crt.private_subject;
        }

        {
            auto issueNorm = AuTime::ToCivilTime(request.iIssuedDateMs, AuTime::ETimezoneShift::eUTC);
            auto issue = fmt::format("{:04}{:02}{:02}{:02}{:02}{:02}",
                               issueNorm.tm_year + 1900, issueNorm.tm_mon + 1, issueNorm.tm_mday,
                               issueNorm.tm_hour, issueNorm.tm_min, issueNorm.tm_sec);

            auto expireNorm = AuTime::ToCivilTime(request.iExpirationDateMs, AuTime::ETimezoneShift::eUTC);
            auto expire = fmt::format("{:04}{:02}{:02}{:02}{:02}{:02}",
                               expireNorm.tm_year + 1900, expireNorm.tm_mon + 1, expireNorm.tm_mday,
                               expireNorm.tm_hour, expireNorm.tm_min, expireNorm.tm_sec);
            if (mbedtls_x509write_crt_set_validity(&crt, issue.c_str(), expire.c_str()) != 0)
            {
                SysPushErrorCrypto("Couldn't update certificate validity");
                goto out;
            }
        }

        hash = ::Crypto::TypeToMbed(request.digest);
        if (hash == mbedtls_md_type_t::MBEDTLS_MD_NONE)
        {
            SysPushErrorCrypto("Unsupported digest: {}", AuHashing::EHashTypeToString(request.digest));
            goto out;
        }

        ::mbedtls_x509write_crt_set_md_alg(&crt, hash);

        if (::mbedtls_mpi_read_string(&serial, request.uSerialRadix, request.sSerial.c_str()) != 0)
        {
            SysPushErrorCrypto("Couldn't parse serial");
            goto out;
        }

        if (::mbedtls_x509write_crt_set_serial(&crt, &serial) != 0)
        {
            SysPushErrorCrypto("Couldn't set serial");
            goto out;
        }
        
        if (request.bSubjectKeyId && request.uVersion == 3)
        {
            if (::mbedtls_x509write_crt_set_subject_key_identifier(&crt) != 0)
            {
                SysPushErrorCrypto("Couldn't set subject key identifier");
                goto out;
            }
        }

        if (request.bIsCA && request.uVersion == 3)
        {
            if (::mbedtls_x509write_crt_set_authority_key_identifier(&crt) != 0)
            {
                SysPushErrorCrypto("Couldn't set authority key identifier");
                goto out;
            }
        }

        if (request.uVersion == 3)
        {
            if (::mbedtls_x509write_crt_set_basic_constraints(&crt, request.bIsCA, request.iMaxPathLength) != 0)
            {
                SysPushErrorCrypto("Couldn't set basic constraints");
                goto out;
            }
        }

        if (request.usage.size())
        {
            mbedtls_asn1_sequence *tail {};

            for (const auto A : request.usage)
            {
            #define SET_OID(x, oid) \
    do { x.len = MBEDTLS_OID_SIZE(oid); x.p = (unsigned char*)oid; x.tag = MBEDTLS_ASN1_OID; } while( 0 )

                auto ext_key_usage = (mbedtls_asn1_sequence *)::Aurora::Memory::_ZAlloc(sizeof(mbedtls_asn1_sequence));
                if (tail)
                {
                    ext_key_usage->next = tail;
                }
                tail = ext_key_usage;

                switch (A)
                {
                case EExtendedUsage::eServerAuth:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_SERVER_AUTH);
                    break;
                case EExtendedUsage::eClientAuth:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_CLIENT_AUTH);
                    break;
                case EExtendedUsage::eCodeSigning:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_CODE_SIGNING);
                    break;
                case EExtendedUsage::eEmailProtection:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_EMAIL_PROTECTION);
                    break;
                case EExtendedUsage::eTimeStamping:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_TIME_STAMPING);
                    break;
                case EExtendedUsage::eOCSPSigning:
                    SET_OID(ext_key_usage->buf, MBEDTLS_OID_OCSP_SIGNING);
                    break;
                default:
                {
                    tail = nullptr;
                    break;
                }
                }
            }
        
            if (tail)
            {
                iRet = ::mbedtls_x509write_crt_set_ext_key_usage(&crt, tail);
                if (iRet != 0)
                {
                    SysPushErrorCrypto("Couldn't write usage {}", iRet);
                    goto out;
                }
            }
        }

        if (!AuTryResize(buffer, 10 * 1024))
        {
            goto out;
        }
        
        iRet = ::mbedtls_x509write_crt_der(&crt,
                                           buffer.base,
                                           buffer.length,
                                           ::mbedtls_ctr_drbg_random,
                                           &IO::TLS::gCtrDrbg);
        if (iRet < 0)
        {
            SysPushErrorCrypto("Couldn't write x509 cert: {}", iRet);
            goto out;
        }

        AuMemmove(buffer.base, buffer.base + buffer.length - iRet, iRet);
        buffer.writePtr += iRet;

        crt.private_issuer = nullptr;
        mbedtls_x509write_crt_free(&crt);
        mbedtls_pk_free(&ctxPv);
        mbedtls_pk_free(&ctxPb);
        return buffer;

    out:

        crt.private_issuer = nullptr;
        mbedtls_x509write_crt_free(&crt);
        mbedtls_pk_free(&ctxPv);
        mbedtls_pk_free(&ctxPb);
        return {};
    }
}