/**
 * @file sdrv_crypto_rsapadding.c
 * @brief crypto rsapadding interface
 *
 * @copyright Copyright (c) 2021  Semidrive Semiconductor.
 *            All rights reserved.
 */

#include <sdrv_crypto_rsapadding.h>

sdrv_crypto_error_status_e gen_rnd_bytes(uint8_t *dst, uint32_t size)
{
    sdrv_crypto_error_status_e RetVal;
    uint8_t *Random_String;

    uint32_t i = 0;
    uint32_t index = 0;
    uint32_t index_max = 0;
    uint32_t cp_left = 0;

    if (dst == NULL) {
        return E_PARAM_ERROR;
    }

    Random_String = ShareMem_GetBlock(CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX);

    index = 0;
    /*64 byte one block,CRYPTO_SHARE_MEM_NODE_COMM_SMALL_BUFF_MAX*/
    index_max = size >> 6;
    cp_left = size & 0x3f;

    for (i = 0; i < index_max; i++) {
        arch_invalidate_cache_range((addr_t)Random_String,
                                    CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX);
        RetVal = cmd_trng_get_rand(Random_String,
                                   CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX);
        arch_invalidate_cache_range((addr_t)Random_String,
                                    CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX);

        memcpy_(dst + index, (void *)Random_String,
                CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX);
        index = index + CRYPTO_SHARE_MEM_COMM_SMALL_BUFF_MAX;
    }

    if (cp_left > 0) {
        arch_invalidate_cache_range((addr_t)Random_String, cp_left);
        RetVal = cmd_trng_get_rand(Random_String, cp_left);
        arch_invalidate_cache_range((addr_t)Random_String, cp_left);
        memcpy_(dst + index, (void *)Random_String, cp_left);
    }

    for (; size > 0; size--) {
        if (dst[size - 1] == 0)
            dst[size - 1] = 0x1;
    }

    ShareMem_ReleaseBlock(Random_String);

    return RetVal;
}

/**
 * @brief Get first mask
 * @param n0 input
 * @return mask corresponding to \p n0
 */
static uint32_t getFirstMask(uint32_t n0)
{
    if (n0 & 0x80) {
        return 0x7F;
    } else if (n0 & 0x40) {
        return 0x3F;
    } else if (n0 & 0x20) {
        return 0x1F;
    } else if (n0 & 0x10) {
        return 0x0F;
    } else if (n0 & 0x08) {
        return 0x07;
    } else if (n0 & 0x04) {
        return 0x03;
    } else if (n0 & 0x02) {
        return 0x01;
    } else {
        return 0x00;
    }
}

/**
 * @brief perform buff = buff XOR mask
 * @param buff input/output buffer
 * @param mask to xor with
 * @param n size of \p buff and \p mask
*/
static void mask(uint8_t *buff, uint8_t *mask, uint32_t n)
{
    for (; n > 0; n--) {
        buff[n - 1] = buff[n - 1] ^ mask[n - 1];
    }
}

/**
 * Mask Generation Function as defined by RFC-8017 B.2.1.
 * @param hash_alg hash function to use
 * @param seed input
 * @param seedLen length of \p seed
 * @param mask output
 * @param maskLen length of \p mask
*/
sdrv_crypto_error_status_e MGF1(uint8_t hash_alg,
                                uint8_t *seed,
                                uint32_t seedLen,
                                uint8_t *mask,
                                uint32_t maskLen)
{
    /*
     * reduce size (get warning "stack protector not protecting function:
     * all local arrays are less than 8 bytes long")
     */
    uint8_t cnt[8];
    uint32_t c = 0;
    uint32_t hashLen;
    uint8_t *hash_in;
    uint32_t tmpLen = maskLen;
    uint8_t *tmpmask = mask;
    sdrv_crypto_error_status_e RetVal = E_OK;

    hash_in = ShareMem_GetBlock(68 + 28);

    hashLen = cmd_get_hash_digest_bytes(hash_alg);

    memcpy_(hash_in, seed, seedLen);

    /*
     * modifying cnt content via array_blk[1].
     * addr to make static analyzer happy
     */
    while (tmpLen) {
        cnt[0] = ((c >> 24)  & 0xFF);
        cnt[1] = ((c >> 16)  & 0xFF);
        cnt[2] = ((c >> 8)   & 0xFF);
        cnt[3] = ((c)        & 0xFF);
        memcpy_(hash_in + seedLen, cnt, 4);

        arch_clean_cache_range((addr_t)hash_in, seedLen + 4);
        RetVal = cmd_hash_calc(hash_alg, hash_in, seedLen + 4, tmpmask);

        if (CMD_RETURN_SUCCESS != RetVal)
            break;

        c++;
        tmpLen -= hashLen >= tmpLen ? tmpLen : hashLen;
        tmpmask += hashLen;
    }
    arch_invalidate_cache_range((addr_t)mask, maskLen);

    if (CMD_RETURN_SUCCESS != RetVal)
        RetVal = E_NOT_OK;
    else
        RetVal = E_OK;

    ShareMem_ReleaseBlock(hash_in);

    return RetVal;
}

