test_crypto.cpp

#include <windows.h>
#include <iostream>
#include <cassert>
#include <string>
#include <vector>
#include <winrt/Windows.Foundation.h>
#include "CryptoService.h"
#include "StorageService.h"

// Define simple check macro
#define ASSERT_TRUE(x) if (!(x)) { std::cerr << "Assertion failed: " << #x << " at line " << __LINE__ << std::endl; exit(1); }

int main()
{
    std::cout << "Starting Aegis Cryptography and Storage tests..." << std::endl;

    // Initialize WinRT for JSON parsing
    winrt::init_apartment();

    // 1. Base64 Tests
    std::string testStr = "Aegis Secure Cryptography Suite";
    std::vector<uint8_t> testBytes(testStr.begin(), testStr.end());
    std::string b64 = Aegis::Cryptography::CryptoService::BytesToBase64(testBytes);
    std::cout << "Base64 encoded: " << b64 << std::endl;
    std::vector<uint8_t> decodedBytes = Aegis::Cryptography::CryptoService::Base64ToBytes(b64);
    std::string decodedStr(decodedBytes.begin(), decodedBytes.end());
    ASSERT_TRUE(decodedStr == testStr);
    std::cout << "Base64 Test passed!" << std::endl;

    // 2. PQ Keygen KEM (ML-KEM-768)
    std::cout << "Generating ML-KEM-768..." << std::endl;
    Aegis::Cryptography::KeyPair kemPair = Aegis::Cryptography::CryptoService::GeneratePQKey(L"ML-KEM-768");
    std::cout << "Public key size: " << kemPair.PublicKey.size() << " bytes" << std::endl;
    std::cout << "Secret key size: " << kemPair.PrivateKey.size() << " bytes" << std::endl;
    ASSERT_TRUE(!kemPair.PublicKey.empty() && !kemPair.PrivateKey.empty());
    std::cout << "ML-KEM-768 Keygen passed!" << std::endl;

    // 3. Hybrid Encryption / Decryption
    std::cout << "Testing Hybrid Encryption (ML-KEM-768 + AES-256-GCM)..." << std::endl;
    std::string secretMessage = "Quantum computer resistant secret message!";
    std::vector<uint8_t> msgBytes(secretMessage.begin(), secretMessage.end());
    Aegis::Cryptography::HybridCiphertext ct = Aegis::Cryptography::CryptoService::EncryptHybrid(L"ML-KEM-768", kemPair.PublicKey, msgBytes);
    
    std::vector<uint8_t> decrypted = Aegis::Cryptography::CryptoService::DecryptHybrid(L"ML-KEM-768", kemPair.PrivateKey, ct);
    std::string decryptedStr(decrypted.begin(), decrypted.end());
    std::cout << "Decrypted message: " << decryptedStr << std::endl;
    ASSERT_TRUE(decryptedStr == secretMessage);
    std::cout << "Hybrid Encrypt/Decrypt passed!" << std::endl;

    // 4. PQ Signing (ML-DSA-65)
    std::cout << "Generating ML-DSA-65 Key..." << std::endl;
    Aegis::Cryptography::KeyPair sigPair = Aegis::Cryptography::CryptoService::GeneratePQKey(L"ML-DSA-65");
    std::string document = "This is a document signed with ML-DSA-65.";
    std::vector<uint8_t> docBytes(document.begin(), document.end());
    std::vector<uint8_t> signature = Aegis::Cryptography::CryptoService::SignPQ(L"ML-DSA-65", sigPair.PrivateKey, docBytes);
    
    bool verified = Aegis::Cryptography::CryptoService::VerifyPQ(L"ML-DSA-65", sigPair.PublicKey, docBytes, signature);
    ASSERT_TRUE(verified);
    std::cout << "PQ Sign/Verify passed!" << std::endl;

    // 5. OpenPGP / Traditional Keygen (RSA)
    std::cout << "Generating RSA-2048..." << std::endl;
    Aegis::Cryptography::KeyPair rsaPair = Aegis::Cryptography::CryptoService::GeneratePGPKey(L"RSA", 2048);
    ASSERT_TRUE(!rsaPair.PublicKey.empty() && !rsaPair.PrivateKey.empty());
    
    // PEM Formatting
    std::string pemPub = Aegis::Cryptography::CryptoService::ExportPEM(L"RSA", rsaPair.PublicKey, false);
    std::cout << "RSA Public Key PEM:\n" << pemPub << std::endl;
    std::vector<uint8_t> importedPub = Aegis::Cryptography::CryptoService::ImportPEM(pemPub, false);
    ASSERT_TRUE(importedPub == rsaPair.PublicKey);
    std::cout << "PEM Export/Import passed!" << std::endl;

    // OpenPGP Formatting
    std::string pgpArmor = Aegis::Cryptography::CryptoService::ExportOpenPGP(L"RSA", L"Alice Smith", L"alice@example.com", rsaPair.PublicKey, rsaPair.PrivateKey, false);
    std::cout << "OpenPGP ASCII Armored Key:\n" << pgpArmor << std::endl;
    std::cout << "OpenPGP Formatting passed!" << std::endl;

    // 6. Storage Service Tests
    std::cout << "Testing Storage Service..." << std::endl;
    auto& storage = Aegis::Storage::StorageService::GetInstance();
    
    // Ensure locked initially
    ASSERT_TRUE(storage.IsLocked());

    // Create new keyring
    std::wstring masterPass = L"AegisMasterPassword123!";
    bool keyringCreated = storage.CreateKeyring(masterPass);
    ASSERT_TRUE(keyringCreated);
    ASSERT_TRUE(!storage.IsLocked());

    // Add key to keyring
    Aegis::Storage::KeyRecord record;
    record.Id = L"test-key-id-111";
    record.Name = L"Work PQ Key";
    record.Email = L"work@company.com";
    record.Type = L"pq";
    record.Algorithm = L"ML-KEM-768";
    record.Fingerprint = L"FINGERPRINT123456789";
    record.CreatedAt = L"2026-06-01T12:00:00Z";
    record.PublicKey = Aegis::Cryptography::CryptoService::BytesToBase64(kemPair.PublicKey);
    record.PrivateKey = Aegis::Cryptography::CryptoService::BytesToBase64(kemPair.PrivateKey);

    bool added = storage.AddKey(record);
    ASSERT_TRUE(added);

    // Lock keyring
    storage.LockKeyring();
    ASSERT_TRUE(storage.IsLocked());

    // Verify retrieval fails when locked
    std::vector<Aegis::Storage::KeyRecord> keysLocked = storage.GetKeys();
    ASSERT_TRUE(keysLocked.empty());

    // Unlock keyring
    bool unlocked = storage.UnlockKeyring(masterPass);
    ASSERT_TRUE(unlocked);
    ASSERT_TRUE(!storage.IsLocked());

    // Retrieve keys
    std::vector<Aegis::Storage::KeyRecord> keysUnlocked = storage.GetKeys();
    ASSERT_TRUE(keysUnlocked.size() == 1);
    ASSERT_TRUE(keysUnlocked[0].Name == L"Work PQ Key");
    ASSERT_TRUE(keysUnlocked[0].PrivateKey.empty()); // Check that private key is hidden in general lists

    // Retrieve decrypted private key
    std::vector<uint8_t> decryptedPrivKey = storage.GetDecryptedPrivateKey(L"test-key-id-111");
    ASSERT_TRUE(decryptedPrivKey == kemPair.PrivateKey);

    // Delete key
    bool deleted = storage.DeleteKey(L"test-key-id-111");
    ASSERT_TRUE(deleted);
    ASSERT_TRUE(storage.GetKeys().empty());

    std::cout << "All Storage tests passed!" << std::endl;
    std::cout << "SUCCESS!" << std::endl;

    return 0;
}