#include "hcu_cases.h"
#include "trng_driver.h"

#define RSA_1024_WORDS                (32U)
#define RSA_PADDING_MESSAGE_BYTES     (32U)
#define RSA_PADDING_RANDOM_BYTES      (32U)
#define RSAES_PKCS1_V15_PS_BYTES      (128U - RSA_PADDING_MESSAGE_BYTES - 3U)

/* Set to 1 to generate the OAEP seed, PKCS1-v1_5 PS, and PSS salt by TRNG. */
#ifndef SEED_GEN_BY_TRNG
#define SEED_GEN_BY_TRNG              (0U)
#endif

static const hcu_rsa_config_t rsa_1024 = {
    .publicKeyLen = 1,
    .e = {0x10001},
    .opLength = RSA_1024,
    .n = {
        0x5A6B9AA1, 0xA0CCF421, 0x087D0869, 0x0F2442BB,
        0x5ADCBE62, 0xD6BB2B74, 0xE95758F9, 0x6C64026D,
        0x285C625A, 0xEFD6C72E, 0x95902D9E, 0x1598F870,
        0xF8313369, 0xCBEB071B, 0x91F37775, 0xE941DCB4,
        0x0F4B8CBD, 0x72436C7F, 0x4194AB72, 0xEB6B6FA4,
        0x98FE812A, 0x0011E5EC, 0x1BD3FD19, 0x5FA2FD8A,
        0xD210AC87, 0x59D9E3AA, 0x676BC6A1, 0x1090AEEC,
        0x1FEA4EF1, 0x5E654EEC, 0x20FADD5D, 0x8CE2FA9F,
    },
    .privateKeyLen = 32,
    .d = {
        0xF8E9F421, 0x720E6468, 0x5240470E, 0x8BC2BAB7,
        0xFB62FA07, 0xBC9C2607, 0x88104B99, 0xF02B35DA,
        0xACAE6C90, 0x1EB03B65, 0x29AF5130, 0xF0776789,
        0xD071F6F4, 0x8E147324, 0x6BD265AF, 0xC934FA70,
        0x0ECD75F3, 0xEB735517, 0x0D4BF2CD, 0x7698DB5A,
        0x62AE3A15, 0x124C9AB9, 0x3AFCD88C, 0xBF4B98E8,
        0x644085EA, 0xB1BC37A3, 0xB53B91B8, 0xCAC436C2,
        0xEE2D939B, 0x264AB31F, 0x04343D40, 0x15102234,
    },
    /* RSA padding APIs must use polling mode. */
    .rsaCarryType = HCU_RSA_USING_POLLING,
};

static const uint32_t s_message[RSA_PADDING_MESSAGE_BYTES / sizeof(uint32_t)] = {
    0x6BC1BEE2, 0x2E409F96, 0xE93D7E11, 0x7393172A,
    0xAE2D8A57, 0x1E03AC9C, 0x9EB76FAC, 0x45AF8E51,
};

#if !SEED_GEN_BY_TRNG
/* Bytes 00..1F are used as both the fixed OAEP seed and PSS salt. */
static const uint32_t s_fixedRandom[RSA_PADDING_RANDOM_BYTES / sizeof(uint32_t)] = {
    0x03020100, 0x07060504, 0x0B0A0908, 0x0F0E0D0C,
    0x13121110, 0x17161514, 0x1B1A1918, 0x1F1E1D1C,
};

static const uint32_t s_expectedOaepCipher[RSA_1024_WORDS] = {
    0xED06B20F, 0xAD425B4D, 0xB67AFEBA, 0xB0A0264A,
    0xF7206668, 0xDA67A816, 0xB29BEC12, 0x7893E33F,
    0xBBCA3BB2, 0x614BB1CB, 0x2103F3FB, 0x4E6B7CC0,
    0x326CB0C2, 0xA3B7C6FA, 0x342CC688, 0x9277ECB4,
    0xFDA4944E, 0x359994B5, 0x03F10E25, 0x93FE9E4A,
    0x5840E030, 0x0C748311, 0x5D7F29A4, 0x0630979A,
    0xA4637744, 0x5A42B3F2, 0xAB2BE3F5, 0x877E097A,
    0xEFABE6EA, 0x4577FB0D, 0x1065DE77, 0x2B13691D,
};

static const uint32_t s_expectedRsaesPkcs1Cipher[RSA_1024_WORDS] = {
    0x9B462077, 0x2FAFAC0F, 0xB7053193, 0xAEA7CB06,
    0xAF876816, 0x82214EA1, 0xC172FD4F, 0xD3F765C4,
    0x56685D9F, 0x01D2C536, 0xBB05A454, 0x71F60BDF,
    0x2F0D40A7, 0x0DF63A33, 0xBE24AF75, 0xACC119E4,
    0xCB0287F7, 0x0EF8B07E, 0x89A9A2F7, 0x92FF2974,
    0xB5F109B0, 0xD9028F1F, 0x77A59E07, 0x9061FAB8,
    0x08190B30, 0xF667820F, 0x974AF9A0, 0x94E621A4,
    0x0DBB0A3D, 0xA4662AE8, 0xAB22B592, 0x3ED71B88,
};

