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

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

#include <mbedtls/x509.h>
#include <mbedtls/oid.h>
//#include <mbedtls/certs.h>
#include <mbedtls/x509_crt.h>

#include <mbedtls/oid.h>
#include <mbedtls/asn1.h>

#define AURORA_OID_ID_AD MBEDTLS_OID_PKIX "\x30"
#define AURORA_OID_ID_PE_ONE MBEDTLS_OID_PKIX "\x01"
#define AURORA_OID_ID_PE_ONE_AIA AURORA_OID_ID_PE_ONE "\x01"

#define AURORA_OID_ID_AD_CA_ISSUERS AURORA_OID_ID_AD "\x02"
#define AURORA_OID_ID_AD_OCSP AURORA_OID_ID_AD "\x01"

namespace Aurora::Crypto::X509
{
#pragma region functions copied from mbedtls, modified to do extract the asn fields we care about
    static int x509_get_crt_ext(mbedtls_x509_crt *crt, const char *oid, int oidLength, AuFunction<void(mbedtls_x509_buf &ex, unsigned char **, unsigned char *)> cb);

    static int x509_get_ca_id(mbedtls_x509_crt *crt, AuByteBuffer &key)
    {
        bool ok = false;
        return x509_get_crt_ext(crt, MBEDTLS_OID_AUTHORITY_KEY_IDENTIFIER, sizeof(MBEDTLS_OID_AUTHORITY_KEY_IDENTIFIER) - 1, [&](mbedtls_x509_buf &ex, unsigned char **p, unsigned char *end)
        {
            int ret = 0;
            size_t len;

            if ((ret = mbedtls_asn1_get_tag(p, end, &len,
                MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE)) != 0)
            {
                ok = false;
                return;
            }

            auto tag = **p;
            (*p)++;
            if ((ret = mbedtls_asn1_get_len(p, end, &len)) != 0)
            {
                ok = false;
                return;
            }


            if (!AuTryResize(key, len))
            {
                ok = false;
                return;
            }

            memcpy(key.data(), *p, key.size());
            ok = true;
        }) == 0 && ok;
    }

    static int x509_get_subject_id(mbedtls_x509_crt *crt, AuByteBuffer &key)
    {
        bool ok = false;
        return x509_get_crt_ext(crt, MBEDTLS_OID_SUBJECT_KEY_IDENTIFIER, sizeof(MBEDTLS_OID_SUBJECT_KEY_IDENTIFIER) - 1, [&](mbedtls_x509_buf &ex, unsigned char **p, unsigned char *end)
        {
            int ret = 0;
            size_t len;


            auto tag = **p;
            (*p)++;
            if ((ret = mbedtls_asn1_get_len(p, end, &len)) != 0)
            {
                ok = false;
                return;
            }

            if (!AuTryResize(key, len))
            {
                ok = false;
                return;
            }

            memcpy(key.data(), *p, key.size());
            ok = true;
        }) == 0 && ok;
    }

