/*******************************************************************************
 * 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 <stdexcept>
#include <algorithm>
#include <numeric>
#include <cstring>
#include <iostream>
#include <iomanip>
#include <sstream>

#include "des.h"

using namespace std;

// ==============================================================================
// 1. DES PERMUTATION AND SUBSTITUTION TABLES (UNMODIFIED)
// ==============================================================================
// (Table definitions remain here, outside the class and namespace, or inside an anonymous namespace)
// ... [IP, FP, PC1, SHIFTS, PC2, E, P, S tables remain unchanged]
namespace netplus{
    // ... [Permutation and Substitution Tables]
// Initial Permutation (IP) table
const int IP[] = {
    58, 50, 42, 34, 26, 18, 10, 2, 60, 52, 44, 36, 28, 20, 12, 4,
    62, 54, 46, 38, 30, 22, 14, 6, 64, 56, 48, 40, 32, 24, 16, 8,
    57, 49, 41, 33, 25, 17, 9, 1, 59, 51, 43, 35, 27, 19, 11, 3,
    61, 53, 45, 37, 29, 21, 13, 5, 63, 55, 47, 39, 31, 23, 15, 7
};

// Final Permutation (Inverse IP) table
const int FP[] = {
    40, 8, 48, 16, 56, 24, 64, 32, 39, 7, 47, 15, 55, 23, 63, 31,
    38, 6, 46, 14, 54, 22, 62, 30, 37, 5, 45, 13, 53, 21, 61, 29,
    36, 4, 44, 12, 52, 20, 60, 28, 35, 3, 43, 11, 51, 19, 59, 27,
    34, 2, 42, 10, 50, 18, 58, 26, 33, 1, 41, 9, 49, 17, 57, 25
};

// Permuted Choice 1 (PC-1) - selects 56 bits from the 64-bit key
const int PC1[] = {
    57, 49, 41, 33, 25, 17, 9, 1, 58, 50, 42, 34, 26, 18, 10, 2,
    59, 51, 43, 35, 27, 19, 11, 3, 60, 52, 44, 36, 63, 55, 47, 39,
    31, 23, 15, 7, 62, 54, 46, 38, 30, 22, 14, 6, 61, 53, 45, 37,
    29, 21, 13, 5, 28, 20, 12, 4
};

// Shift schedule for the 16 rounds
const int SHIFTS[] = {1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1};

// Permuted Choice 2 (PC-2) - selects 48 bits from the 56-bit shifted key
const int PC2[] = {
    14,17,11,24, 1, 5, 3,28,15, 6,21,10,
    23,19,12, 4,26, 8,16, 7,27,20,13, 2,
    41,52,31,37,47,55,30,40,51,45,33,48,
    44,49,39,56,34,53,46,42,50,36,29,32
};

// Expansion Permutation (E) - expands 32 bits to 48 bits
const int E[] = {
    32, 1, 2, 3, 4, 5, 4, 5, 6, 7, 8, 9, 8, 9, 10, 11,
    12, 13, 12, 13, 14, 15, 16, 17, 16, 17, 18, 19, 20, 21, 20, 21,
    22, 23, 24, 25, 24, 25, 26, 27, 28, 29, 28, 29, 30, 31, 32, 1
};

// Permutation P
const int P[] = {
    16, 7, 20, 21, 29, 12, 28, 17, 1, 15, 23, 26, 5, 18, 31, 10,
    2, 8, 24, 14, 32, 27, 3, 9, 19, 13, 30, 6, 22, 11, 4, 25
};

// S-Boxes (8 boxes, each 4x16 -> 64 values)
const int S[8][64] = {
    {14, 4, 13, 1, 2, 15, 11, 8, 3, 10, 6, 12, 5, 9, 0, 7, 0, 15, 7, 4, 14, 2, 13, 1, 10, 6, 12, 11, 9, 5, 3, 8, 4, 1, 14, 8, 13, 6, 2, 11, 15, 12, 9, 7, 3, 10, 5, 0, 15, 12, 8, 2, 4, 9, 1, 7, 5, 11, 3, 14, 10, 0, 6, 13},
    {15, 1, 8, 14, 6, 11, 3, 4, 9, 7, 2, 13, 12, 0, 5, 10, 3, 13, 4, 7, 15, 2, 8, 14, 12, 0, 1, 10, 6, 9, 11, 5, 0, 14, 7, 11, 10, 4, 13, 1, 5, 8, 12, 6, 9, 3, 2, 15, 13, 8, 10, 1, 3, 15, 4, 2, 11, 6, 7, 12, 0, 5, 14, 9},
    {10, 0, 9, 14, 6, 3, 15, 5, 1, 13, 12, 7, 11, 4, 2, 8, 13, 7, 0, 9, 3, 4, 6, 10, 2, 8, 5, 14, 12, 11, 15, 1, 13, 6, 4, 9, 8, 15, 3, 0, 11, 1, 2, 12, 5, 10, 14, 7, 1, 10, 13, 0, 6, 9, 8, 7, 4, 15, 14, 3, 11, 5, 2, 12},
    {7, 13, 14, 3, 0, 6, 9, 10, 1, 2, 8, 5, 11, 12, 4, 15, 13, 8, 11, 5, 6, 15, 0, 3, 4, 7, 2, 12, 1, 10, 14, 9, 10, 6, 9, 0, 12, 11, 7, 13, 15, 1, 3, 14, 5, 2, 8, 4, 3, 15, 0, 6, 10, 1, 13, 8, 9, 4, 5, 11, 12, 7, 2, 14},
    {2,12, 4, 1, 7,10,11, 6, 8, 5, 3,15,13, 0,14, 9,14,11, 2,12, 4, 7,13, 1, 5, 0,15,10, 3, 9, 8, 6,4, 2, 1,11,10,13, 7, 8,15, 9,12, 5, 6, 3, 0,14,11, 8,12, 7, 1,14, 2,13, 6,15, 0, 9,10, 4, 5, 3},
    {12, 1, 10, 15, 9, 2, 6, 8, 0, 13, 3, 4, 14, 7, 5, 11, 10, 15, 4, 2, 7, 12, 9, 5, 6, 1, 13, 14, 0, 11, 3, 8, 9, 14, 15, 5, 2, 8, 12, 3, 7, 0, 4, 10, 1, 13, 11, 6, 4, 3, 2, 12, 9, 5, 15, 10, 11, 14, 1, 7, 6, 0, 8, 13},
    {4, 11, 2, 14, 15, 0, 8, 13, 3, 12, 9, 7, 5, 10, 6, 1, 13, 0, 11, 7, 4, 9, 1, 10, 14, 3, 5, 12, 2, 15, 8, 6, 1, 4, 11, 13, 12, 3, 7, 14, 10, 15, 6, 8, 0, 5, 9, 2, 6, 11, 13, 8, 1, 4, 10, 7, 9, 5, 0, 15, 14, 2, 3, 12},
    {13, 2, 8, 4, 6, 15, 11, 1, 10, 9, 3, 14, 5, 0, 12, 7, 1, 15, 13, 8, 10, 3, 7, 4, 12, 5, 6, 11, 0, 14, 9, 2, 7, 11, 4, 1, 9, 12, 14, 2, 0, 6, 10, 13, 15, 3, 5, 8, 2, 1, 14, 7, 4, 10, 8, 13, 15, 12, 9, 0, 3, 5, 6, 11}
};
} // end anonymous namespace

// ==============================================================================
// 2. BIT MANIPULATION HELPERS (UPDATED TO USE netplus::block64)
// ==============================================================================

namespace netplus {
    /**
     * @brief Converts a 64-bit vector (block64) into a 64-bit unsigned integer (uint64_t).
     */
    uint64_t vecToU64(const block64& vec) {
        uint64_t val = 0;
        for (size_t i = 0; i < 8; ++i) {
             val = (val << 8) | vec[i]; // This correctly reads vec[0] as MSB
        }
        return val;
    }

    /**
     * @brief Converts a 64-bit unsigned integer (uint64_t) into a 64-bit vector (block64).
     */
    block64 u64ToVec(uint64_t val) {
        block64 vec(8);
        for (int i = 7; i >= 0; --i) {
            vec[i] = static_cast<uint8_t>(val & 0xFF);
            val >>= 8;
        }
        return vec;
    }

    static inline uint64_t getBitBytesMSB(const uint8_t* bytes, int bit_pos_1based) {
        int bit_index = bit_pos_1based - 1;   // 0..63
        int byte_index = bit_index / 8;        // 0..7
        int bit_in_byte = bit_index % 8;       // 0..7 (MSB first)

        return (bytes[byte_index] >> (7 - bit_in_byte)) & 1ULL;
    }

    static inline uint64_t getBitMSB(uint64_t x, int bit_pos_1based, int total_bits) {
        return (x >> (total_bits - bit_pos_1based)) & 1ULL;
    }

    uint64_t permuteBytes(const uint8_t* in, const int* table, int out_len) {
        uint64_t result = 0;
        for (int i = 0; i < out_len; ++i) {
            result <<= 1;
            result |= getBitBytesMSB(in, table[i]);
        }
        return result;
    }

    uint64_t permute(uint64_t input, const int* table, int out_len, int in_len) {
        uint64_t result = 0;
        for (int i = 0; i < out_len; ++i) {
            result <<= 1;
            result |= getBitMSB(input, table[i], in_len);
        }
        return result;
    }

    uint64_t permuteU64(uint64_t input, const int* table, int out_len, int in_len = 64) {
        if (in_len < 64) input <<= (64 - in_len); // left-align window

        uint8_t bytes[8] = {
            (uint8_t)((input >> 56) & 0xFF),
            (uint8_t)((input >> 48) & 0xFF),
            (uint8_t)((input >> 40) & 0xFF),
            (uint8_t)((input >> 32) & 0xFF),
            (uint8_t)((input >> 24) & 0xFF),
            (uint8_t)((input >> 16) & 0xFF),
            (uint8_t)((input >>  8) & 0xFF),
            (uint8_t)((input >>  0) & 0xFF)
        };

        uint64_t result = 0;
        for (int i = 0; i < out_len; ++i) {
            result <<= 1;
            result |= getBitBytesMSB(bytes, table[i]); // NO offset
        }
        return result;
    }

    netplus::block64 byte_reverse_64(const netplus::block64& block) {
        if (block.size() != 8) {
            // Handle error or return copy
            return block;
        }
        netplus::block64 reversed_block = block;
        // std::reverse swaps bytes: 01 23 45 67 89 AB CD EF -> EF CD AB 89 67 45 23 01
        std::reverse(reversed_block.begin(), reversed_block.end());
        return reversed_block;
    }

    des::des(const block64& key) {
        if (key.size() != 8) {
            throw std::invalid_argument("DES key must be exactly 8 bytes (64 bits).");
        }
        m_masterKey = vecToU64(key);
        generateSubKeys();
    }

    des::~des() = default;

    // --- Key Schedule ---

    void netplus::des::generateSubKeys() {
        // 1. Apply PC-1 (64-bit key -> 56-bit key)
        uint64_t key56 = permute(m_masterKey, PC1, 56, 64);

        // CRITICAL FIX: Ensure key56 is exactly 56 bits for clean split
        const uint64_t MASK_56_BIT = 0x00FFFFFFFFFFFFFFULL;
        key56 &= MASK_56_BIT;

        // Split into C0 (MSB 28 bits) and D0 (LSB 28 bits)
        uint32_t C = static_cast<uint32_t>((key56 >> 28) & 0x0FFFFFFF);
        uint32_t D = static_cast<uint32_t>(key56 & 0x0FFFFFFF);

        const uint32_t MASK_28_BIT = 0x0FFFFFFF;

        for (int i = 0; i < 16; ++i) {
            int shift = SHIFTS[i];

            // CANONICAL 28-BIT CIRCULAR LEFT SHIFT
            // Performs rotation and masks the result to guarantee only 28 bits remain.
            C = ((C << shift) | (C >> (28 - shift))) & MASK_28_BIT;
            D = ((D << shift) | (D >> (28 - shift))) & MASK_28_BIT;

            // 3. Combine C and D (56 bits)
            uint64_t CD56 = ((static_cast<uint64_t>(C) << 28) | D) & 0x00FFFFFFFFFFFFFFULL;
            // 4. Apply PC-2 (56-bit key -> 48-bit round key)
            uint64_t roundKey48 = permute(CD56, PC2, 48, 56);
            // Store the 48-bit key as a 6-byte vector (block64), MSB first
            m_roundKeys48[i] = roundKey48 & 0xFFFFFFFFFFFFULL;
        }
    }

    // In des::initialPermutation:
    uint64_t des::initialPermutation(const block64& block) const {
        return permuteBytes(block.data(), IP, 64);
    }

    // In des::finalPermutation:
    block64 des::finalPermutation(const uint64_t preoutput) const {
        block64 out(8);
        uint64_t p = permuteU64(preoutput, FP, 64); // or use byte variant
        for (int i = 7; i >= 0; --i) {
            out[i] = p & 0xFF;
            p >>= 8;
        }
        return out;
    }

    // --- Feistel Function (F) ---

    // This method is part of the netplus::DES class in des.cpp
    netplus::block64 netplus::des::feistelFunction(const netplus::block64& rightHalf,uint64_t roundKey48) const {

        // Right half is 32 bits (4 bytes)
        uint32_t R = (static_cast<uint32_t>(rightHalf[0]) << 24) |
                    (static_cast<uint32_t>(rightHalf[1]) << 16) |
                    (static_cast<uint32_t>(rightHalf[2]) << 8) |
                    static_cast<uint32_t>(rightHalf[3]);

        // 1. Expansion Permutation (E) - 32 bits to 48 bits
        uint64_t E_expanded = permute(static_cast<uint64_t>(R), E, 48, 32) & 0xFFFFFFFFFFFFULL;

        uint64_t K_round = roundKey48 & 0xFFFFFFFFFFFFULL;
        uint64_t B = (E_expanded & 0xFFFFFFFFFFFFULL) ^ K_round;

        // 3. S-Box Substitution (48 -> 32), MSB-first
        uint32_t S_output = 0;

        for (int i = 0; i < 8; ++i) {
            // S1 takes the top 6 bits, S8 takes the bottom 6 bits
            uint8_t box_input = static_cast<uint8_t>((B >> (42 - 6*i)) & 0x3F);

            int row = ((box_input & 0x20) >> 4) | (box_input & 0x01);
            int col = (box_input & 0x1E) >> 1;

            uint32_t sub_val = static_cast<uint32_t>(S[i][row * 16 + col]);

            // pack nibbles MSB-first into 32-bit
            S_output |= (sub_val << (28 - 4*i));
        }


        // 4. P-Box Permutation (32 bits)
        uint32_t P_permuted = static_cast<uint32_t>(permute(static_cast<uint64_t>(S_output), P, 32, 32));


        // Convert 32-bit result back to block64 (4 bytes)
        netplus::block64 result(4);
        result[0] = static_cast<uint8_t>((P_permuted >> 24) & 0xFF);
        result[1] = static_cast<uint8_t>((P_permuted >> 16) & 0xFF);
        result[2] = static_cast<uint8_t>((P_permuted >> 8) & 0xFF);
        result[3] = static_cast<uint8_t>(P_permuted & 0xFF);

        return result;
    }

    netplus::block64 des::encrypt(const netplus::block64& plaintext) const {
        if (plaintext.size() != 8) {
            throw std::invalid_argument("Plaintext block must be 8 bytes (64 bits).");
        }

        netplus::block64 ipBlock = u64ToVec(initialPermutation(plaintext));

        uint32_t L = (uint32_t(ipBlock[0]) << 24) |
                    (uint32_t(ipBlock[1]) << 16) |
                    (uint32_t(ipBlock[2]) <<  8) |
                    uint32_t(ipBlock[3]);

        uint32_t R = (uint32_t(ipBlock[4]) << 24) |
                    (uint32_t(ipBlock[5]) << 16) |
                    (uint32_t(ipBlock[6]) <<  8) |
                    uint32_t(ipBlock[7]);

        for (int i = 0; i < 16; ++i) {
            netplus::block64 Rbytes(4);
            Rbytes[0] = uint8_t((R >> 24) & 0xFF);
            Rbytes[1] = uint8_t((R >> 16) & 0xFF);
            Rbytes[2] = uint8_t((R >>  8) & 0xFF);
            Rbytes[3] = uint8_t((R >>  0) & 0xFF);

            netplus::block64 f = feistelFunction(Rbytes, m_roundKeys48[i]);
            uint32_t F = (uint32_t(f[0]) << 24) |
                        (uint32_t(f[1]) << 16) |
                        (uint32_t(f[2]) <<  8) |
                        uint32_t(f[3]);

            uint32_t newL = R;
            uint32_t newR = L ^ F;
            L = newL;
            R = newR;
        }

        uint64_t preoutput = (uint64_t(R) << 32) | uint64_t(L);
        return finalPermutation(preoutput);
    }

    netplus::block64 des::decrypt(const netplus::block64& ciphertext) const {
        if (ciphertext.size() != 8) {
            throw std::invalid_argument("Ciphertext block must be 8 bytes (64 bits).");
        }

        // IP
        netplus::block64 ipBlock = u64ToVec(initialPermutation(ciphertext));

        // Split into L and R (32 bits each)
        uint32_t L = (uint32_t(ipBlock[0]) << 24) |
                    (uint32_t(ipBlock[1]) << 16) |
                    (uint32_t(ipBlock[2]) <<  8) |
                    uint32_t(ipBlock[3]);

        uint32_t R = (uint32_t(ipBlock[4]) << 24) |
                    (uint32_t(ipBlock[5]) << 16) |
                    (uint32_t(ipBlock[6]) <<  8) |
                    uint32_t(ipBlock[7]);

        // 16 rounds with keys reversed
        for (int i = 15; i >= 0; --i) {
            // Build R half as 4 bytes for feistelFunction (expects block64 of size 4)
            netplus::block64 Rbytes(4);
            Rbytes[0] = uint8_t((R >> 24) & 0xFF);
            Rbytes[1] = uint8_t((R >> 16) & 0xFF);
            Rbytes[2] = uint8_t((R >>  8) & 0xFF);
            Rbytes[3] = uint8_t((R >>  0) & 0xFF);

            // F(R, K_i)
            netplus::block64 f = feistelFunction(Rbytes, m_roundKeys48[i]);
            uint32_t F = (uint32_t(f[0]) << 24) |
                        (uint32_t(f[1]) << 16) |
                        (uint32_t(f[2]) <<  8) |
                        uint32_t(f[3]);

            // Standard Feistel step: (L, R) := (R, L XOR F(R, K_i))
            uint32_t newL = R;
            uint32_t newR = L ^ F;
            L = newL;
            R = newR;
        }

        // Final swap: R16 || L16
        uint64_t preoutput = (uint64_t(R) << 32) | uint64_t(L);

        // FP
        return finalPermutation(preoutput);
    }


    // Revert to direct assignment (which passed the 3DES Roundtrip Test)
    trippledes::trippledes(const block64& key1, const block64& key2, const block64& key3)
        : m_des1(key1), m_des2(key2), m_des3(key3) {}



    /**
     * @brief Encrypts a 64-bit block using Triple DES EDE (Encrypt-Decrypt-Encrypt) mode: C = E_K3(D_K2(E_K1(P))).
     *
     * @param plaintext The 8-byte (64-bit) block to encrypt.
     * @return The 8-byte (64-bit) ciphertext block.
     */
    block64 trippledes::encrypt(const block64& plaintext) const {
        if (plaintext.size() != 8) {
            throw std::invalid_argument("Plaintext block must be 8 bytes (64 bits) for Triple DES.");
        }
        // Matches dump label K1-K2-K3: C = E_K3(D_K2(E_K1(P)))
        return m_des3.encrypt(m_des2.decrypt(m_des1.encrypt(plaintext)));
    }




    /**
     * @brief Decrypts a 64-bit block using Triple DES EDE (Encrypt-Decrypt-Encrypt) mode (Decryption): P = D_K1(E_K2(D_K3(C))).
     *
     * @param ciphertext The 8-byte (64-bit) block to decrypt.
     * @return The 8-byte (64-bit) plaintext block.
     */
    block64 trippledes::decrypt(const block64& ciphertext) const {
        if (ciphertext.size() != 8) {
            throw std::invalid_argument("Ciphertext block must be 8 bytes (64 bits) for Triple DES.");
        }
        // Inverse: P = D_K1(E_K2(D_K3(C)))
        return m_des1.decrypt(m_des2.encrypt(m_des3.decrypt(ciphertext)));
    }
} // namespace netplus