/* message len <= key len -41 */
sdrv_crypto_error_status_e rsa_pad_eme_oaep_encode(uint32_t k,
                                                   uint32_t hash_alg,
                                                   uint8_t *EM,
                                                   uint8_t *message,
                                                   uint32_t mLen)
{
    uint8_t *dbMask;
    uint8_t *seedMask;
    uint8_t *tmp;
    uint32_t hLen;
    sdrv_crypto_error_status_e RetVal;

    hLen = cmd_get_hash_digest_bytes(hash_alg);

    if (mLen + 2 * hLen + 2 > k) {
        return E_NOT_OK;
    }

    tmp = ShareMem_GetBlock(hLen);

    arch_invalidate_cache_range((addr_t)tmp, hLen);
    /* get label hash -> no label so NULL_blk, get empty string hash */
    RetVal = cmd_hash_calc(hash_alg, NULL, 0, tmp);

    memcpy_((uint8_t *)(EM + hLen + 1), tmp, hLen);

    /*Assemply of DB*/
    gen_rnd_bytes(EM + 1, hLen);
    EM[0] = 0x00;
    memset_((uint8_t *)(EM + 2 * hLen + 1), 0x00, k - mLen - 2 * hLen - 2);
    *((uint8_t *)(EM + k - mLen - 1)) = 0x01;
    memcpy_(EM + k - mLen, message, mLen);

    dbMask = ShareMem_GetBlock(RSA_MAX_SIZE);

    seedMask = ShareMem_GetBlock(MAX_DIGESTSIZE);

    MGF1(hash_alg, EM + 1, hLen, dbMask, k - hLen - 1);
    mask(EM + 1 + hLen, dbMask, k - hLen - 1);
    MGF1(hash_alg, EM + 1 + hLen, k - hLen - 1, seedMask, hLen);
    mask(EM + 1, seedMask, hLen);

    ShareMem_ReleaseBlock(dbMask);
    ShareMem_ReleaseBlock(seedMask);
    ShareMem_ReleaseBlock(tmp);

    if (CMD_RETURN_SUCCESS != RetVal)
        RetVal = E_NOT_OK;
    else
        RetVal = E_OK;

    return RetVal;
}

