/*******************************************************************************
 * 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 "rc4.h"
#include <cstring>

// Helper function to swap two elements
void netplus::rc4::swap(uint8_t& a, uint8_t& b) {
    uint8_t temp = a;
    a = b;
    b = temp;
}

/**
 * Key-Scheduling Algorithm (KSA)
 * Initializes the state vector S based on the key.
 */
void netplus::rc4::ksa(const uint8_t* key, size_t key_len) {
    if (key_len == 0 || key_len > SBOX_SIZE) {
        throw std::invalid_argument("rc4 Key-Länge muss zwischen 1 und 256 Bytes liegen.");
    }

    // 1. Initialization of the S-Box: S[i] = i
    for (int i = 0; i < SBOX_SIZE; ++i) {
        S[i] = (uint8_t)i;
    }

    // 2. Initial Permutation of S
    int j = 0;
    for (int i = 0; i < SBOX_SIZE; ++i) {
        // j = (j + S[i] + K[i mod key_len]) mod 256
        j = (j + S[i] + key[i % key_len]) % SBOX_SIZE;
        swap(S[i], S[j]);
    }

    // After KSA, the PRGA state is reset to 0.
    reset_prga_state();
}

/**
 * Konstruktor: Initialisiert KSA.
 */
netplus::rc4::rc4(const std::vector<uint8_t>& key) {
    // Ruft KSA mit den Schlüsselbytes auf
    ksa(key.data(), key.size());
}

/**
 * Pseudo-Random Generation Algorithm (PRGA) - generates one keystream byte
 * Generates the keystream output based on the S-box.
 */
uint8_t netplus::rc4::PRGA_byte() {
    // i = (i + 1) mod 256
    i_prga = (i_prga + 1) % SBOX_SIZE;

    // j = (j + S[i]) mod 256
    j_prga = (j_prga + S[i_prga]) % SBOX_SIZE;

    // Swap S[i] and S[j]
    swap(S[i_prga], S[j_prga]);

    // t = (S[i] + S[j]) mod 256
    int t = (S[i_prga] + S[j_prga]) % SBOX_SIZE;

    // Keystream-Byte = S[t]
    return S[t];
}

/**
* Encrypts/Decrypts data
*/
std::vector<uint8_t> netplus::rc4::crypt(const std::vector<uint8_t>& data) {
    std::vector<uint8_t> output;
    output.reserve(data.size());

    for (uint8_t data_byte : data) {
        // Keystream-Byte generieren und mit dem Datenbyte XORen
        output.push_back(data_byte ^ PRGA_byte());
    }

    return output;
}

/**
 * Resets the PRGA state
*/
void netplus::rc4::reset_prga_state() {
    i_prga = 0;
    j_prga = 0;
}
