use hex;
use rand;
use std::fs::File;
use std::io::{Write, Read};
use std::io::{Seek, SeekFrom};
use std::env;

mod encoder;
mod reorder;
mod cpu_data;

fn usage(args: &Vec<String>) {
        println!("Usage: {} -i shellcode.bin -o shellcode.out.bin", args[0]);
        std::process::exit(1);
}

fn main() {

    let number_of_loops = 1;

    // Parsing the Arguments
    let args: Vec<String> = env::args().collect();

    // variable if code should be reordered
    let mut reorder = false;
    if args.contains(&"--reorder".to_string()) {
        reorder = true;
    }

    if args.len() < 5 {
        usage(&args);
    }

    if args[1] != "-i" || args[3] != "-o" {
        usage(&args);
    }

    // Get the Input/Output File-Paths
    let input_file = &args[2];
    let output_file = &args[4];


    let mut shellcode = match File::open(input_file) {
        Ok(v) => v,
        Err(_) => {
            println!("[-] Input File not found, exiting");
            std::process::exit(1);
        }
    };


    // Obtain the Size of the input
    let size = shellcode.seek(SeekFrom::End(0)).unwrap();
    println!("[+] Reading 0x{:x} Bytes from {}", size, input_file);
    shellcode.seek(SeekFrom::Start(0)).unwrap();

    // Input Buffer for the shellcode
    let mut payload = vec![0; size as usize];
    shellcode.read(&mut payload).unwrap();

    let prog = reorder::AsmProgram::new(&payload);

    if reorder {
         payload = prog.reorder();
    }

    // The Encoding Phase
    for i in 0..number_of_loops {

        println!("[*] Encoding Round {}", i);

        // Generate random key
        let key: u8 = rand::random();

        // encode the payload
        let mut encoded_payload = encoder::encode(payload.clone(), key);

        //println!("[*] Generating Stub...");
        let mut stub = encoder::get_decoder_stub(key, encoded_payload.len());

        // Append the encoded_payload to the stub
        stub.append(&mut encoded_payload);

        // Set the [stub+payload] as the new raw payload
        payload = stub.clone();
    }

    //Write Payload to File:
    let mut file = match File::create(output_file) {
        Ok(v) => v,
        Err(_) => {
            println!("[-] Cannot open Output File, outputting hex encoded payload:");
            println!("{}", hex::encode(payload));
            std::process::exit(1);
        }
    };
    file.write(&payload).unwrap();
    println!("[+] Done, Generated encoded payload @ {}", output_file);

}