sdrv_crypto_error_status_e rsa_pad_eme_oaep_decode(uint32_t k,
                                                   uint32_t hash_alg,
                                                   uint8_t *EM,
                                                   uint8_t **message,
                                                   uint32_t *mLen)
{
    uint32_t hLen;
    uint8_t *dbMask;
    uint8_t *seedMask;
    uint32_t i;
    int chkLHash;
    sdrv_crypto_error_status_e RetVal;

    dbMask = ShareMem_GetBlock(RSA_MAX_SIZE);

    seedMask = ShareMem_GetBlock(MAX_DIGESTSIZE);

    hLen = cmd_get_hash_digest_bytes(hash_alg);

    MGF1(hash_alg, EM + hLen + 1, k - hLen - 1, seedMask, hLen);
    mask(EM + 1, seedMask, hLen);
    MGF1(hash_alg, EM + 1, hLen, dbMask, k - hLen - 1);
    mask(EM + hLen + 1, dbMask, k - hLen - 1);

    *mLen = 0;

    for (i = 2 * hLen + 1; i < k; i++) {
        if (*(uint8_t *)(EM + i) == 0x01) {
            *mLen = k - i - 1;
            break;
        }
    }
    arch_invalidate_cache_range((addr_t)(seedMask), hLen);
    /* get label hash -> no label so NULL_blk, get empty string hash*/
    RetVal = cmd_hash_calc(hash_alg, NULL, 0, seedMask);

    chkLHash = memcmp_(seedMask, EM + hLen + 1,
                       cmd_get_hash_digest_bytes(hash_alg));

    ShareMem_ReleaseBlock(dbMask);
    ShareMem_ReleaseBlock(seedMask);

    if (chkLHash || *mLen == 0 || EM[0]) {
        return E_NOT_OK;
    }

    *message = (uint8_t *)((EM + k) - *mLen);

    if (CMD_RETURN_SUCCESS != RetVal)
        RetVal = E_NOT_OK;
    else
        RetVal = E_OK;

    return RetVal;
}

/* message len <= key len -11 */
sdrv_crypto_error_status_e rsa_pad_eme_pkcs_encode(uint32_t k,
                                                   uint8_t *EM,
                                                   uint8_t *message,
                                                   uint32_t mLen)
{
    if (mLen > (uint32_t)(k - 11))
        return E_NOT_OK;

    /* Assemply of DB, PS (first written for alignment purpose)*/
    gen_rnd_bytes(EM + 2, k - mLen - 3);
    for (int i = 0; i < (k - mLen - 3); i++) {
        if (EM[2 + i] == 0)
            EM[2 + i]++;
    }

    EM[0] = 0x00;
    EM[1] = 0x02;

    *((uint8_t *)(EM + k - mLen - 1)) = 0x00;
    memcpy_(EM + k - mLen, message, mLen);

    return E_OK;
}

sdrv_crypto_error_status_e rsa_pad_eme_pkcs_decode(uint32_t k,
                                 uint8_t *EM,
                                 uint8_t **message,
                                 uint32_t *mLen)
{
    uint32_t PSLen;
    uint32_t i;

    *mLen = 0;

    for (i = 2; i < k; i++) {
        if (*(uint8_t *)(EM + i) == 0x00) {
            *mLen = k - i - 1;
            break;
        }
    }

    PSLen = k - 3 - *mLen;

    if (*mLen == 0 || EM[0] || EM[1] != 0x02 || PSLen < 8) {
        return E_NOT_OK;
    }

    *message = (uint8_t *)((EM + k) - *mLen);

    return E_OK;
}

sdrv_crypto_error_status_e rsa_pad_emsa_pkcs_encode(uint32_t emLen,
                                                    uint32_t hash_type,
                                                    uint8_t *EM,
                                                    uint8_t *hash)
{
    uint8_t hash_algDer[19] = {
        0x30, 0x21, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86,
        0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x06, 0x05,
        0x00, 0x04, 0x20
    };
    uint32_t tLen;
    uint32_t hLen;
    uint32_t dLen = 19;

    /*adapt DER encoded hash_alg*/
    switch (hash_type) {
    case HASH_ALG_SHA224:
        hLen = 28;
        hash_algDer[1] = 0x2d;
        hash_algDer[14] = 0x04;
        hash_algDer[18] = 0x1c;
        break;

    case HASH_ALG_SHA256:
        hLen = 32;
        hash_algDer[1] = 0x31;
        hash_algDer[14] = 0x01;
        hash_algDer[18] = 0x20;
        break;

    case HASH_ALG_SHA384:
        hLen = 4 * 12;
        hash_algDer[1] = 0x41;
        hash_algDer[14] = 0x02;
        hash_algDer[18] = 0x30;
        break;

    case HASH_ALG_SHA512:
        hLen = 4 * 16;
        hash_algDer[1] = 0x51;
        hash_algDer[14] = 0x03;
        hash_algDer[18] = 0x40;
        break;

    default:
        return E_NOT_OK;
    }

    tLen = hLen + dLen;

    if (emLen < tLen + 11) {
        return E_NOT_OK;
    }

    *EM = 0x00;
    *(uint8_t *)(EM + 1) = 0x01;
    memset_((uint8_t *)(EM + 2), 0xff, emLen - tLen - 3);
    *(uint8_t *)(EM + emLen - tLen - 1) = 0x00;
    memcpy_(EM + emLen - tLen, hash_algDer, dLen);
    memcpy_(EM + emLen - hLen, hash, hLen);

    return E_OK;
}

