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

    File: x509.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, AuList<AuUInt8> &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, AuList<AuUInt8> &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, std::function<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, CertName &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);
    }

    static bool ParseCert(const Certificate &der, AuFunction<void(mbedtls_x509_crt &crt)> cb)
    {
        bool ret = false;
        mbedtls_x509_crt crt {};

        mbedtls_x509_crt_init(&crt);

        auto status = mbedtls_x509_crt_parse(&crt,
                                             reinterpret_cast<const unsigned char *>(der.data()),
                                             der.size());

        if (status < 0) goto out;

        cb(crt);

        ret = true;
    out:
        mbedtls_x509_crt_free(&crt);
        return ret;
    }

    static AuInt64 ConvertTime(const mbedtls_x509_time &time)
    {
        AuTime::tm tm = {};
        tm.tm_year = time.year - 1900;
        tm.tm_mon  = time.mon - 1;
        tm.tm_mday = time.day;
        tm.tm_sec  = time.sec;
        tm.tm_min  = time.min;
        tm.tm_hour = time.hour;
        return AuTime::FromCivilTime(tm, true);
    }

    AUKN_SYM bool Decode(const Certificate &der, DecodedCertificate &out)
    {
        bool failed = false;

        return ParseCert(der,
                         [&](mbedtls_x509_crt &crt)
        {
            auto &issuer = crt.issuer;
            auto &subject = crt.subject;

            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(&crt, out.issuer.id);
            x509_get_subject_id(&crt, out.subject.id);

            out.serialNumber.resize(crt.serial.len);
            memcpy(out.serialNumber.data(), crt.serial.p, out.serialNumber.size());

            out.algorithmOid.resize(crt.sig_oid.len);
            memcpy(out.algorithmOid.data(), crt.sig_oid.p, out.algorithmOid.size());

            AuList<AuString> oscp;
            x509_get_aia(&crt, oscp, out.AIAs);
        }) && !failed;

    }

    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;
        }

        // Winnie the Pooh has been caught with his paws in the honey jar one too many times
        // At least other nation state attacks have been covert enough for us to not know about any *TLS* MITM attacks originating from them
        // Pro-china cuckolds at mozilla suggested we merely block certs issued issued by compromised cn CAs using their issue date. lol no
        // https://arstechnica.com/information-technology/2015/03/google-warns-of-unauthorized-tls-certificates-trusted-by-almost-all-oses/
        // https://www.zdnet.com/article/china-is-now-blocking-all-encrypted-https-traffic-using-tls-1-3-and-esni/
        // https://www.popularmechanics.com/military/news/a28510/china-secret-plan-invade-taiwan/
        // https://www.theguardian.com/world/2021/mar/10/china-could-invade-taiwan-in-next-six-years-top-us-admiral-warns
        // https://en.wikipedia.org/wiki/Lazarus_Group (use chinese connections)
        if ((!gRuntimeConfig.crypto.allowChineseCerts) && 
            (stricmp(issuer.c_str(), "cn") == 0) ||  // mainland and countries the communists decided are theirs. often used by the dprk
            (stricmp(issuer.c_str(), "mo") == 0) ||  // macau, china special admin region
            (stricmp(issuer.c_str(), "hk") == 0) ||  // hong kong 
            /*(stricmp(issuer.c_str(), "tw") == 0)*/ false) // tehe yes we will be ready to invade by 2020. i suppose taiwan is safe enough for.now
        {
            SysPushErrorCrypt("The funny western yellow bear in the east is behind you");
            return true;
        }

        // The intersection between technical threats and countries we can't do business with legally
        // https://en.wikipedia.org/wiki/United_States_sanctions
        // https://en.wikipedia.org/wiki/Communications_in_Iran
        // https://www.gov.uk/government/collections/financial-sanctions-regime-specific-consolidated-lists-and-releases
        // https://en.wikipedia.org/wiki/Lazarus_Group
        // https://www.bbc.co.uk/news/stories-57520169
        if ((stricmp(issuer.c_str(), "ir") == 0) ||  // iran - has own intranet. IXPs are owned by the state
            (stricmp(issuer.c_str(), "iq") == 0) ||  // iraq
            (stricmp(issuer.c_str(), "kp") == 0))    // north korea - has own intranet. IXPs are owned by the state
        {
            SysPushErrorCrypt("Service is unavailable in your country for legal and/or technical safety concerns");
            return true;
        }
        
        // Russia is somewhat misunderstood and have their hands tied in a lot of cases, not to excuse blatent corruption
        // Moscow wants MSK-IX IXP to be independent and have ran tests to validate their intranet could work without the rest of the world
        // I wonder if Pavel Durov would trust russian certificates, probably not, right? Major states, especially Russia, shouldn't be trusted with foreign services
        // allowRussianCerts is true by default; however, services targeting a market of close political enemies of russia should consider disallowing russian certs
        // Further, private infrastructure around Russian territories should consider a heightened global pinning security policy
        if ((!gRuntimeConfig.crypto.allowRussianCerts) && 
            (stricmp(issuer.c_str(), "ru") == 0))   
        {
            SysPushErrorCrypt("Service is unavailable in your country for technical safety concerns");
            return true;
        }

        // The 5 eyes are smart enough to know MITM attacks would send alarms ringing at major services providers 
        // They are known to MITM plain text traffic and hoard vulns; however, we can design around these issues
        // Do not use NIST curves, do not use plain text across public IXPs, don't save data in plain text in EU/US datacenters, etc
        // There is nothing we can do about western powers. They seem to be playing... not fair... but smart
        {
            // obama bin listening
        }

        return false;
    }

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

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

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

                failed = 0 > mbedtls_x509_crt_verify_restartable(&crt, &ca, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
            });
        }) && !failed;
    }
}