//
// Created by mht on 20.04.25.
//
#include <fstream>
#include <vector>
#include <iostream>
#include <cstring>
#include <string>
#include <sstream>
#include "keya.h"

using namespace std;

// Helper to pad data to 16 bytes
void pad_data(vector<unsigned char>& data)
{
    size_t pad_len = 16 - (data.size() % 16);
    if (pad_len < 16) {
        data.insert(data.end(), pad_len, 0); // simple zero padding
    }
}

// Read file into vector
vector<unsigned char> read_file(const string& filename)
{
    ifstream file(filename, ios::binary);
    return vector<unsigned char>((istreambuf_iterator<char>(file)), istreambuf_iterator<char>());
}

// Write vector into file
void write_file(const string& filename, const vector<unsigned char>& data)
{
    ofstream file(filename, ios::binary);
    file.write(reinterpret_cast<const char*>(data.data()), data.size());
}

// Encrypt buffer using KeyA
int encrypt_buffer(vector<unsigned char>& data, unsigned char* kp, unsigned char* kr, unsigned char* kw)
{
    KeyA keya;
    int result = keya.OpenDevice(0);
    if (result != 0) return result;

    result = keya.Init(kp, kr, kw, NULL);
    if (result != 0) {
        keya.CloseDevice();
        return result;
    }

    unsigned char temp[16];
    for (size_t i = 0; i < data.size(); i += 16) {
        memcpy(temp, &data[i], 16);
        result = keya.EncByKV(0, temp, temp);
        if (result != 0) {
            keya.CloseDevice();
            return result;
        }
        memcpy(&data[i], temp, 16);
    }

    keya.CloseDevice();
    return 0;
}

// Decrypt buffer using KeyA
int decrypt_buffer(vector<unsigned char>& data, unsigned char* kp, unsigned char* kr, unsigned char* kw)
{
    KeyA keya;
    int result = keya.OpenDevice(0);
    if (result != 0) return result;

    result = keya.Init(kp, kr, kw, NULL);
    if (result != 0) {
        keya.CloseDevice();
        return result;
    }

    unsigned char temp[16];
    for (size_t i = 0; i < data.size(); i += 16) {
        memcpy(temp, &data[i], 16);
        result = keya.DecByKV(0, temp, temp);
        if (result != 0) {
            keya.CloseDevice();
            return result;
        }
        memcpy(&data[i], temp, 16);
    }

    keya.CloseDevice();
    return 0;
}

// Save size as first 8 bytes (uint64_t) of encrypted file
void prepend_size(vector<unsigned char>& data, uint64_t size)
{
    vector<unsigned char> result(8);
    for (int i = 0; i < 8; ++i) {
        result[i] = (size >> (i * 8)) & 0xFF;
    }
    result.insert(result.end(), data.begin(), data.end());
    data = result;
}

// Extract original size from first 8 bytes
uint64_t extract_size(vector<unsigned char>& data)
{
    uint64_t size = 0;
    for (int i = 7; i >= 0; --i) {
        size = (size << 8) | data[i];
    }
    data.erase(data.begin(), data.begin() + 8);
    return size;
}

// Save list of hashes into a text file
void save_hashes_to_file(const vector<string>& hashes, const string& filename)
{
    ofstream file(filename);
    for (const auto& hash : hashes) {
        file << hash << "\n";
    }
}

// Read hashes from decrypted data
vector<string> read_hashes_from_memory(const vector<unsigned char>& data)
{
    vector<string> hashes;
    istringstream iss(string(data.begin(), data.end()));
    string line;
    while (getline(iss, line)) {
        if (!line.empty())
            hashes.push_back(line);
    }
    return hashes;
}

// Check if a hash exists in list
bool hash_exists(const string& input_hash, const vector<string>& hashes)
{
    for (const auto& hash : hashes) {
        if (hash == input_hash) {
            return true;
        }
    }
    return false;
}

