#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <cstring>
#include <openssl/evp.h>
#include <openssl/aes.h>
#include <openssl/rand.h>
#include <openssl/err.h>
#include <sys/stat.h>
#include <unistd.h>

class AESCBCDecryptor {
private:
    std::vector<uint8_t> key;

public:
    bool deriveKey(const std::string& password, const std::vector<uint8_t>& salt) {
        // Match the Rust key derivation: SHA256(password + salt), take first 16 bytes
        EVP_MD_CTX* sha256 = EVP_MD_CTX_new();
        if (!sha256) return false;

        if (EVP_DigestInit_ex(sha256, EVP_sha256(), nullptr) != 1) {
            EVP_MD_CTX_free(sha256);
            return false;
        }

        if (EVP_DigestUpdate(sha256, password.c_str(), password.length()) != 1) {
            EVP_MD_CTX_free(sha256);
            return false;
        }

        if (EVP_DigestUpdate(sha256, salt.data(), salt.size()) != 1) {
            EVP_MD_CTX_free(sha256);
            return false;
        }

        std::vector<uint8_t> hash(32);
        if (EVP_DigestFinal_ex(sha256, hash.data(), nullptr) != 1) {
            EVP_MD_CTX_free(sha256);
            return false;
        }

        EVP_MD_CTX_free(sha256);

        // Use first 16 bytes for AES-128
        key.assign(hash.begin(), hash.begin() + 16);

        return true;
    }

    std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext,
                              const std::vector<uint8_t>& iv,
                              const std::vector<uint8_t>& salt,
                              const std::string& password) {
        if (!deriveKey(password, salt)) {
            throw std::runtime_error("Key derivation failed");
        }

        EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
        if (!ctx) throw std::runtime_error("Failed to create cipher context");

        std::vector<uint8_t> plaintext(ciphertext.size());

        // Use ECB mode and implement CBC manually to match Rust implementation
        if (EVP_DecryptInit_ex(ctx, EVP_aes_128_ecb(), nullptr, key.data(), nullptr) != 1) {
            EVP_CIPHER_CTX_free(ctx);
            throw std::runtime_error("Failed to initialize decryption");
        }

        EVP_CIPHER_CTX_set_padding(ctx, 0);

        std::vector<uint8_t> current_iv = iv;
        int len;

        // Decrypt block by block implementing CBC manually
        for (size_t i = 0; i < ciphertext.size(); i += 16) {
            std::vector<uint8_t> encrypted_block(ciphertext.begin() + i, ciphertext.begin() + i + 16);
            std::vector<uint8_t> decrypted_block(16);

            // Decrypt the block (ECB mode)
            if (EVP_DecryptUpdate(ctx, decrypted_block.data(), &len, encrypted_block.data(), 16) != 1) {
                EVP_CIPHER_CTX_free(ctx);
                throw std::runtime_error("Decryption update failed");
            }

            // XOR with current IV
            for (size_t j = 0; j < 16; ++j) {
                decrypted_block[j] ^= current_iv[j];
            }

            // Copy to output
            std::copy(decrypted_block.begin(), decrypted_block.end(), plaintext.begin() + i);

            // Update IV for next block (use original ciphertext block)
            current_iv = encrypted_block;
        }

        EVP_CIPHER_CTX_free(ctx);

        // Remove PKCS7 padding manually
        if (!plaintext.empty()) {
            uint8_t padding_size = plaintext.back();
            if (padding_size <= 16 && padding_size > 0) {
                // Verify padding is correct
                bool padding_valid = true;
                for (size_t i = plaintext.size() - padding_size; i < plaintext.size(); ++i) {
                    if (plaintext[i] != padding_size) {
                        padding_valid = false;
                        break;
                    }
                }
                if (padding_valid) {
                    plaintext.resize(plaintext.size() - padding_size);
                }
            }
        }

        return plaintext;
    }
};

std::vector<uint8_t> readFile(const std::string& filename) {
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Cannot open file: " + filename);
    }
    
    file.seekg(0, std::ios::end);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    
    std::vector<uint8_t> data(size);
    file.read(reinterpret_cast<char*>(data.data()), size);
    return data;
}

void writeFile(const std::string& filename, const std::vector<uint8_t>& data) {
    std::ofstream file(filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Cannot create file: " + filename);
    }
    file.write(reinterpret_cast<const char*>(data.data()), data.size());
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        std::cout << "Usage: " << argv[0] << " <password>" << std::endl;
        return 1;
    }
    
    try {
        // Read encrypted DLL file
        std::vector<uint8_t> data = readFile("encrypted_dll.dll");

        if (data.size() < 52) {
            throw std::runtime_error("File too small");
        }

        // Parse: salt (32) + iv (16) + size (4) + ciphertext
        std::vector<uint8_t> salt(data.begin(), data.begin() + 32);
        std::vector<uint8_t> iv(data.begin() + 32, data.begin() + 48);
        uint32_t ciphertext_len = *reinterpret_cast<const uint32_t*>(&data[48]);
        std::vector<uint8_t> ciphertext(data.begin() + 52, data.begin() + 52 + ciphertext_len);

        // Decrypt
        AESCBCDecryptor decryptor;
        std::string password = argv[1];
        std::vector<uint8_t> plaintext = decryptor.decrypt(ciphertext, iv, salt, password);
        
        // Write decrypted binary
        writeFile("decrypted_binary", plaintext);
        
        std::cout << "Decryption successful!" << std::endl;
        std::cout << "Output: decrypted_binary" << std::endl;
        std::cout << "Size: " << plaintext.size() << " bytes" << std::endl;
        
        // Make executable on Unix systems
        chmod("decrypted_binary", 0755);
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}