/*******************************************************************************
 * Copyright (c) 2025, Jan Koester jan.koester@gmx.net
 * All rights reserved.
 * * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 * Neither the name of the <organization> nor the
 * names of its contributors may be used to endorse or promote products
 * derived from this software without specific prior written permission.
 * * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *******************************************************************************/

#include <iostream>
#include <vector>
#include <algorithm>

#include "../src/crypto/rc4.h"

#define ASSERT_TRUE(condition, message) if (!(condition)) { \
    std::cerr << "Assertion Failed: " << message << std::endl; \
    throw std::runtime_error("Assertion Failed"); \
}

// --- Helper Functions ---

// Converts string literal to byte vector
std::vector<uint8_t> string_to_bytes(const std::string& str) {
    return std::vector<uint8_t>(str.begin(), str.end());
}

// Converts a hex string (e.g., "BBF316...") to a byte vector
std::vector<uint8_t> hex_to_bytes(const std::string& hex_str) {
    std::vector<uint8_t> bytes;
    for (size_t i = 0; i < hex_str.length(); i += 2) {
        std::string byte_str = hex_str.substr(i, 2);
        bytes.push_back((uint8_t)std::stoul(byte_str, nullptr, 16));
    }
    return bytes;
}

// --- Unit Tests ---

void test_RC4_StandardVectorTest() {
    // Standard RC4 Test Vector:
    // Key: "Key"
    // Plaintext: "Plaintext"
    // Expected Ciphertext (Hex): BBF316E8D940AF0AD3
    try {
        std::vector<uint8_t> key = string_to_bytes("Key");
        std::vector<uint8_t> plaintext = string_to_bytes("Plaintext");
        std::vector<uint8_t> expected_ciphertext = hex_to_bytes("BBF316E8D940AF0AD3");

        netplus::rc4 cipher_enc(key);
        std::vector<uint8_t> actual_ciphertext = cipher_enc.crypt(plaintext);

        // Test 1: Check if ciphertext matches the standard vector
        ASSERT_TRUE(actual_ciphertext.size() == expected_ciphertext.size(),
                    "Ciphertext size mismatch.");
        ASSERT_TRUE(std::equal(actual_ciphertext.begin(), actual_ciphertext.end(),
                            expected_ciphertext.begin()),
                    "Ciphertext content mismatch against standard test vector.");

        // Test 2: Check decryption (must reset PRGA state first)

        netplus::rc4 cipher_dec(key);
        std::vector<uint8_t> actual_decrypted = cipher_dec.crypt(actual_ciphertext);

        ASSERT_TRUE(std::equal(actual_decrypted.begin(), actual_decrypted.end(),
                            plaintext.begin()),
                    "Decryption failed. Result does not match original plaintext.");

        std::cout << "PASSED [RC4_StandardVectorTest]" << std::endl;
    }catch(const std::exception& e){
        std::cerr << "FAILED [test_RC4_StandardVectorTest]: Exception: " << e.what() << std::endl;
    }
}

void test_RC4_LongKeyTest() {
    try {
        // RC4 works with keys up to 256 bytes.
        std::string long_key(256, 'A');
        std::vector<uint8_t> key = string_to_bytes(long_key);
        std::vector<uint8_t> data = string_to_bytes("Test Data");

        // ENCRYPTION: Use a cipher instance
        netplus::rc4 cipher_enc(key);
        std::vector<uint8_t> encrypted = cipher_enc.crypt(data);

        // Check if encryption resulted in some data
        ASSERT_TRUE(encrypted.size() == data.size(), "Long key encryption failed.");

        // DECRYPTION: Use a brand new cipher instance to reset the S-Box
        netplus::rc4 cipher_dec(key); // <--- FIX: Create new object for clean state
        std::vector<uint8_t> decrypted = cipher_dec.crypt(encrypted);

        ASSERT_TRUE(std::equal(decrypted.begin(), decrypted.end(), data.begin()),
                    "Long key decryption failed.");

        std::cout << "PASSED [RC4_LongKeyTest]" << std::endl;
    } catch(const std::exception& e) {
        std::cerr << "FAILED [test_RC4_LongKeyTest]: Exception: " << e.what() << std::endl;
    }
}

// --- Main Runner ---
int main() {
    std::cout << "--- Running RC4 Unit Tests ---" << std::endl;
    test_RC4_StandardVectorTest();
    test_RC4_LongKeyTest();
    std::cout << "--- Tests Complete ---" << std::endl;
    return 0;
}
