//
// Revised by assistant to fix segfaults and robustness issues
// Created from original code: mht 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 (only if needed)
void pad_data(vector<unsigned char>& data)
{
    size_t rem = data.size() % 16;
    if (rem != 0) {
        size_t pad_len = 16 - rem;
        data.insert(data.end(), pad_len, 0); // zero padding
    }
}

// Read file into vector (with simple error handling)
vector<unsigned char> read_file(const string& filename)
{
    ifstream file(filename, ios::binary);
    if (!file) {
        cerr << "Failed to open file for reading: " << filename << endl;
        return {};
    }
    return vector<unsigned char>((istreambuf_iterator<char>(file)), istreambuf_iterator<char>());
}

// Write vector into file (with simple error handling)
bool write_file(const string& filename, const vector<unsigned char>& data)
{
    ofstream file(filename, ios::binary);
    if (!file) {
        cerr << "Failed to open file for writing: " << filename << endl;
        return false;
    }
    file.write(reinterpret_cast<const char*>(data.data()), data.size());
    return true;
}

// 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 -- little endian
void prepend_size(vector<unsigned char>& data, uint64_t size)
{
    vector<unsigned char> result(8);
    for (int i = 0; i < 8; ++i) {
        result[i] = static_cast<unsigned char>((size >> (i * 8)) & 0xFF);
    }
    result.insert(result.end(), data.begin(), data.end());
    data.swap(result);
}

// Extract original size from first 8 bytes (little endian) and remove header.
// Returns pair(success, size). On failure (not enough bytes) returns {false,0}.
pair<bool,uint64_t> extract_size(vector<unsigned char>& data)
{
    if (data.size() < 8) {
        cerr << "Encrypted data too short to contain size header (need 8 bytes)." << endl;
        return {false, 0};
    }
    uint64_t size = 0;
    // reconstruct little-endian
    for (int i = 0; i < 8; ++i) {
        size |= (static_cast<uint64_t>(data[i]) << (i * 8));
    }
    // remove header
    data.erase(data.begin(), data.begin() + 8);
    return {true, size};
}

// Save list of hashes into a text file
void save_hashes_to_file(const vector<string>& hashes, const string& filename)
{
    ofstream file(filename);
    if (!file) {
        cerr << "Failed to open " << filename << " for writing." << endl;
        return;
    }
    for (const auto& hash : hashes) {
        file << hash << "\n";
    }
}

// Read hashes from decrypted data (handles embedded NULs correctly)
vector<string> read_hashes_from_memory(const vector<unsigned char>& data)
{
    cout << "entered read_hashes_from_memory" << endl;
    vector<string> hashes;
    // construct a string with explicit size to handle binary / NULs
    string s(reinterpret_cast<const char*>(data.data()), data.size());
    cout << "data size: " << data.size() << endl;
    istringstream iss(s);
    string line;
    while (getline(iss, line)) {
        if (!line.empty()) {
            cout << "readed hash: " << line << endl;
            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");
    if (hash_data.empty()) {
        cerr << "hashes.txt is empty or could not be read. Aborting encryption." << endl;
        return 1;
    }

    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
        if (!write_file("lic.enc", hash_data)) {
            cerr << "Failed to write lic.enc" << endl;
            return 1;
        }
        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
    // NOTE: read the correct file 'lic.enc' (was 'hashes.enc' in original -> bug)
    vector<unsigned char> enc_data = read_file("lic.enc");
    if (enc_data.empty()) {
        cerr << "Failed to read 'lic.enc' or file empty." << endl;
        return 1;
    }
    cout << "read 'lic.enc' successful" << endl;

    // extract original size
    auto [ok, saved_size] = extract_size(enc_data);
    if (!ok) {
        cerr << "Failed to extract size header from encrypted file." << endl;
        return 1;
    }

    vector<string> decrypted_hashes; // will be filled below

    try {
        if (decrypt_buffer(enc_data, kp, kr, kw) == 0) {
            if (saved_size > enc_data.size()) {
                cerr << "Saved original_size is larger than decrypted data. Adjusting to available size." << endl;
                saved_size = enc_data.size();
            }
            enc_data.resize(saved_size); // trim padding
            try {
                // fill the outer decrypted_hashes (do NOT shadow)
                decrypted_hashes = read_hashes_from_memory(enc_data);
                cout << "decryption of 'lic.enc' successful" << endl;
            }
            catch (exception &e) {
                cout << "Parsing decrypted payload failed: " << e.what() << endl;
                return 1;
            }

            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;
        }
    } catch (exception &e) {
        cout << "Failed to decrypt buffer: " << e.what() << endl;
        return 1;
    }

    return 0;
}