static const uint32_t s_expectedRsassaPkcs1Signature[RSA_1024_WORDS] = {
    0xC3EFF912, 0xEC6F1D22, 0xCA29D850, 0x914F4983,
    0xC622A39E, 0x018059BA, 0xA72F192A, 0xFA1D728B,
    0x3747F4E3, 0xF9ED907A, 0x14FF913F, 0x5BAE9900,
    0x8CE88671, 0x047FD29B, 0x6C213CC0, 0x584C5C81,
    0x53D2C51F, 0xF134DBCE, 0x84695479, 0xE730DD9F,
    0xC6039541, 0xE67C3CCB, 0x126E63FB, 0x034AB0A9,
    0x9D94EA8F, 0x83994AAF, 0x2757686E, 0x81C01134,
    0x6238301B, 0xEF703DA5, 0x7D483AE7, 0x50C81BF8,
};

static const uint32_t s_expectedPssSignature[RSA_1024_WORDS] = {
    0xC32042C0, 0x5E7FCB20, 0x3B017440, 0x73E89DE8,
    0x789B8552, 0xA43F5787, 0xD6668593, 0xD714684D,
    0x1F9FB206, 0x8C542A8A, 0x5786C57D, 0xD6A33206,
    0x72F27C72, 0x71AEE862, 0x8936D235, 0xE1C9034F,
    0x3CD949F5, 0xA82396CE, 0xB363AB92, 0xEE87602A,
    0x20B0B10D, 0xC6888BF2, 0xE7BC4E9D, 0xA72C8AAC,
    0xF315EC37, 0x67176D99, 0x32426BA1, 0xC16F91BB,
    0x1BEF2EC0, 0x83A61E65, 0xCC4A072A, 0x5D3D584B,
};
#endif /* !SEED_GEN_BY_TRNG */

static uint8_t s_paddingRandom[RSAES_PKCS1_V15_PS_BYTES];
static uint32_t s_cipher[RSA_1024_WORDS];
static uint32_t s_recovered[RSA_PADDING_MESSAGE_BYTES / sizeof(uint32_t)];
static uint32_t s_signature[RSA_1024_WORDS];
static pke_user_config_t s_pkeUserConfig;
static pke_state_t s_pkeState;

static bool HcuRsaPaddingWordsMatch(const uint32_t *actual,
                                    const uint32_t *expected,
                                    uint32_t wordCount)
{
    for (uint32_t i = 0U; i < wordCount; i++)
    {
        if (actual[i] != expected[i])
        {
            return false;
        }
    }
    return true;
}

/* Encryption passes only when the complete plaintext is recovered. */
static bool HcuRsaPaddingMessageMatches(uint32_t recoveredLength)
{
    return ((recoveredLength == RSA_PADDING_MESSAGE_BYTES) &&
            HcuRsaPaddingWordsMatch(s_recovered, s_message,
                                    RSA_PADDING_MESSAGE_BYTES / sizeof(uint32_t)));
}

#if SEED_GEN_BY_TRNG
/* Fill a seed/salt, or a nonzero PKCS1-v1_5 PS, from 128-bit TRNG blocks. */
static void HcuRsaPaddingGetRandom(uint8_t *output, uint32_t length,
                                   bool rejectZero)
{
    uint32_t entropy[4];
    uint8_t *bytes = (uint8_t *)entropy;
    uint32_t offset = 0U;

    while (offset < length)
    {
        TRNG_DRV_Get_Ent(0, entropy);
        while (STATUS_SUCCESS != TRNG_DRV_GetStatus(0));
        for (uint32_t i = 0U; (i < sizeof(entropy)) && (offset < length); i++)
        {
            if ((!rejectZero) || (bytes[i] != 0U))
            {
                output[offset++] = bytes[i];
            }
        }
    }
}
#endif /* SEED_GEN_BY_TRNG */