sdrv_crypto_error_status_e rsa_pad_emsa_pss_encode(uint32_t emLen,
                                                   uint32_t hash_alg,
                                                   uint8_t *EM,
                                                   uint8_t *hash,
                                                   uint32_t n0,
                                                   uint32_t sLen)
{
    sdrv_crypto_error_status_e RetVal = E_OK;
    uint32_t hLen;
    uint8_t *dbMask;
    uint8_t *tempm;
    uint8_t *temph;

    hLen = cmd_get_hash_digest_bytes(hash_alg);
    if (!hLen) {
        return E_NOT_OK;
    }

    if (emLen < sLen + hLen + 2) {
        return E_NOT_OK;
    }

    tempm = ShareMem_GetBlock(128 + 8);
    memset_(tempm, 0, 8);
    memcpy_(tempm + 8, hash, hLen);
    gen_rnd_bytes(tempm + 8 + hLen, sLen);

    temph = ShareMem_GetBlock(hLen);
    arch_clean_cache_range((addr_t)tempm, 8 + hLen + sLen);
    arch_invalidate_cache_range((addr_t)temph, hLen);
    RetVal = cmd_hash_calc(hash_alg, tempm, (8 + hLen + sLen), temph);
    arch_invalidate_cache_range((addr_t)temph, hLen);

    memset_(EM, 0x00, emLen - sLen - hLen - 2);
    *((uint8_t *)(EM + emLen - sLen - hLen - 2)) = 0x01;
    memcpy_(EM + emLen - sLen - hLen - 1, tempm + 8 + hLen, sLen);

    dbMask = ShareMem_GetBlock(RSA_MAX_SIZE);
    MGF1(hash_alg, temph, hLen, dbMask, emLen - hLen - 1);
    mask(EM, dbMask, emLen - hLen - 1);

    memcpy_((uint8_t *)(EM + emLen - hLen - 1), temph, hLen);
    EM[emLen - 1] = 0xBC;
    EM[0] = EM[0] & getFirstMask(n0);

    if (CMD_RETURN_SUCCESS != RetVal)
        RetVal = E_NOT_OK;
    else
        RetVal = E_OK;

    ShareMem_ReleaseBlock(dbMask);
    ShareMem_ReleaseBlock(tempm);
    ShareMem_ReleaseBlock(temph);

    return RetVal;
}