    static int x509_get_aia(mbedtls_x509_crt *crt, AuList<AuString> &ocsp, AuList<AuString> &caIssers)
    {
        bool ok = false;
        return x509_get_crt_ext(crt, AURORA_OID_ID_PE_ONE_AIA, sizeof(AURORA_OID_ID_PE_ONE_AIA) - 1, [&](mbedtls_x509_buf &ex, unsigned char **p, unsigned char *end)
        {
            int ret = 0;
            size_t len, sublen;


            /*
               AuthorityInfoAccessSyntax  ::=
               SEQUENCE SIZE (1..MAX) OF AccessDescription
            */

            if ((ret = mbedtls_asn1_get_tag(p, end, &len,
                MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE)) != 0)
            {
                ok = false;
                return;
            }

            while (*p != end)
            {
                /*
                   AccessDescription  ::=  SEQUENCE {
                   accessMethod          OBJECT IDENTIFIER,
                   accessLocation        GeneralName  }
                */

                auto endOfResource = *p + len;
                if ((ret = mbedtls_asn1_get_tag(p, endOfResource, &sublen,
                    MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE)) != 0)
                {
                    ok = false;
                    return;
                }

                // access method id
                mbedtls_x509_buf oid = { 0, 0, NULL };
                if ((ret = mbedtls_asn1_get_tag(p, *p + sublen, &oid.len,
                    MBEDTLS_ASN1_OID)) != 0)
                {
                    ok = false;
                    return;
                }

                oid.tag = MBEDTLS_ASN1_OID;
                oid.p = *p;
                *p += oid.len;

                // general name
                mbedtls_x509_buf name = { 0, 0, NULL };

                auto tag = **p;
                (*p)++;
                if ((ret = mbedtls_asn1_get_len(p, endOfResource, &name.len)) != 0)
                {
                    ok = false;
                    return;
                }

                if ((tag & MBEDTLS_ASN1_TAG_CLASS_MASK) !=
                    MBEDTLS_ASN1_CONTEXT_SPECIFIC)
                {
                    ok = false;
                    return;
                }

                name.tag = tag;
                name.p = *p;

                if (oid.len == sizeof(AURORA_OID_ID_AD_OCSP) - 1)
                {
                    if (memcmp(oid.p, AURORA_OID_ID_AD_OCSP, oid.len) == 0)
                    {
                        ocsp.push_back(AuString(reinterpret_cast<const char *>(name.p), reinterpret_cast<const char *>(name.p) + name.len));
                    }
                }

                if (oid.len == sizeof(AURORA_OID_ID_AD_CA_ISSUERS) - 1)
                {
                    if (memcmp(oid.p, AURORA_OID_ID_AD_CA_ISSUERS, oid.len) == 0)
                    {
                        caIssers.push_back(AuString(reinterpret_cast<const char *>(name.p), reinterpret_cast<const char *>(name.p) + name.len));
                    }
                }

                *p += name.len;
            }
        }) == 0 && ok;
    }

    static int x509_get_crt_ext(mbedtls_x509_crt *crt, const char *oid, int oidLength, AuFunction<void(mbedtls_x509_buf &ex, unsigned char **, unsigned char *)> cb)
    {
        int ret = 0;
        size_t len;
        unsigned char *end_ext_data, *start_ext_octet, *end_ext_octet;


        unsigned char *scre = crt->v3_ext.p;
        unsigned char **p = &scre;
        auto end = crt->v3_ext.p + crt->v3_ext.len;

        if ((ret = mbedtls_asn1_get_tag(p, end, &len,
            MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE)) != 0)
            return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS + ret);

