use aes::cipher::{generic_array::GenericArray, BlockEncrypt, KeyInit};
use aes::cipher::generic_array::typenum::U16;
use aes::Aes128;
use rand::rngs::StdRng;
use rand::{RngCore, SeedableRng, Rng};
use sha2::{Sha256, Digest};
use argon2::{Params, Argon2};
use std::fs::read;
use std::fs::File;
use std::io::prelude::*;

fn main() -> std::io::Result<()> {
    // Read input file into a vector RENAME example.exe to what you want and put it in this "crypt" folder
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        println!("Run with {} <inputfile.exe>", args.get(0).unwrap());
        return Err(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            "file input not found",
        ));
    }
    let fname = args.get(1).unwrap();
    let plaintext_bytes = read(fname).expect("Failed to read file");

    // Create single combined output file
    let mut encrypted_file = File::create("encrypted_payload.bin")?;

    // Master password (in production, this should be securely provided)
    let master_password = "YourSecureMasterPassword123!"; // Change this!
    
    // Derive encryption key using simple PBKDF2-like approach
    let salt: [u8; 32] = rand::thread_rng().gen();

    let mut key_hasher = Sha256::new();
    key_hasher.update(master_password.as_bytes());
    key_hasher.update(&salt);
    let password_hash = key_hasher.finalize();

    let mut key = [0u8; 16];
    key.copy_from_slice(&password_hash[..16]);

    // Generate random IV for AES-CBC
    let mut iv = [0u8; 16];
    rand::thread_rng().fill_bytes(&mut iv);

    // Pad the plaintext with PKCS7
    let block_size = 16;
    let padding_size = block_size - (plaintext_bytes.len() % block_size);
    let mut padded_plaintext = plaintext_bytes.clone();
    padded_plaintext.extend(vec![padding_size as u8; padding_size]);

    // Encrypt with AES-CBC using manual implementation
    let cipher = Aes128::new(GenericArray::from_slice(&key));
    let mut ciphertext = Vec::new();
    let mut current_iv = iv;

    for chunk in padded_plaintext.chunks(16) {
        let mut block = GenericArray::clone_from_slice(chunk);

        // XOR with current IV
        for i in 0..16 {
            block[i] ^= current_iv[i];
        }

        // Encrypt
        cipher.encrypt_block(&mut block);

        // Use ciphertext as next IV
        current_iv.copy_from_slice(&block);

        ciphertext.extend_from_slice(&block);
    }

    // Write combined file: metadata + encrypted data
    encrypted_file.write_all(&salt)?;
    encrypted_file.write_all(&iv)?;
    encrypted_file.write_all(&(ciphertext.len() as u32).to_le_bytes())?;
    encrypted_file.write_all(&ciphertext)?;
    
    println!("Encryption complete!");
    println!("Encrypted file: encrypted_payload.bin");
    println!("Master password: {}", master_password);
    Ok(())
}