/* Steps from rfc8017 9.1.2 Verification operation */
sdrv_crypto_error_status_e rsa_pad_emsa_pss_decode(uint32_t emLen,
                                                   uint32_t hash_alg,
                                                   uint8_t *EM,
                                                   uint8_t *hash,
                                                   uint32_t sLen,
                                                   uint32_t n0)
{
    uint8_t *maskedDB;
    uint32_t dbLen;
    uint8_t *H;
    uint32_t mlen;
    uint8_t *salt;
    uint8_t *H_;
    int chkLHash;
    uint8_t *dbMask;
    uint8_t *hash_in;
    uint32_t hLen;
    uint8_t *DB;
    sdrv_crypto_error_status_e RetVal;

    hLen = cmd_get_hash_digest_bytes(hash_alg);

    /* 3.  If emLen < hLen + sLen + 2, output "inconsistent" and stop. */
    /* 4.  If the rightmost octet of EM does not have hexadecimal value
        0xbc, output "inconsistent" and stop. */
    if (emLen < hLen + sLen + 2 || EM[emLen - 1] != 0xbc) {
        return E_NOT_OK;
    }

    /* 5.  Let maskedDB be the leftmost emLen - hLen - 1 octets of EM, and
        let H be the next hLen octets. */
    maskedDB = EM;
    dbLen = emLen - hLen - 1;
    H = EM + emLen - hLen - 1;

    /*FIXME: This check is not performed because it fails a few known-good test cases.*/
    /* 6.   If the leftmost 8emLen - emBits bits of the leftmost octet in
            maskedDB are not all equal to zero, output "inconsistent" and
            stop. */

    /* 7.  Let dbMask = MGF(H, emLen - hLen - 1). */
    /*This buffer is later reused to store H'*/

    mlen = emLen - hLen - 1;
    mlen = mlen + ((mlen % hLen) ? hLen - (mlen % hLen) : 0);

    if (mlen > RSA_MAX_SIZE) {
        return E_NOT_OK;
    }

    dbMask = ShareMem_GetBlock(RSA_MAX_SIZE);

    MGF1(hash_alg, H, hLen, dbMask, mlen);

    /* 8.  Let DB = maskedDB \xor dbMask. */
    mask(maskedDB, dbMask, emLen - hLen - 1);
    /*from here maskedDB = RFC's DB*/
    DB = maskedDB;
    /* 9.  Set the leftmost 8emLen - emBits bits of the leftmost octet in DB
           to zero. */
    /* emLen = (emBits + 7) / 8 */
    DB[0] = DB[0] & getFirstMask(n0);

    /* 10. If the emLen - hLen - sLen - 2 leftmost octets of DB are not zero
        or if the octet at position emLen - hLen - sLen - 1 (the leftmost
        position is "position 1") does not have hexadecimal value 0x01,
        output "inconsistent" and stop. */
    for (uint32_t i = 1; i < (emLen - hLen - sLen - 2); i++) {
        if (*(uint8_t *)(DB + i)) {
            RetVal = E_NOT_OK;
            goto end;
        }
    }

    if (*(uint8_t *)(DB + emLen - hLen - sLen - 2) != 0x01) {
        RetVal = E_NOT_OK;
        goto end;
    }

    /* 11.  Let salt be the last sLen octets of DB. */
    salt = DB + dbLen - sLen;

    /* 12.  Let M' = (0x)00 00 00 00 00 00 00 00 || mHash || salt ;
        M' is an octet string of length 8 + hLen + sLen with eight
        initial zero octets. */
    hash_in = ShareMem_GetBlock(128);

    memset_(hash_in, 0, 8);
    memcpy_(hash_in + 8, hash, hLen);
    memcpy_(hash_in + 8 + hLen, salt, sLen);

    /* 13. Let H' = Hash(M'), an octet string of length hLen. */
    /*Hash is placed in the unused big buffer. First byte of EM 1 byte to small.*/
    H_ = dbMask;

    arch_clean_cache_range((addr_t)hash_in, 8 + hLen + sLen);
    arch_invalidate_cache_range((addr_t)H_, hLen);
    RetVal = cmd_hash_calc(hash_alg, hash_in, (8 + hLen + sLen), H_);
    if (CMD_RETURN_SUCCESS != RetVal) {
        RetVal = E_NOT_OK;
        goto end;
    }

    arch_invalidate_cache_range((addr_t)H_, hLen);

    /* 14. If H = H', output "consistent." Otherwise, output "inconsistent." */
    chkLHash = memcmp_(H, H_, hLen);

    if (chkLHash)
        RetVal = E_NOT_OK;
    else
        RetVal = E_OK;

end:
    ShareMem_ReleaseBlock(dbMask);
    ShareMem_ReleaseBlock(hash_in);

    return RetVal;
}

void rsa_pad_zeros(uint8_t *EM, uint32_t emLen, uint8_t *hash, uint32_t hashLen)
{
    memset_(EM, 0x00, emLen - hashLen);

    if (hashLen) {
        memcpy_(EM + emLen - hashLen, hash, hashLen);
    }
}