        if (end != *p + len)
            return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS +
                   MBEDTLS_ERR_ASN1_LENGTH_MISMATCH);

        while (*p < end)
        {
            /*
            * Extension  ::=  SEQUENCE  {
            *      extnID      OBJECT IDENTIFIER,
            *      critical    BOOLEAN DEFAULT FALSE,
            *      extnValue   OCTET STRING  }
             */
            mbedtls_x509_buf extn_oid = { 0, 0, NULL };
            int is_critical = 0; /* DEFAULT FALSE */
            int ext_type = 0;

            if ((ret = mbedtls_asn1_get_tag(p, end, &len,
                MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE)) != 0)
                return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS + ret);

            end_ext_data = *p + len;

            /* Get extension ID */
            if ((ret = mbedtls_asn1_get_tag(p, end_ext_data, &extn_oid.len,
                MBEDTLS_ASN1_OID)) != 0)
                return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS + ret);

            extn_oid.tag = MBEDTLS_ASN1_OID;
            extn_oid.p = *p;
            *p += extn_oid.len;

            /* Get optional critical */
            if ((ret = mbedtls_asn1_get_bool(p, end_ext_data, &is_critical)) != 0 &&
                (ret != MBEDTLS_ERR_ASN1_UNEXPECTED_TAG))
                return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS + ret);

            /* Data should be octet string type */
            if ((ret = mbedtls_asn1_get_tag(p, end_ext_data, &len,
                MBEDTLS_ASN1_OCTET_STRING)) != 0)
                return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS + ret);

            start_ext_octet = *p;
            end_ext_octet = *p + len;

            if (end_ext_octet != end_ext_data)
                return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS +
                       MBEDTLS_ERR_ASN1_LENGTH_MISMATCH);


            if (extn_oid.len != oidLength)
            {
                *p = end_ext_octet;
                continue;
            }

            if (memcmp(extn_oid.p, oid, oidLength))
            {
                *p = end_ext_octet;
                continue;
            }

            cb(extn_oid, p, end_ext_octet);
            return 0;
        }

        if (*p != end)
            return(MBEDTLS_ERR_X509_INVALID_EXTENSIONS +
                   MBEDTLS_ERR_ASN1_LENGTH_MISMATCH);

        return MBEDTLS_ERR_X509_INVALID_EXTENSIONS;
    }


    template<int C>
    static bool find_oid_value_in_name(const mbedtls_x509_name *name, const char(&oid)[C], AuString &value)
    {
        const char *short_name = NULL;
        size_t retval = 0;

        while (name != NULL)
        {
            if (!name->oid.p)
            {
                name = name->next;
                continue;
            }

            if (name->oid.len != (C - 1))
            {
                name = name->next;
                continue;
            }

            if (memcmp(oid, name->oid.p, (C - 1)))
            {
                name = name->next;
                continue;
            }

            value = AuString(reinterpret_cast<const char *>(name->val.p), reinterpret_cast<const char *>(name->val.p) + name->val.len);
            return true;
        }

        return false;
    }