int main()
{
    // Your keys
    unsigned char g_Kp[16]  = {0x74,0xD8,0xCD,0x0F,0xC6,0x6A,0x09,0x69,0xB1,0x4A,0xC5,0xD1,0xE9,0x7A,0x9D,0x07};
    unsigned char g_Kr[16]  = {0xDA,0x58,0x86,0x3C,0xD6,0x65,0xDF,0xE2,0x84,0x04,0xD3,0x4E,0x74,0xEB,0xBD,0x16};
    unsigned char g_Kw[16]  = {0xC8,0xA8,0x16,0x3C,0x59,0xEF,0x9C,0xC4,0xC0,0xEC,0x48,0x41,0x4B,0xE5,0x11,0x37};

    unsigned char* kp = g_Kp;
    unsigned char* kr = g_Kr;
    unsigned char* kw = g_Kw;

    // Step 1: Save some hashes to file
    vector<string> hashes = {
            "8d523c784e2f81de923e2fad5960983b735aca25b2490a1cd50d91dc4b166795b27e59f8dfb76ef5120460224cc9ab5d489dc9ae135a732588dd43f71048771c",
            "dd590a463303424547f76562fbaeca485c1b95581b7340ecdd02e308166ddcfcac757ed23b11931d23ca46e1b6498fe3c95bf2fd7f760337b41d9e7f4e396678",
            "75863d940c8f298c92ccdcfeca6d809d2e5bf67d23c588d2bd62f4b8f86be8a882e30b9b5e524b79e8cf618290b18eec8b621d9c9d73f5300246bd3e15c99561",
            "b6d7ec0a2520e2a2ba8f7bb4807b4f3641b80f1062075028600201b1b01e4ac102e078061cd4618f2e847348b118a11be1341a9d3a41fcfc443f92c69c40ffa1",
            "4347b58249eea37f74c8a036ca29d6b2a362b9a8743bbd4f99de75c92b46fd876bf98a65041c28c9f10eb548df7645079fe616b00e42711d52135034fe5dbc55",
            "a254cc5e08d0d11da6074fccbbc5fbff7a2ce99e6a36f6294f66e271455d3f0663f70864a74ea7d307c0ad7d76fa30694b52dcd9385a8de21b9a708a791d6a8c",
            "f9f0ce3eed54de9e057fd66ae17e4b2dccb5b4b6c17e18a278e6628a8a5ac4e357f06c510370e06da2a62241a3200f3d5334102e045c36b5a74954f2ff070a79",
            "a1475fdb4b02de84a57f85babc010e5f12281eedea1c75efeac54af7e40c0dfa1659eaa0b106c6aa3a605c7ec442f477d425db24b75c8a818750d9fe06e3727a",
            "183d402bb712787f5aa8bf3f0cbe91d3383fdd460edac3db01ea756ae5afe9ca60d8917d749358a012b2b45eaae04ae5fc10fe17aa42409841dc5b475c2794a1",
            "a29fc7db04adb89736832cd94dd1fd3fe405e18db670ef3f055e3310dd2f0fc4dd3c4535adcd1ec3fcd16f45a8b814558586bfc67b1d01d664cd51ee1db05c28",
            "0092d1a3a78bab2b85abc55d5a851c3c351d0dacd402d77ad57b6d3f08836262bc7dff8b958368cddedd3aa8246ec4919f77c325210cabeadd2dfde5ce8d0412",
            "0485a2fd8aceaddfeb45f4072f64637bec8ab5962e088c4f8f0a46196d6b80c1a2924e1d43fa530b2183e5240507115a3db8d49aae9d100195a6ba1eb8f1adcf",
            "ec23e51688c082584821a20038510a818c58a0201dbb7ee133d0b11dfc6456a2ad1de4dd944e8083e378af4e58362b5c3ecf3f242493e0cfc1c1cde58e7fc8f1"

    };
    save_hashes_to_file(hashes, "hashes.txt");
    cout << "Hashes saved to 'hashes.txt'." << endl;

    // Step 2: Encrypt the hashes file
    vector<unsigned char> hash_data = read_file("hashes.txt");

    uint64_t original_size = hash_data.size(); // Save original size
    pad_data(hash_data); // pad to multiple of 16

    if (encrypt_buffer(hash_data, kp, kr, kw) == 0) {
        prepend_size(hash_data, original_size); // Put size in front
        write_file("lic.enc", hash_data);
        cout << "Encryption successful. Encrypted file saved as 'lic.enc'." << endl;
    } else {
        cout << "Encryption failed." << endl;
        return 1;
    }

    // Step 3: Later... Decrypt and check for input hash
    vector<unsigned char> enc_data = read_file("hashes.enc");
    uint64_t saved_size = extract_size(enc_data);

    if (decrypt_buffer(enc_data, kp, kr, kw) == 0) {
        enc_data.resize(saved_size); // trim padding
        vector<string> decrypted_hashes = read_hashes_from_memory(enc_data);

        string input_hash;
        cout << "Enter a hash to search for: ";
        cin >> input_hash;

        if (hash_exists(input_hash, decrypted_hashes)) {
            cout << "Hash FOUND in list." << endl;
        } else {
            cout << "Hash NOT FOUND." << endl;
        }
    } else {
        cout << "Decryption failed." << endl;
        return 1;
    }

    return 0;
}