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

    File: LTCImportEx.c
    Date: 2021-6-24
    Author: Reece
***/
#include <tomcrypt.h>
#include "tomcrypt_private.h"
#include "LTCExtensions.h"

int rsa_import_ex(const unsigned char *in, unsigned long inlen, rsa_key *key, int flags)
{
    int           err;
    void *zero;
    char scratch[4096];
    const unsigned char *rsa = scratch;
    unsigned long rsa_len = sizeof(scratch);
    unsigned long  len;

    LTC_ARGCHK(in != NULL);
    LTC_ARGCHK(key != NULL);
    LTC_ARGCHK(ltc_mp.name != NULL);

    /* init key */
    if ((err = rsa_init(key)) != CRYPT_OK)
    {
        return err;
    }

    len = 0;

    if ((flags & kRsaFlagPKCS1) != 0)
    {
        rsa = in;
        rsa_len = inlen;
    }
    else
    {
        unsigned long oid[16];
        const char *rsaoid;
        ltc_asn1_list alg_seq[2], top_seq[3];
        ltc_asn1_list *l = NULL;

        /* get RSA alg oid */
        err = pk_get_oid(PKA_RSA, &rsaoid);
        if (err != CRYPT_OK)
        {
            goto LBL_ERR;
        }

        if ((err = mp_init(&zero)) != CRYPT_OK)
        {
            goto LBL_ERR;
        }

        LTC_SET_ASN1(alg_seq, 0, LTC_ASN1_OBJECT_IDENTIFIER, oid, 16UL);
        LTC_SET_ASN1(alg_seq, 1, LTC_ASN1_NULL, NULL, 0UL);// << key
        if ((flags & kRsaFlagPublic) == 0)
        {
            LTC_SET_ASN1(top_seq, 0, LTC_ASN1_INTEGER, zero, 1UL);
            LTC_SET_ASN1(top_seq, 1, LTC_ASN1_SEQUENCE, alg_seq, 2UL);
            LTC_SET_ASN1(top_seq, 2, LTC_ASN1_OCTET_STRING, rsa, rsa_len);

            err = der_decode_sequence(in, inlen, top_seq, 3UL);
        }
        else
        {
            LTC_SET_ASN1(top_seq, 0, LTC_ASN1_SEQUENCE, alg_seq, 1UL);
            LTC_SET_ASN1(top_seq, 1, LTC_ASN1_BIT_STRING, rsa, rsa_len);

            err = der_decode_sequence(in, inlen, top_seq, 2UL);
        }

        mp_clear(zero);

        rsa_len = top_seq[2].size;

        if (err != CRYPT_OK)
        {
            goto LBL_ERR;
        }

        /* check alg oid */
        if ((err = pk_oid_cmp_with_asn1(rsaoid, &alg_seq[0])) != CRYPT_OK)
        {
            goto LBL_ERR;
        }
    }

    if ((flags & kRsaFlagPublic) != 0)
    {
        /* it's a public key and we lack e */
        if ((err = der_decode_sequence_multi(rsa, rsa_len,
            LTC_ASN1_INTEGER, 1UL, key->N,
            LTC_ASN1_INTEGER, 1UL, key->e,
            LTC_ASN1_EOL, 0UL, NULL)) != CRYPT_OK)
        {
            goto LBL_ERR;
        }
        key->type = PK_PUBLIC;
    }
    else
    {
        err = der_decode_sequence_multi(rsa, rsa_len, LTC_ASN1_INTEGER, 1UL, key->N,
                                        LTC_ASN1_EOL, 0UL, NULL);

        if (err != CRYPT_OK && err != CRYPT_INPUT_TOO_LONG)
        {
            goto LBL_ERR;
        }

        if (mp_cmp_d(key->N, 0) == LTC_MP_EQ)
        {
            if ((err = mp_init(&zero)) != CRYPT_OK)
            {
                goto LBL_ERR;
            }

            /* it's a private key */
            if ((err = der_decode_sequence_multi(rsa, rsa_len,
                LTC_ASN1_INTEGER, 1UL, zero,
                LTC_ASN1_INTEGER, 1UL, key->N,
                LTC_ASN1_INTEGER, 1UL, key->e,
                LTC_ASN1_INTEGER, 1UL, key->d,
                LTC_ASN1_INTEGER, 1UL, key->p,
                LTC_ASN1_INTEGER, 1UL, key->q,
                LTC_ASN1_INTEGER, 1UL, key->dP,
                LTC_ASN1_INTEGER, 1UL, key->dQ,
                LTC_ASN1_INTEGER, 1UL, key->qP,
                LTC_ASN1_EOL, 0UL, NULL)) != CRYPT_OK)
            {
                mp_clear(zero);
                goto LBL_ERR;
            }

            mp_clear(zero);

            key->type = PK_PRIVATE;
        }
        else if (mp_cmp_d(key->N, 1) == LTC_MP_EQ)
        {
            /* we don't support multi-prime RSA */
            err = CRYPT_PK_INVALID_TYPE;
            goto LBL_ERR;
        }
        else
        {
            err = CRYPT_PK_INVALID_TYPE;
            goto LBL_ERR;
        }
    }

    return CRYPT_OK;

LBL_ERR:
    rsa_free(key);
    return err;
}