ExploitBuilder 0.1.0

use crate::CPU::Endian;
use crate::CPU::Mips32::{Mips32Instruction, Mips32RTypeInstruction, Mips32ITypeInstruction, Mips32Register, Mips32InstructionOpcode, Mips32InstructionFunctionCode};


// Legends:
// (C)[x](S):
// 	- S: simple, true mips instruction
// 	- PD: pseudo, 1-1 map to mips instruction
// 	- PD: pseudo, not 1-1 map to mips instruction
// 	- CD: preverse all registers
// 	- CA:preverse registers according to calling abi
// 	- CR: don't preverse registers
// 	- C[x](S): use syscall
#[derive (Debug, Clone)]
pub struct Mips32Shell
{
	endian: Endian,
	instructions: Vec<Mips32Instruction>,
}

impl Mips32Shell
{
	pub fn New (endian: Endian) -> Self
	{
		Self
		{
			endian,
			instructions: Vec::new (),
		}
	}

	pub fn Lui (&mut self, rt: Mips32Register, number: u16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::LUI, Mips32Register::ZERO, rt, number)));
	}

	pub fn Ori (&mut self, rt: Mips32Register, rs: Mips32Register, number: u16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::ORI, rs, rt, number)));
	}

	pub fn Or (&mut self, rd: Mips32Register, rt: Mips32Register, rs: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::OR)));
	}

	pub fn Li (&mut self, register: Mips32Register, number: u32)
	{
		let upperHalf = ((number&0xffff_0000)>>16) as u16;
		let lowerHalf = (number&0x0000_ffff) as u16;
		self.Lui (register, upperHalf);
		self.Ori (register, register, lowerHalf);
	}

	pub fn Add (&mut self, rd: Mips32Register, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::ADD)));
	}

	pub fn Addi (&mut self, rt: Mips32Register, rs: Mips32Register, imm: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::ADDI, rs, rt, imm as u16)))
	}

	pub fn Addiu (&mut self, rt: Mips32Register, rs: Mips32Register, imm: u16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::ADDIU, rs, rt, imm)))
	}

	pub fn Sub (&mut self, rd: Mips32Register, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::SUB)));
	}

	pub fn Subu (&mut self, rd: Mips32Register, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::SUBU)));
	}

	pub fn Subi (&mut self, rt: Mips32Register, rs: Mips32Register, imm: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::ADDI, rs, rt, -imm as u16)))
	}

	pub fn Xori (&mut self, rt: Mips32Register, rs: Mips32Register, imm: u16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::XORI, rs, rt, imm as u16)))
	}

	pub fn Mult (&mut self, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, Mips32Register::ZERO, 0, Mips32InstructionFunctionCode::MULT)));
	}

	pub fn Multu (&mut self, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, Mips32Register::ZERO, 0, Mips32InstructionFunctionCode::MULTU)));
	}

	pub fn Mfhi (&mut self, rd: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, Mips32Register::ZERO, Mips32Register::ZERO, rd, 0, Mips32InstructionFunctionCode::MFHI)));
	}

	pub fn Mflo (&mut self, rd: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, Mips32Register::ZERO, Mips32Register::ZERO, rd, 0, Mips32InstructionFunctionCode::MFLO)));
	}

	pub fn Lb (&mut self, rt: Mips32Register, rs: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::LB, rs, rt, offset as u16)));
	}

	pub fn Lhu (&mut self, rt: Mips32Register, rs: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::LHU, rs, rt, offset as u16)));
	}

	pub fn Lw (&mut self, rt: Mips32Register, rs: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::LW, rs, rt, offset as u16)));
	}

	pub fn Sh (&mut self, rt: Mips32Register, rs: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::SH, rs, rt, offset as u16)));
	}

	pub fn Sw (&mut self, rt: Mips32Register, rs: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::SW, rs, rt, offset as u16)));
	}

	pub fn Slt (&mut self, rd: Mips32Register, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::SLT)));
	}

	pub fn Sltu (&mut self, rd: Mips32Register, rs: Mips32Register, rt: Mips32Register)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, 0, Mips32InstructionFunctionCode::SLTU)));
	}

	pub fn Slti (&mut self, rt: Mips32Register, rs: Mips32Register, imm: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::SLTI, rs, rt, imm as u16)));
	}

	pub fn Sltiu (&mut self, rt: Mips32Register, rs: Mips32Register, imm: u16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::SLTIU, rs, rt, imm)));
	}

	/// Offset is instruction count, aka byte offset>>2
	pub fn Bne (&mut self, rs: Mips32Register, rt: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::BNE, rs, rt, offset as u16)));
	}

	/// Offset is instruction count, aka byte offset>>2
	pub fn Beq (&mut self, rs: Mips32Register, rt: Mips32Register, offset: i16)
	{
		self.instructions.push (Mips32Instruction::IType(Mips32ITypeInstruction::New (Mips32InstructionOpcode::BEQ, rs, rt, offset as u16)));
	}

	pub fn Sync (&mut self)
	{
		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, Mips32Register::T9, Mips32Register::T9, Mips32Register::T9, 0, Mips32InstructionFunctionCode::SYNC)));
	}

	/// Call syscall with provided 20 bits imm value
	/// imm value can be anything, 0x40404 is normally used to avoid null byte
	pub fn Syscall (&mut self, imm: usize)
	{
		// get the fill value
		let imm = imm&0xfffff;
		let rs = ((imm&0b11111_00000_00000_00000)>>15).try_into ().unwrap ();
		let rt = ((imm&0b00000_11111_00000_00000)>>10).try_into ().unwrap ();
		let rd = ((imm&0b00000_00000_11111_00000)>>5).try_into ().unwrap ();
		let shamt = (imm&0b00000_00000_00000_11111).try_into ().unwrap ();

		self.instructions.push (Mips32Instruction::RType(Mips32RTypeInstruction::New (Mips32InstructionOpcode::REG, rs, rt, rd, shamt, Mips32InstructionFunctionCode::SYSCALL)));
	}

	/// Push 32 bits number
	/// Endian depend on shellcode endian
	/// Subtract $sp by 4
	pub fn PushU32 (&mut self, value: u32)
	{
		// save $t0
		self.Subi (Mips32Register::SP, Mips32Register::SP, 8);
		self.Sw (Mips32Register::T0, Mips32Register::SP, 0);

		// load value to $t0 and push $t0 to stack
		self.Li (Mips32Register::T0, value);
		self.Sw (Mips32Register::T0, Mips32Register::SP, 4);
		self.Addi (Mips32Register::SP, Mips32Register::SP, 4);
	}

	/// Push 64 bits number
	/// Endian depend on shellcode endian
	/// Subtract $sp by 8
	pub fn PushU64 (&mut self, value: u64)
	{
		match self.endian
		{
			Endian::BIG => self.PushByteArray (&value.to_be_bytes ()),
			Endian::LITTLE => self.PushByteArray (&value.to_le_bytes ()),
		};
	}

	/// Push 128 bits number
	/// Endian depend on shellcode endian
	/// Subtract $sp by 16
	pub fn PushU128 (&mut self, value: u128)
	{
		match self.endian
		{
			Endian::BIG => self.PushByteArray (&value.to_be_bytes ()),
			Endian::LITTLE => self.PushByteArray (&value.to_le_bytes ()),
		};
	}

	/// Push a byte array to stack
	/// Always pad data array to 8 bytes align with null byte
	/// Return the number of $sp subtracted
	pub fn PushByteArray (&mut self, data: &[u8]) -> usize
	{
		let mut data = data.to_vec ();

		// align data
		let paddingLen = ((data.len () + 8 - 1)/8)*8 - data.len ();
		data.append (&mut vec! [0u8; paddingLen]);

		// alloc space in stack
		self.Subi (Mips32Register::SP, Mips32Register::SP, (8 + data.len ()).try_into ().unwrap ());

		// save $t0
		self.Sw (Mips32Register::T0, Mips32Register::SP, 0);

		for i in (0..data.len ()).step_by (4).rev ()
		{
			let num = match self.endian
			{
				Endian::BIG => u32::from_be_bytes (data[i..i + 4].try_into ().unwrap ()),
				Endian::LITTLE => u32::from_le_bytes (data[i..i + 4].try_into ().unwrap ()),
			};

			self.Li (Mips32Register::T0, num);
			self.Sw (Mips32Register::T0, Mips32Register::SP, (8 + i).try_into ().unwrap ());
		}

		// restore stack and $t0
		self.Lw (Mips32Register::T0, Mips32Register::SP, 0);
		self.Addi (Mips32Register::SP, Mips32Register::SP, 8);

		data.len ()
	}

	/// Push a string to stack, with null byte appended
	pub fn PushString (&mut self, string: &str) -> usize
	{
		let mut tmp = string.as_bytes ().to_vec ();
		tmp.push (0);   // add null byte
		self.PushByteArray (&tmp)
	}

	/// Get a label to use in later branch call
	/// Label is number of instruction so far
	/// Label only valid in one shellcode instance
	pub fn GetLabel (&self) -> usize
	{
		self.instructions.len ()
	}

	/// Branch to label if rs is equa to rt
	pub fn BeqLabel (&mut self, rs: Mips32Register, rt: Mips32Register, label: usize)
	{
		let offset = label as isize - (self.instructions.len () as isize + 1);
		self.Beq (rs, rt, offset as i16);
	}

	/// Branch to label if rs is not equa to rt
	pub fn BneLabel (&mut self, rs: Mips32Register, rt: Mips32Register, label: usize)
	{
		let offset = label as isize - (self.instructions.len () as isize + 1);
		self.Bne (rs, rt, offset as i16);
	}

	/// Return the shellcode as byte array
	pub fn ToByteArray (&self) -> Vec<u8>
	{
		let mut result = Vec::new ();

		for instruction in &self.instructions
		{
			result.append (&mut instruction.ToByteArray (self.endian));
		}

		result
	}

	/// Return the assembly source of this shellcode
	pub fn ToAssemblySource (&self) -> String
	{
		let mut result = String::new ();

		for instruction in &self.instructions
		{
			result.push_str (&instruction.ToAssemblySource ());
		}

		result
	}

	// open file
	// create file
	// connect
	// sh
	// back connect
	// execve
	// read from file
	// write to file
	// download file http
	// bind port
}