use crate::cpu_data::{x64, Arch as EncArch};
use keystone::*;

pub struct Encoder {
    arch: EncArch,
    shellcode: Vec<u8>,
    num_rounds: i32,
    reorder: bool,
}

impl Encoder {
    pub fn new(arch: EncArch, payload: Vec<u8>) -> Self {
        Encoder {
            arch: arch,
            shellcode: payload,
            num_rounds: 10,
            reorder: false,
        }
    }

    // main method for encoding the payload
    pub fn encode(&self) -> Vec<u8> {

        let mut payload = self.shellcode.clone();

        for i in 0..self.num_rounds {
            println!("[*] Encoding Round {}", i);

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

            // encode the payload
            let mut encoded_payload = self.encode_payload(payload.clone(), key);

            //println!("[*] Generating Stub...");
            let mut stub = self.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();
        }
        Vec::new()
    }

    /// This function does the plain shellcode encoding (XOR)
    pub fn encode_payload(&self, 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
    }
    /// Generates a decoder stub, each with using different registers
    fn get_decoder_stub(&self, 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 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
}