void hcu_rsa_padding_demo(void)
{
    const hcu_rsa_oaep_config_t oaepConfig = {
        .label = NULL,
        .labelLength = 0U,
        .shaType = HCU_SHA_256,
    };
    const hcu_rsa_pss_config_t pssConfig = {
        .saltLength = RSA_PADDING_RANDOM_BYTES,
        .shaType = HCU_SHA_256,
    };
    const uint8_t *seed;
    const uint8_t *salt;
    uint32_t recoveredLength;
    status_t status;
    bool passed;

#if SEED_GEN_BY_TRNG
    PRINTF("Show four RSA-1024 padding modes using TRNG.\r\n");
    HcuRsaPaddingGetRandom(s_paddingRandom, RSA_PADDING_RANDOM_BYTES, false);
    seed = s_paddingRandom;
#else
    PRINTF("Show four RSA-1024 padding modes with fixed vectors.\r\n");
    seed = (const uint8_t *)s_fixedRandom;
#endif

    s_pkeUserConfig.rsaConfig = &rsa_1024;
    status = PKE_DRV_Init(&s_pkeUserConfig, &s_pkeState);
    if (status != STATUS_SUCCESS)
    {
        PRINTF("RSA padding PKE init failed, status=%x.\r\n", status);
        return;
    }

    status = HCU_DRV_RSAES_OAEP_Encrypt(
        &oaepConfig, (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES,
        seed, (uint8_t *)s_cipher);
    passed = (status == STATUS_SUCCESS);
#if !SEED_GEN_BY_TRNG
    passed = (passed &&
              HcuRsaPaddingWordsMatch(s_cipher, s_expectedOaepCipher,
                                      RSA_1024_WORDS));
#endif
    recoveredLength = sizeof(s_recovered);
    if (passed)
    {
        status = HCU_DRV_RSAES_OAEP_Decrypt(
            &oaepConfig, (const uint8_t *)s_cipher, (uint8_t *)s_recovered,
            &recoveredLength);
        passed = ((status == STATUS_SUCCESS) &&
                  HcuRsaPaddingMessageMatches(recoveredLength));
    }
    PRINTF("RSAES-OAEP-SHA256 %s.\r\n", passed ? "success" : "failed");

#if SEED_GEN_BY_TRNG
    HcuRsaPaddingGetRandom(s_paddingRandom, sizeof(s_paddingRandom), true);
#else
    for (uint32_t i = 0U; i < sizeof(s_paddingRandom); i++)
    {
        s_paddingRandom[i] = (uint8_t)(i + 1U);
    }
#endif
    status = HCU_DRV_RSAES_PKCS1_V15_Encrypt(
        (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES, s_paddingRandom,
        (uint8_t *)s_cipher);
    passed = (status == STATUS_SUCCESS);
#if !SEED_GEN_BY_TRNG
    passed = (passed &&
              HcuRsaPaddingWordsMatch(s_cipher, s_expectedRsaesPkcs1Cipher,
                                      RSA_1024_WORDS));
#endif
    recoveredLength = sizeof(s_recovered);
    if (passed)
    {
        status = HCU_DRV_RSAES_PKCS1_V15_Decrypt(
            (const uint8_t *)s_cipher, (uint8_t *)s_recovered,
            &recoveredLength);
        passed = ((status == STATUS_SUCCESS) &&
                  HcuRsaPaddingMessageMatches(recoveredLength));
    }
    PRINTF("RSAES-PKCS1-v1_5 %s.\r\n", passed ? "success" : "failed");

    status = HCU_DRV_RSASSA_PKCS1_V15_GenerateSignature(
        (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES, HCU_SHA_256,
        (uint8_t *)s_signature);
    passed = (status == STATUS_SUCCESS);
#if !SEED_GEN_BY_TRNG
    passed = (passed &&
              HcuRsaPaddingWordsMatch(s_signature,
                                      s_expectedRsassaPkcs1Signature,
                                      RSA_1024_WORDS));
#endif
    if (passed)
    {
        status = HCU_DRV_RSASSA_PKCS1_V15_VerifySignature(
            (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES, HCU_SHA_256,
            (const uint8_t *)s_signature);
        passed = (status == STATUS_SUCCESS);
    }
    PRINTF("RSASSA-PKCS1-v1_5-SHA256 %s.\r\n",
           passed ? "success" : "failed");

#if SEED_GEN_BY_TRNG
    HcuRsaPaddingGetRandom(s_paddingRandom, RSA_PADDING_RANDOM_BYTES, false);
    salt = s_paddingRandom;
#else
    salt = (const uint8_t *)s_fixedRandom;
#endif
    status = HCU_DRV_RSASSA_PSS_GenerateSignature(
        &pssConfig, (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES,
        salt, (uint8_t *)s_signature);
    passed = (status == STATUS_SUCCESS);
#if !SEED_GEN_BY_TRNG
    passed = (passed &&
              HcuRsaPaddingWordsMatch(s_signature, s_expectedPssSignature,
                                      RSA_1024_WORDS));
#endif
    if (passed)
    {
        status = HCU_DRV_RSASSA_PSS_VerifySignature(
            &pssConfig, (const uint8_t *)s_message, RSA_PADDING_MESSAGE_BYTES,
            (const uint8_t *)s_signature);
        passed = (status == STATUS_SUCCESS);
    }
    PRINTF("RSASSA-PSS-SHA256 %s.\r\n", passed ? "success" : "failed");

    status = PKE_DRV_DeInit(&s_pkeState);
    if (status != STATUS_SUCCESS)
    {
        PRINTF("RSA padding PKE deinit failed, status=%x.\r\n", status);
    }
}