#pragma endregion

    static void FindCommonNames(const mbedtls_x509_name &name, CertificateName &out)
    {
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_CN, out.commonName);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_COUNTRY, out.countryCode);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_ORGANIZATION, out.organization);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_ORG_UNIT, out.department);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_STATE, out.state);
        find_oid_value_in_name(&name, MBEDTLS_OID_PKCS9_EMAIL, out.email);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_TITLE, out.title);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_GIVEN_NAME, out.name);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_POSTAL_CODE, out.postcode);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_POSTAL_ADDRESS, out.address);
        find_oid_value_in_name(&name, MBEDTLS_OID_AT_LOCALITY, out.locality);
    }

    bool ParseCert(const AuMemoryViewRead &der, const AuFunction<void(mbedtls_x509_crt &crt)> &cb)
    {
        bool bRet {};
        mbedtls_x509_crt crt {};

        mbedtls_x509_crt_init(&crt);

        auto status = mbedtls_x509_crt_parse(&crt,
                                             der.Begin<unsigned char>(),
                                             der.Size());

        if (status < 0)
        {
            goto out;
        }

        try
        {
            if (cb)
            {
                cb(crt);
            }
            bRet = true;
        }
        catch (...)
        {
            SysPushErrorCatch();
        }

    out:
        mbedtls_x509_crt_free(&crt);
        return bRet;
    }

    AuInt64 ConvertTime(const mbedtls_x509_time &time)
    {
        AuTime::tm tm = {};
        tm.year = time.year;
        tm.mon  = time.mon - 1;
        tm.mday = time.day - 1;
        tm.sec  = time.sec;
        tm.min  = time.min;
        tm.hour = time.hour;
        return AuTime::FromCivilTime(tm, AuTime::ETimezoneShift::eUTC);
    }

    static void FindUsage(CertificateDecoded &out,
                          const mbedtls_x509_sequence *extended_key_usage)
    {
        const mbedtls_x509_sequence *pCur = extended_key_usage;

        while (pCur)
        {

        #define CHECK_CHECK(enum, oid)                          \
            if (pCur->buf.len == sizeof(oid) - 1)               \
            {                                                   \
                if (!AuMemcmp(oid, pCur->buf.p, pCur->buf.len)) \
                {                                               \
                    out.usage.push_back(enum);                  \
                }                                               \
            }

            CHECK_CHECK(EExtendedUsage::eServerAuth, MBEDTLS_OID_SERVER_AUTH);
            CHECK_CHECK(EExtendedUsage::eClientAuth, MBEDTLS_OID_CLIENT_AUTH);
            CHECK_CHECK(EExtendedUsage::eCodeSigning, MBEDTLS_OID_CODE_SIGNING);
            CHECK_CHECK(EExtendedUsage::eEmailProtection, MBEDTLS_OID_EMAIL_PROTECTION);
            CHECK_CHECK(EExtendedUsage::eTimeStamping, MBEDTLS_OID_TIME_STAMPING);
            CHECK_CHECK(EExtendedUsage::eOCSPSigning, MBEDTLS_OID_OCSP_SIGNING);

        #undef CHECK_CHECK
         
            pCur = pCur->next;
        }
    }

    bool DecodeInternal(const mbedtls_x509_crt &crt, CertificateDecoded &out)
    {
        auto &issuer = crt.issuer;
        auto &subject = crt.subject;

        out.version  = crt.version;
        out.iMaxPath = crt.private_max_pathlen;
        out.bIsCA    = crt.private_ca_istrue;

        FindUsage(out, &crt.ext_key_usage);

        FindCommonNames(issuer, out.issuer);
        FindCommonNames(subject, out.subject);

        out.validity.issued = ConvertTime(crt.valid_from);
        out.validity.expire = ConvertTime(crt.valid_to);

        x509_get_ca_id((mbedtls_x509_crt *)&crt, out.issuer.id);
        x509_get_subject_id((mbedtls_x509_crt *)&crt, out.subject.id);
        x509_get_aia((mbedtls_x509_crt *)&crt, out.OCSPs, out.AIAs);

        if (out.serialNumber.Write(crt.serial.p, crt.serial.len) != crt.serial.len)
        {
            return false;
        } 

        if (out.algorithmOid.Write(crt.sig_oid.p, crt.sig_oid.len) != crt.sig_oid.len)
        {
            return false;
        } 

        if (out.publicKey.Write(crt.pk_raw.p, crt.pk_raw.len) != crt.pk_raw.len)
        {
            return false;
        } 

        return true;
    }

    AUKN_SYM bool Decode(const AuMemoryViewRead &der, CertificateDecoded &out)
    {
        bool bRet { true };
        return ParseCert(der,
                         [&](mbedtls_x509_crt &crt)
        {
            bRet = DecodeInternal(crt, out);
        }) && bRet;
    }

    static bool IsHighRiskStateIssuer(const mbedtls_x509_crt &ca)
    {
        AuString issuer;
        
        if (!find_oid_value_in_name(&ca.issuer, MBEDTLS_OID_AT_COUNTRY, issuer))
        {
            // FAIL!
            return true;
        }

        if (issuer.empty())
        {
            // FAIL!
            return true;
        }

        // TODO: removed <on pastebin>
        // We need some kind of shim objects for IPinCertificate &
        //         some kind of plan on how intrastructure that has to trust public exchange points will deal with state sponsored attacks.
        // this old check isnt it.

        return false;
    }

    AUKN_SYM bool Validate(const AuMemoryViewRead &der, const AuMemoryViewRead &parentDer)
    {
        bool failed = false;

        // gross
        return ParseCert(der,
                         [&](mbedtls_x509_crt &crt)
        {
            ParseCert(parentDer,
                      [&](mbedtls_x509_crt &ca)
            {

                if ((failed |= IsHighRiskStateIssuer(ca)))
                {
                    return;
                }

                uint32_t flags {};
                failed |= 0 > mbedtls_x509_crt_verify_restartable(&crt, &ca, NULL, &mbedtls_x509_crt_profile_default, NULL, &flags, NULL, NULL, NULL);
                failed |= flags != 0;
            });
        }) && !failed;
    }
}