AES加解密接口

#include <windows.h>
#include <bcrypt.h>
#include <iostream>
#include <vector>
#include <string>
#include <iomanip>
#include <sstream>

#ifndef NT_SUCCESS
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#endif

#pragma comment(lib, "bcrypt.lib")

std::string BytesToHex(const std::vector<BYTE>& data) {
    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (BYTE b : data) {
        oss << std::setw(2) << static_cast<int>(b);
    }
    return oss.str();
}

// AES-CBC 加密接口
NTSTATUS AesCbcEncrypt(
    const std::vector<BYTE>& key,
    const std::vector<BYTE>& iv,
    const std::vector<BYTE>& plaintext,
    std::vector<BYTE>& ciphertext)
{
    // === 所有变量在函数顶部声明 ===
    NTSTATUS status = 0;
    BCRYPT_ALG_HANDLE hAlg = NULL;
    BCRYPT_KEY_HANDLE hKey = NULL;
    const wchar_t* chainingMode = BCRYPT_CHAIN_MODE_CBC;
    ULONG cipherBufferSize = 0;
    ULONG resultSize = 0;
    std::vector<BYTE> iv_copy; // 用于防止 IV 被 CNG 修改

    // 验证输入
    if (iv.size() != 16) {
        return STATUS_INVALID_PARAMETER;
    }

    // 创建 IV 副本（关键！）
    iv_copy = iv;

    status = BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_AES_ALGORITHM, NULL, 0);
    if (!NT_SUCCESS(status)) goto cleanup;

    status = BCryptSetProperty(
        hAlg,
        BCRYPT_CHAINING_MODE,
        reinterpret_cast<PUCHAR>(const_cast<wchar_t*>(chainingMode)),
        static_cast<ULONG>((wcslen(chainingMode) + 1) * sizeof(wchar_t)),
        0
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    status = BCryptGenerateSymmetricKey(
        hAlg,
        &hKey,
        NULL, 0,
        const_cast<BYTE*>(key.data()),
        static_cast<ULONG>(key.size()),
        0
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    // 第一次调用：查询所需输出缓冲区大小
    status = BCryptEncrypt(
        hKey,
        const_cast<BYTE*>(plaintext.data()),
        static_cast<ULONG>(plaintext.size()),
        NULL,
        iv_copy.data(),
        static_cast<ULONG>(iv_copy.size()),
        NULL, 0,
        &cipherBufferSize,
        BCRYPT_BLOCK_PADDING
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    // 分配并执行加密
    ciphertext.resize(cipherBufferSize);
    status = BCryptEncrypt(
        hKey,
        const_cast<BYTE*>(plaintext.data()),
        static_cast<ULONG>(plaintext.size()),
        NULL,
        iv_copy.data(),
        static_cast<ULONG>(iv_copy.size()),
        ciphertext.data(),
        static_cast<ULONG>(ciphertext.size()),
        &resultSize,
        BCRYPT_BLOCK_PADDING
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    ciphertext.resize(resultSize);

cleanup:
    if (hKey) BCryptDestroyKey(hKey);
    if (hAlg) BCryptCloseAlgorithmProvider(hAlg, 0);
    return status;
}

// AES-CBC 解密接口
NTSTATUS AesCbcDecrypt(
    const std::vector<BYTE>& key,
    const std::vector<BYTE>& iv,
    const std::vector<BYTE>& ciphertext,
    std::vector<BYTE>& plaintext)
{
    // === 所有变量在函数顶部声明 ===
    NTSTATUS status = 0;
    BCRYPT_ALG_HANDLE hAlg = NULL;
    BCRYPT_KEY_HANDLE hKey = NULL;
    const wchar_t* chainingMode = BCRYPT_CHAIN_MODE_CBC;
    ULONG plainBufferSize = 0;
    ULONG resultSize = 0;
    std::vector<BYTE> iv_copy; // 用于防止 IV 被 CNG 修改

    // 验证输入
    if (iv.size() != 16) {
        return STATUS_INVALID_PARAMETER;
    }

    // 创建 IV 副本（关键！）
    iv_copy = iv;

    status = BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_AES_ALGORITHM, NULL, 0);
    if (!NT_SUCCESS(status)) goto cleanup;

    status = BCryptSetProperty(
        hAlg,
        BCRYPT_CHAINING_MODE,
        reinterpret_cast<PUCHAR>(const_cast<wchar_t*>(chainingMode)),
        static_cast<ULONG>((wcslen(chainingMode) + 1) * sizeof(wchar_t)),
        0
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    status = BCryptGenerateSymmetricKey(
        hAlg,
        &hKey,
        NULL, 0,
        const_cast<BYTE*>(key.data()),
        static_cast<ULONG>(key.size()),
        0
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    // 第一次调用：查询所需输出缓冲区大小
    status = BCryptDecrypt(
        hKey,
        const_cast<BYTE*>(ciphertext.data()),
        static_cast<ULONG>(ciphertext.size()),
        NULL,
        iv_copy.data(),
        static_cast<ULONG>(iv_copy.size()),
        NULL, 0,
        &plainBufferSize,
        BCRYPT_BLOCK_PADDING
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    // 分配并执行解密
    plaintext.resize(plainBufferSize);
    status = BCryptDecrypt(
        hKey,
        const_cast<BYTE*>(ciphertext.data()),
        static_cast<ULONG>(ciphertext.size()),
        NULL,
        iv_copy.data(),
        static_cast<ULONG>(iv_copy.size()),
        plaintext.data(),
        static_cast<ULONG>(plaintext.size()),
        &resultSize,
        BCRYPT_BLOCK_PADDING
    );
    if (!NT_SUCCESS(status)) goto cleanup;

    plaintext.resize(resultSize);

cleanup:
    if (hKey) BCryptDestroyKey(hKey);
    if (hAlg) BCryptCloseAlgorithmProvider(hAlg, 0);
    return status;
}

// ==================== 主函数：演示调用 ====================
int main() {
    std::string plaintext_str = "Top secret message for AES encryption!1111144444";
    std::vector<BYTE> plaintext(plaintext_str.begin(), plaintext_str.end());

    // AES-256 密钥（32 字节）
    std::vector<BYTE> key = {
        0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,
        0x08,0x09,0x0A,0x0B,0x0C,0x0D,0x0E,0x0F,
        0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,
        0x18,0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F
    };

    // 初始向量 IV（16 字节）
    std::vector<BYTE> iv = {
        0x10,0x11,0x12,0x13,0x14,0x15,0x16,0x17,
        0x18,0x19,0x1A,0x1B,0x1C,0x1D,0x1E,0x1F
    };

    std::vector<BYTE> ciphertext;
    std::vector<BYTE> decrypted;

    // 调用加密接口
    NTSTATUS status = AesCbcEncrypt(key, iv, plaintext, ciphertext);
    if (!NT_SUCCESS(status)) {
        std::cerr << "[ERROR] Encryption failed: 0x" << std::hex << status << std::dec << "\n";
        return 1;
    }

    std::cout << "Plaintext:  " << plaintext_str << "\n";
    std::cout << "Ciphertext: " << BytesToHex(ciphertext) << "\n";

    // 注意：此时外部的 `iv` 仍然完好（未被修改）

    // 调用解密接口
    status = AesCbcDecrypt(key, iv, ciphertext, decrypted);
    if (!NT_SUCCESS(status)) {
        std::cerr << "[ERROR] Decryption failed: 0x" << std::hex << status << std::dec << "\n";
        return 1;
    }

    std::string decrypted_str(decrypted.begin(), decrypted.end());
    std::cout << "Decrypted:  " << decrypted_str << "\n";

    if (plaintext_str == decrypted_str) {
        std::cout << "[SUCCESS] Modular AES-CBC verified!\n";
    }
    else {
        std::cerr << "[ERROR] Mismatch!\n";
        return 1;
    }

    return 0;
}