/*
 * This file is part of the openHiTLS project.
 *
 * openHiTLS is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>

#include "crypt_modes_cfb.h"   

// 十六进制打印
static void dump_hex(const char *title, const uint8_t *buf, size_t len) {
    printf("%s (len=%zu):\n", title, len);
    for (size_t i = 0; i < len; ++i) {
        printf("%02X", buf[i]);
        if ((i + 1) % 16 == 0 || i + 1 == len) printf("\n");
        else printf(" ");
    }
}

// 使用给定切分进行流式CFB处理；enc=true加密/false解密
static int stream_cfb(bool enc,
                      const uint8_t *key, uint32_t keyLen,
                      const uint8_t *iv,  uint32_t ivLen,
                      const uint8_t *in,  uint32_t inLen,
                      const uint32_t *chunks, size_t nChunks,
                      uint8_t *out /* >= inLen */)
{
    MODES_CFB_Ctx m;
    int32_t rc = SM4_CFB_InitCtx(&m, key, keyLen, iv, ivLen, enc);
    if (rc != 0) {
        fprintf(stderr, "SM4_CFB_InitCtx failed: %d\n", rc);
        return rc;
    }

    uint32_t produced = 0, consumed = 0;
    for (size_t i = 0; i < nChunks && consumed < inLen; ++i) {
        uint32_t take = chunks[i];
        if (consumed + take > inLen) take = inLen - consumed;

        uint32_t outLen = 0;
        rc = SM4_CFB_Update(&m, in + consumed, take, out + produced, &outLen);
        if (rc != 0) {
            fprintf(stderr, "SM4_CFB_Update failed at chunk #%zu: %d\n", i, rc);
            return rc;
        }
        if (outLen != take) {
            fprintf(stderr, "Unexpected outLen %u (expect %u)\n", outLen, take);
            return -1;
        }
        consumed += take;
        produced += outLen;
    }
    // 若切分未覆盖全部数据，收尾一次
    if (consumed < inLen) {
        uint32_t tail = inLen - consumed, outLen = 0;
        rc = SM4_CFB_Update(&m, in + consumed, tail, out + produced, &outLen);
        if (rc != 0 || outLen != tail) {
            fprintf(stderr, "SM4_CFB_Update tail failed: rc=%d, outLen=%u\n", rc, outLen);
            return (rc != 0) ? rc : -1;
        }
        produced += outLen;
    }
    return (produced == inLen) ? 0 : -1;
}

int main(void) {
    // 16字节密钥：01..10
    uint8_t key[16]; for (int i = 0; i < 16; ++i) key[i] = (uint8_t)(i + 1);
    // 16字节IV：00..0F
    uint8_t iv[16];  for (int i = 0; i < 16; ++i) iv[i]  = (uint8_t)i;

    // 非16倍数长度明文，便于测试流式
    const uint8_t plaintext[] =
        "The quick brown fox jumps over the lazy dog. CFB mode streaming test!";
    const uint32_t pt_len = (uint32_t)strlen((const char*)plaintext);

    uint8_t ct_stream[1024] = {0};
    uint8_t ct_one[1024]    = {0};
    uint8_t pt_back[1024]   = {0};

    dump_hex("PLAINTEXT", plaintext, pt_len);

    // 1) 流式加密（不规则分段）
    const uint32_t enc_chunks[] = {5, 7, 13, 1, 3};
    if (stream_cfb(true, key, sizeof(key), iv, sizeof(iv),
                   plaintext, pt_len, enc_chunks, sizeof(enc_chunks)/sizeof(enc_chunks[0]),
                   ct_stream) != 0) {
        return 2;
    }
    dump_hex("CIPHERTEXT (streamed)", ct_stream, pt_len);

    // 2) 一次性整段加密，作为校验基准（仍只用 Init+Update）
    {
        MODES_CFB_Ctx m;
        uint8_t iv2[16]; memcpy(iv2, iv, 16);
        if (SM4_CFB_InitCtx(&m, key, sizeof(key), iv2, sizeof(iv2), true) != 0) return 3;
        uint32_t outLen = 0;
        if (SM4_CFB_Update(&m, plaintext, pt_len, ct_one, &outLen) != 0 || outLen != pt_len) return 3;
    }
    dump_hex("CIPHERTEXT (one-shot)", ct_one, pt_len);

    // 3) 校验：流式 == 一次性
    if (memcmp(ct_stream, ct_one, pt_len) != 0) {
        fprintf(stderr, "[X] Encrypt mismatch: streamed vs one-shot.\n");
        return 4;
    }
    printf("[OK] Encrypt check passed.\n");

    // 4) 流式解密（用不同切分）
    const uint32_t dec_chunks[] = {8, 2, 19, 4, 1};
    if (stream_cfb(false, key, sizeof(key), iv, sizeof(iv),
                   ct_stream, pt_len, dec_chunks, sizeof(dec_chunks)/sizeof(dec_chunks[0]),
                   pt_back) != 0) {
        return 5;
    }
    dump_hex("DECRYPTED", pt_back, pt_len);

    // 5) 校验：解密结果 == 原文
    if (memcmp(pt_back, plaintext, pt_len) != 0) {
        fprintf(stderr, "[X] Decrypt mismatch: decrypted != plaintext.\n");
        return 6;
    }
    printf("[OK] Decrypt check passed.\n");

    puts("\nAll SM4-CFB streaming tests passed.");
    return 0;
}