diff --git a/sleighcraft/Cargo.toml b/sleighcraft/Cargo.toml index 3a2ef18..a3131dd 100644 --- a/sleighcraft/Cargo.toml +++ b/sleighcraft/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "sleighcraft" -version = "0.1.1-dev3" +version = "0.2.0-dev1" authors = ["Anciety "] edition = "2018" description = "Binary Analysis Craft" @@ -14,6 +14,7 @@ repository = "https://github.com/ret2lab/bincraft/" cxx = "1.0" once_cell = "1.6.0" num_enum = "0.5.1" +downcast-rs = "1.2.0" sleighcraft_util_macro = { path = "../sleighcraft_util_macro" } [dependencies.pyo3] diff --git a/sleighcraft/src/cpp/bridge/disasm.cpp b/sleighcraft/src/cpp/bridge/disasm.cpp index 5ae0eb6..6bd2b8f 100644 --- a/sleighcraft/src/cpp/bridge/disasm.cpp +++ b/sleighcraft/src/cpp/bridge/disasm.cpp @@ -45,12 +45,77 @@ void SleighProxy::setSpecFromPath(const rust::Str path,int mode) { this->ctx.setVariableDefault("opsize",mode); // Operand size is 32-bit } -unique_ptr new_sleigh_proxy(RustLoadImage &ld) { - unique_ptr proxy(new SleighProxy(ld)); - return proxy; +unique_ptr new_sleigh_proxy(rust::Box ld) { + return std::make_unique(std::move(ld)); } -void SleighProxy::decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_emit, uint64_t start) { +rust::Box& SleighProxy::get_loader_mut() { + return this->loader->load_image; +} + +const rust::Box& SleighProxy::get_loader() const { + return this->loader->load_image; +} + +void SleighProxy::set_loader(rust::Box ld) { + this->loader = make_unique(std::move(ld)); + this->translator.reset(this->loader.get(), &this->ctx); + this->translator.initialize(storage); +} + +rust::Box& SleighProxy::get_asm_emit_mut() { + return this->asm_emit->assemblyEmit; +} + +rust::Box& SleighProxy::get_pcode_emit_mut() { + return this->pcode_emit->rustPcodeEmit; +} + +const rust::Box& SleighProxy::get_asm_emit() const { + return this->asm_emit->assemblyEmit; +} + +const rust::Box& SleighProxy::get_pcode_emit() const { + return this->pcode_emit->rustPcodeEmit; +} + +void SleighProxy::set_asm_emit(rust::Box asm_emit) { + this->asm_emit = std::make_unique(std::move(asm_emit)); +} + +void SleighProxy::set_pcode_emit(rust::Box pcode_emit) { + this->pcode_emit = std::make_unique(std::move(pcode_emit)); +} + +int32_t SleighProxy::decode_asm_at(uint64_t start) { + Address address(translator.getDefaultCodeSpace(), start); + + try { + auto length = translator.printAssembly(*asm_emit, address); + return length; + } catch (BadDataError& e) { + throw std::invalid_argument("bad data when decode asm: " + e.explain); + } catch (UnimplError& e) { + throw std::logic_error("pcode not implemented"); + } + + return 0; +} + +void SleighProxy::decode_pcode_at(uint64_t start) { + Address address(translator.getDefaultCodeSpace(), start); + + try { + translator.oneInstruction(*pcode_emit, address); + } catch (BadDataError& e) { + throw std::invalid_argument("bad data when decode pcode: " + e.explain); + } catch (UnimplError& e) { + throw std::logic_error("pcode not implemented"); + } +} + +/* +void SleighProxy::decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_emit, uint64_t start, uint64_t inst_size) { auto assemblyEmit = RustAssemblyEmitProxy{asm_emit}; auto pcodeEmit = RustPcodeEmitProxy{pcode_emit}; @@ -59,7 +124,8 @@ void SleighProxy::decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_e auto length = 0; auto buf_used = 0; - auto buf_size = loader.bufSize(); + auto buf_size = loader->bufSize(); + auto total_insts = 0; while (buf_used < buf_size) { try { @@ -67,6 +133,10 @@ void SleighProxy::decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_e translator.oneInstruction(pcodeEmit, address); address = address + length; buf_used = buf_used + length; + total_insts ++; + if (inst_size > 0 && total_insts >= inst_size) { + break; + } } catch (BadDataError &e) { throw std::invalid_argument("BadDataError"); @@ -79,6 +149,7 @@ void SleighProxy::decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_e } } +*/ // RustLoadImageProxy void RustLoadImageProxy::loadFill(uint1 *ptr, int4 size, const Address &address) { @@ -86,11 +157,11 @@ void RustLoadImageProxy::loadFill(uint1 *ptr, int4 size, const Address &address) uint8_t* array = (uint8_t*)ptr; rust::Slice<::std::uint8_t> slice{array,(unsigned long)size}; const auto addr_proxy = AddressProxy{addr}; - load_image.load_fill(slice, addr_proxy); + load_image->load_fill(slice, addr_proxy); } void RustLoadImageProxy::adjustVma(long adjust) { - this->load_image.adjust_vma(adjust); + this->load_image->adjust_vma(adjust); } string RustLoadImageProxy::getArchType(void) const { @@ -98,17 +169,17 @@ string RustLoadImageProxy::getArchType(void) const { } int4 RustLoadImageProxy::bufSize() { - return load_image.buf_size(); + return load_image->buf_size(); } -std::unique_ptr from_rust(RustLoadImage& load_image) { - return unique_ptr(new RustLoadImageProxy(load_image)); +std::unique_ptr from_rust(rust::Box load_image) { + return std::make_unique(std::move(load_image)); } void RustAssemblyEmitProxy::dump(const Address &address, const string &mnemonic, const string &body) { Address addr = const_cast (address); const auto addr_proxy = AddressProxy{addr}; - assemblyEmit.dump(addr_proxy, mnemonic, body); + assemblyEmit->dump(addr_proxy, mnemonic, body); } @@ -124,5 +195,5 @@ void RustPcodeEmitProxy::dump(const Address &addr, OpCode opc, VarnodeData *outv vars_vec.push_back(VarnodeDataProxy{&vars[i]}); } - rustPcodeEmit.dump(addr_proxy, opcodes, outvar_proxy, vars_vec); + rustPcodeEmit->dump(addr_proxy, opcodes, outvar_proxy, vars_vec); } diff --git a/sleighcraft/src/cpp/bridge/disasm.h b/sleighcraft/src/cpp/bridge/disasm.h index 3ed173e..b1ab3e9 100644 --- a/sleighcraft/src/cpp/bridge/disasm.h +++ b/sleighcraft/src/cpp/bridge/disasm.h @@ -30,11 +30,8 @@ class OpCodeProxy; class VarnodeDataProxy; class RustLoadImageProxy: public LoadImage { public: - RustLoadImage& load_image; -// string(load_image->get_filename()) - RustLoadImageProxy(RustLoadImage &load_image): load_image(load_image), LoadImage("nofile") {} - - RustLoadImageProxy(RustLoadImage *load_image): load_image(*load_image), LoadImage("nofile") {} + rust::Box load_image; + RustLoadImageProxy(rust::Box load_image): load_image(std::move(load_image)), LoadImage("nofile") {} virtual void loadFill(uint1 *ptr, int4 size, const Address &address); @@ -72,9 +69,8 @@ struct InstructionProxy { class RustAssemblyEmitProxy: public AssemblyEmit { public: - RustAssemblyEmit& assemblyEmit; - RustAssemblyEmitProxy(RustAssemblyEmit& assemblyEmit): assemblyEmit(assemblyEmit){} - RustAssemblyEmitProxy(RustAssemblyEmit* assemblyEmit): assemblyEmit(*assemblyEmit){} + rust::Box assemblyEmit; + RustAssemblyEmitProxy(rust::Box assemblyEmit): assemblyEmit(std::move(assemblyEmit)){} virtual void dump(const Address &address, const string &mnemonic, const string &body); @@ -84,9 +80,8 @@ class RustAssemblyEmitProxy: public AssemblyEmit { class RustPcodeEmitProxy: public PcodeEmit { public: - RustPcodeEmit& rustPcodeEmit; - RustPcodeEmitProxy(RustPcodeEmit& rustPcodeEmit): rustPcodeEmit(rustPcodeEmit){} - RustPcodeEmitProxy(RustPcodeEmit* rustPcodeEmit): rustPcodeEmit(*rustPcodeEmit){} + rust::Box rustPcodeEmit; + RustPcodeEmitProxy(rust::Box rustPcodeEmit): rustPcodeEmit(std::move(rustPcodeEmit)){} virtual void dump(const Address &addr,OpCode opc,VarnodeData *outvar,VarnodeData *vars,int4 isize); @@ -94,14 +89,29 @@ class RustPcodeEmitProxy: public PcodeEmit { class SleighProxy { public: - SleighProxy(RustLoadImage &ld): loader(ld), translator(&loader, &this->ctx) {} + SleighProxy(rust::Box ld): loader(std::make_unique(std::move(ld))), translator(loader.get(), &this->ctx) {} + + void set_asm_emit(rust::Box asm_emit); + void set_pcode_emit(rust::Box pcode_emit); + rust::Box& get_asm_emit_mut(); + rust::Box& get_pcode_emit_mut(); + const rust::Box& get_asm_emit() const; + const rust::Box& get_pcode_emit() const; void setSpecFromPath(const rust::Str path, int mode); void set_spec(const rust::Str spec_content, int mode); - void decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_emit, uint64_t start); + int32_t decode_asm_at(uint64_t start); + void decode_pcode_at(uint64_t start); + //void decode_with(RustAssemblyEmit& asm_emit, RustPcodeEmit& pcode_emit, uint64_t start, uint64_t inst_size); + void set_loader(rust::Box ld); + rust::Box& get_loader_mut(); + const rust::Box& get_loader() const; private: - RustLoadImageProxy loader; + std::unique_ptr loader; + std::unique_ptr asm_emit; + std::unique_ptr pcode_emit; + Sleigh translator; ContextInternal ctx; DocumentStorage storage; @@ -109,7 +119,7 @@ class SleighProxy { //unique_ptr proxy_from_spec(rust::Str path, RustLoadImage &ld, RustAssemblyEmit &asm_emit, RustPcodeEmit &rustPcodeEmit); //unique_ptr proxy_from_spec_path(rust::Str spec_content, RustLoadImage &ld, RustAssemblyEmit &asm_emit, RustPcodeEmit &rustPcodeEmit); -std::unique_ptr from_rust(RustLoadImage& load_image); -unique_ptr new_sleigh_proxy(RustLoadImage &ld); +std::unique_ptr from_rust(rust::Box load_image); +unique_ptr new_sleigh_proxy(rust::Box ld); #endif \ No newline at end of file diff --git a/sleighcraft/src/prelude.rs b/sleighcraft/src/prelude.rs index ebe3578..d98e981 100644 --- a/sleighcraft/src/prelude.rs +++ b/sleighcraft/src/prelude.rs @@ -13,4 +13,4 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub use crate::{arch, CollectingAssemblyEmit, CollectingPcodeEmit, PlainLoadImage, SleighBuilder}; +pub use crate::{CollectingAssemblyEmit, CollectingPcodeEmit, PlainLoadImage, Sleigh}; \ No newline at end of file diff --git a/sleighcraft/src/sleigh.rs b/sleighcraft/src/sleigh.rs index b65e15c..d416277 100644 --- a/sleighcraft/src/sleigh.rs +++ b/sleighcraft/src/sleigh.rs @@ -17,7 +17,8 @@ use crate::error::{Error, Result}; use cxx::{CxxString, UniquePtr}; use once_cell::sync::Lazy; use sleighcraft_util_macro::def_sla_load_preset; -use std::collections::HashMap; +use std::any::Any; +use std::collections::{BTreeSet, HashMap}; #[cxx::bridge] pub mod ffi { @@ -188,14 +189,14 @@ pub mod ffi { } extern "Rust" { - type RustAssemblyEmit<'a>; + type RustAssemblyEmit; fn dump( self: &mut RustAssemblyEmit, address: &AddressProxy, mnem: &CxxString, body: &CxxString, ); - type RustPcodeEmit<'a>; + type RustPcodeEmit; fn dump( self: &mut RustPcodeEmit, address: &AddressProxy, @@ -203,7 +204,7 @@ pub mod ffi { outvar: Pin<&mut VarnodeDataProxy>, vars: &CxxVector, ); - type RustLoadImage<'a>; + type RustLoadImage; fn load_fill(self: &mut RustLoadImage, ptr: &mut [u8], addr: &AddressProxy); //fn get_arch_type(self: &RustLoadImage) -> String; fn adjust_vma(self: &mut RustLoadImage, adjust: isize); @@ -379,7 +380,7 @@ pub mod ffi { type RustLoadImageProxy; - fn from_rust(load_iamge: &mut RustLoadImage) -> UniquePtr; + fn from_rust(load_iamge: Box) -> UniquePtr; // type InstructionProxy; // @@ -390,20 +391,24 @@ pub mod ffi { type SleighProxy; fn set_spec(self: Pin<&mut SleighProxy>, spec_content: &str, mode: i32); - fn new_sleigh_proxy(ld: &mut RustLoadImage) -> UniquePtr; - fn decode_with( - self: Pin<&mut SleighProxy>, - asm_emit: &mut RustAssemblyEmit, - pcode_emit: &mut RustPcodeEmit, - start: u64, - ) -> Result<()>; + fn new_sleigh_proxy(ld: Box) -> UniquePtr; + fn decode_asm_at(self: Pin<&mut SleighProxy>, start: u64) -> Result; + fn decode_pcode_at(self: Pin<&mut SleighProxy>, start: u64) -> Result<()>; + fn set_loader(self: Pin<&mut SleighProxy>, load: Box); + fn set_asm_emit(self: Pin<&mut SleighProxy>, asm_emit: Box); + fn set_pcode_emit(self: Pin<&mut SleighProxy>, pcode_emit: Box); + fn get_loader_mut(self: Pin<&mut SleighProxy>) -> &mut Box; + fn get_asm_emit_mut(self: Pin<&mut SleighProxy>) -> &mut Box; + fn get_pcode_emit_mut(self: Pin<&mut SleighProxy>) -> &mut Box; + + fn get_loader(self: &SleighProxy) -> &Box; + fn get_asm_emit(self: &SleighProxy) -> &Box; + fn get_pcode_emit(self: &SleighProxy) -> &Box; } } -use crate::Mode::MODE16; use ffi::*; use num_enum::TryFromPrimitive; -use std::borrow::BorrowMut; use std::pin::Pin; impl ToString for PcodeOpCode { @@ -497,15 +502,23 @@ pub enum Mode { MODE64 = 2, } +impl Default for Mode { + fn default() -> Self { + Self::MODE16 + } +} + pub trait AssemblyEmit { fn dump(&mut self, addr: &AddressProxy, mnem: &str, body: &str); + fn as_any(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; } -pub struct RustAssemblyEmit<'a> { - internal: &'a mut dyn AssemblyEmit, +pub struct RustAssemblyEmit { + internal: Box, } -impl<'a> RustAssemblyEmit<'a> { +impl RustAssemblyEmit { pub fn dump(&mut self, address: &AddressProxy, mnem: &CxxString, body: &CxxString) { let mnem = mnem.to_str().unwrap(); let body = body.to_str().unwrap(); @@ -513,9 +526,17 @@ impl<'a> RustAssemblyEmit<'a> { self.internal.dump(address, mnem, body); } - pub fn from_internal(internal: &'a mut dyn AssemblyEmit) -> Self { + pub fn from_internal(internal: Box) -> Self { Self { internal } } + + pub fn internal_ref(&self) -> &Box { + &self.internal + } + + pub fn internal_mut(&mut self) -> &mut Box { + &mut self.internal + } } #[derive(Debug, Default)] @@ -534,8 +555,93 @@ impl AssemblyEmit for CollectingAssemblyEmit { }; self.asms.push(asm) } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + + +impl CollectingAssemblyEmit { + pub fn correspond<'a>(&'a mut self, pcodes: &'a CollectingPcodeEmit) -> CorrespondedCollectedAsm<'a> { + let mut idx_map = HashMap::new(); + for asm_idx in 0..self.asms.len() { + let asm = &self.asms[asm_idx]; + let mut pcode_indices = BTreeSet::new(); + for i in 0..pcodes.pcode_asms.len() { + let addr = &pcodes.pcode_asms[i].addr; + if pcodes.pcode_asms[i].addr == asm.addr { + pcode_indices.insert(PcodeAddrIndex { + pcode_vec_idx: i, + offset: addr.offset, + seq: pcodes.pcode_asms[i].seq + }); + } + } + + idx_map.insert(asm_idx, pcode_indices); + } + + CorrespondedCollectedAsm { + asm_emit: self, + pcode_emit: pcodes, + idx_map + } + } } +#[derive(Clone, Copy, PartialEq, Eq)] +struct PcodeAddrIndex { + /// index into pcode table + pcode_vec_idx: usize, + /// addr sequence, to properly order the pcodes + offset: u64, + seq: u64 +} + +impl PartialOrd for PcodeAddrIndex { + fn partial_cmp(&self, other: &Self) -> Option { + if self.offset == other.offset { + self.seq.partial_cmp(&other.seq) + } else { + self.offset.partial_cmp(&other.offset) + } + } +} + +impl Ord for PcodeAddrIndex { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.partial_cmp(other).unwrap() + } +} + +#[derive(Clone)] +pub struct CorrespondedCollectedAsm<'a> { + asm_emit: &'a CollectingAssemblyEmit, + pcode_emit: &'a CollectingPcodeEmit, + /// asm => range of inst in pcode table. + idx_map: HashMap>, +} + +impl<'a> CorrespondedCollectedAsm<'a> { + pub fn asms(&self) -> &Vec { + &self.asm_emit.asms + } + + pub fn pcodes(&self) -> &Vec { + &self.pcode_emit.pcode_asms + } + + pub fn pcode_indices_of_asm(&self, asm_idx: usize) -> BTreeSet { + self.idx_map[&asm_idx].iter().map(|i| i.pcode_vec_idx).collect() + } +} + + #[derive(Debug)] pub struct PcodeVarnodeData { pub space: String, @@ -560,6 +666,7 @@ impl PcodeVarnodeData { #[derive(Debug)] pub struct PcodeInstruction { pub addr: Address, + pub seq: u64, pub opcode: PcodeOpCode, pub vars: Vec, pub out_var: Option, @@ -578,11 +685,15 @@ pub trait PcodeEmit { outvar: Option<&VarnodeDataProxy>, vars: &[&VarnodeDataProxy], ); + fn as_any(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; } #[derive(Debug, Default)] pub struct CollectingPcodeEmit { pub pcode_asms: Vec, + seq: u64, + last_offset: u64, } impl PcodeEmit for CollectingPcodeEmit { @@ -593,10 +704,6 @@ impl PcodeEmit for CollectingPcodeEmit { outvar: Option<&VarnodeDataProxy>, vars: &[&VarnodeDataProxy], ) { - //let space = String::from(outvar.get_space().get_name().to_str().unwrap()); - //let offset = outvar.get_offset() as u64; - // let data = format!("{}{}{}{}{:x}{}{}{}{}", "(", space, ",", "0x", of, ",", size, ")", "="); - let space = String::from(addr.get_space().get_name().to_str().unwrap()); let offset = addr.get_offset() as u64; let mut pcode_vars = vec![]; @@ -608,24 +715,48 @@ impl PcodeEmit for CollectingPcodeEmit { } else { None }; + + if self.last_offset != offset { + self.seq = 0; + } + self.pcode_asms.push(PcodeInstruction { addr: Address { space, offset }, opcode: opcode, vars: pcode_vars, + seq: self.seq, out_var, }); + + self.seq += 1; + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self } } -pub struct RustPcodeEmit<'a> { - pub internal: &'a mut dyn PcodeEmit, +pub struct RustPcodeEmit { + pub internal: Box, } -impl<'a> RustPcodeEmit<'a> { - pub fn from_internal(internal: &'a mut dyn PcodeEmit) -> Self { +impl RustPcodeEmit { + pub fn from_internal(internal: Box) -> Self { Self { internal } } + pub fn internal_ref(&self) -> &Box { + &self.internal + } + + pub fn internal_mut(&mut self) -> &mut Box { + &mut self.internal + } + pub fn dump( &mut self, address: &AddressProxy, @@ -654,15 +785,23 @@ pub trait LoadImage { fn buf_size(&mut self) -> usize; } -pub struct RustLoadImage<'a> { - internal: &'a mut dyn LoadImage, +pub struct RustLoadImage { + internal: Box, } -impl<'a> RustLoadImage<'a> { - pub fn from_internal(internal: &'a mut dyn LoadImage) -> Self { +impl RustLoadImage { + pub fn from_internal(internal: Box) -> Self { Self { internal } } + pub fn internal_ref(&self) -> &Box { + &self.internal + } + + pub fn internal_mut(&mut self) -> &mut Box { + &mut self.internal + } + pub fn load_fill(&mut self, ptr: &mut [u8], addr: &AddressProxy) { self.internal.load_fill(ptr, addr) } @@ -716,6 +855,23 @@ pub struct Address { pub offset: u64, } +impl PartialOrd for Address { + fn partial_cmp(&self, other: &Self) -> Option { + if self.space == other.space { + self.offset.partial_cmp(&other.offset) + } else { + None + } + } +} + +impl Ord for Address { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.partial_cmp(other) + .expect("incorrect address: not comparable") + } +} + #[derive(Debug)] pub struct Instruction { pub addr: Address, @@ -743,96 +899,246 @@ def_sla_load_preset!("sleighcraft/sla/", fn load_preset() -> HashMap<&'static st const PRESET: Lazy> = Lazy::new(|| load_preset()); -pub struct Sleigh<'a> { - sleigh_proxy: UniquePtr, - asm_emit: RustAssemblyEmit<'a>, - pcode_emit: RustPcodeEmit<'a>, - _load_image: Pin>>, +pub struct SleighSetting { + asm_emit: Box, + pcode_emit: Box, + loader: Box, + spec: String, + mode: Mode, } -impl<'a> Sleigh<'a> { - pub fn decode(&mut self, start: u64) -> Result<()> { - // self.load_image.set_buf(bytes); - let assembly_emit = self.asm_emit.borrow_mut(); - let pcodes_emit = self.pcode_emit.borrow_mut(); - self.sleigh_proxy - .as_mut() - .unwrap() - .decode_with(assembly_emit, pcodes_emit, start) - .map_err(|e| Error::CppException(e)) +impl SleighSetting { + pub fn new(loader: Box) -> Self { + let asm_emit = Box::new(CollectingAssemblyEmit::default()); + let pcode_emit = Box::new(CollectingPcodeEmit::default()); + let spec = arch_spec("x86-64") + .expect("unable to find default arch: x86-64") + .to_string(); + let mode = Mode::default(); + + Self { + asm_emit, + pcode_emit, + loader, + spec, + mode, + } } -} -#[derive(Default)] -pub struct SleighBuilder<'a> { - asm_emit: Option>, - pcode_emit: Option>, - load_image: Option>, - spec: Option, - mode: Option, -} -impl<'a> SleighBuilder<'a> { - // TODO: add from_arch(arch_name: &str) -> Self helper function. + pub fn with_spec(spec: &str, loader: Box) -> Self { + let mut setting = SleighSetting::new(loader); + setting.spec(spec); + setting + } - pub fn asm_emit(&mut self, asm_emit: &'a mut dyn AssemblyEmit) -> &mut Self { - self.asm_emit = Some(RustAssemblyEmit::from_internal(asm_emit)); - self + pub fn with_arch(arch: &str, loader: Box) -> Result { + let mut setting = SleighSetting::new(loader); + setting.arch(arch)?; + Ok(setting) } - pub fn pcode_emit(&mut self, pcode_emit: &'a mut dyn PcodeEmit) -> &mut Self { - self.pcode_emit = Some(RustPcodeEmit::from_internal(pcode_emit)); + pub fn asm_emit(&mut self, emit: Box) -> &mut Self { + self.asm_emit = emit; self } - pub fn mode(&mut self, mode: Mode) -> &mut Self { - self.mode = Some(mode); + pub fn pcode_emit(&mut self, emit: Box) -> &mut Self { + self.pcode_emit = emit; self } - pub fn spec(&mut self, spec: &str) -> &mut Self { - // self.load_image = unsafe{Some(Box::from_raw(loader))}; - self.spec = Some(spec.to_string()); + pub fn spec(&mut self, spec_str: &str) -> &mut Self { + self.spec = spec_str.to_string(); self } - pub fn loader(&mut self, loader: &'a mut dyn LoadImage) -> &mut Self { - self.load_image = Some(RustLoadImage::from_internal(loader)); + pub fn arch(&mut self, arch_name: &str) -> Result<&mut Self> { + self.spec = arch_spec(arch_name)?.to_string(); + Ok(self) + } + + pub fn mode(&mut self, mode: Mode) -> &mut Self { + self.mode = mode; self } +} - pub fn try_build(mut self) -> Result> { - let load_image = self - .load_image - .ok_or(Error::MissingArg("load_image".to_string()))?; - let mut load_image = Box::pin(load_image); - let mut sleigh_proxy = new_sleigh_proxy(&mut load_image); +pub struct Sleigh { + sleigh_proxy: UniquePtr, +} - let spec = self.spec.ok_or(Error::MissingArg("spec".to_string()))?; - if self.mode.is_none() { - // Set default address and Operand size - self.mode = Some(MODE16); - }; +impl Sleigh { + pub fn new(loader: Box) -> Result { + let setting = SleighSetting::new(loader); + Self::with_setting(setting) + } + + pub fn with_setting(setting: SleighSetting) -> Result { + let SleighSetting { + asm_emit, + pcode_emit, + loader, + spec, + mode, + } = setting; + + let asm_emit = Box::new(RustAssemblyEmit::from_internal(asm_emit)); + let pcode_emit = Box::new(RustPcodeEmit::from_internal(pcode_emit)); + let loader = Box::new(RustLoadImage::from_internal(loader)); + + let mut sleigh_proxy = new_sleigh_proxy(loader); sleigh_proxy .as_mut() .unwrap() - .set_spec(spec.as_str(), self.mode.unwrap() as i32); + .set_spec(spec.as_str(), mode as i32); - let asm_emit = self - .asm_emit - .ok_or(Error::MissingArg("asm_emit".to_string()))?; - let pcode_emit = self - .pcode_emit - .ok_or(Error::MissingArg("pcode_emit".to_string()))?; + sleigh_proxy.as_mut().unwrap().set_asm_emit(asm_emit); - Ok(Sleigh { - sleigh_proxy, - asm_emit, - pcode_emit, - _load_image: load_image, - }) + sleigh_proxy.as_mut().unwrap().set_pcode_emit(pcode_emit); + + Ok(Self { sleigh_proxy }) + } + + pub fn with_arch(arch_name: &str, loader: Box) -> Result { + let mut setting = SleighSetting::new(loader); + setting.arch(arch_name)?; + Self::with_setting(setting) + } + + pub fn with_spec(spec: &str, loader: Box) -> Result { + let mut setting = SleighSetting::new(loader); + setting.spec(spec); + Self::with_setting(setting) + } + + pub fn decode_asm_at(&mut self, addr: u64) -> Result { + self.sleigh_proxy + .as_mut() + .unwrap() + .decode_asm_at(addr) + .map_err(|e| Error::CppException(e)) + } + + pub fn decode_pcode_at(&mut self, addr: u64) -> Result<()> { + self.sleigh_proxy + .as_mut() + .unwrap() + .decode_pcode_at(addr) + .map_err(|e| Error::CppException(e)) + } + + pub fn decode_at(&mut self, addr: u64) -> Result { + let length = self.decode_asm_at(addr)?; + self.decode_pcode_at(addr)?; + Ok(length) + } + + pub fn decode(&mut self, mut start: u64, inst_size: Option) -> Result<()> { + let inst_size = if inst_size.is_none() { + 0 + } else { + inst_size.unwrap() + }; + let buf_size = self.load_image_mut().buf_size(); + let mut buf_used = 0; + let mut total_inst = 0; + + while buf_used < buf_size { + let length = self.decode_at(start)?; + assert!(length >= 0); + start += length as u64; + buf_used += length as usize; + total_inst += 1; + if inst_size > 0 && total_inst >= inst_size { + break; + } + } + + Ok(()) + } + + pub fn set_asm_emit(&mut self, asm_emit: Box) { + let asm_emit = Box::new(RustAssemblyEmit::from_internal(asm_emit)); + self.sleigh_proxy.as_mut().unwrap().set_asm_emit(asm_emit); + } + + pub fn set_pcode_emit(&mut self, pcode_emit: Box) { + let pcode_emit = Box::new(RustPcodeEmit::from_internal(pcode_emit)); + self.sleigh_proxy + .as_mut() + .unwrap() + .set_pcode_emit(pcode_emit); + } + + pub fn set_loader(&mut self, loader: Box) { + let load_image = Box::new(RustLoadImage::from_internal(loader)); + self.sleigh_proxy.as_mut().unwrap().set_loader(load_image); + } + + pub fn asm_emit(&self) -> &Box { + self.sleigh_proxy + .as_ref() + .unwrap() + .get_asm_emit() + .as_ref() + .internal_ref() + } + + pub fn asm_emit_mut(&mut self) -> &mut Box { + self.sleigh_proxy + .as_mut() + .unwrap() + .get_asm_emit_mut() + .as_mut() + .internal_mut() + } + + pub fn pcode_emit(&self) -> &Box { + self.sleigh_proxy + .as_ref() + .unwrap() + .get_pcode_emit() + .as_ref() + .internal_ref() + } + + pub fn pcode_emit_mut(&mut self) -> &mut Box { + self.sleigh_proxy + .as_mut() + .unwrap() + .get_pcode_emit_mut() + .as_mut() + .internal_mut() + } + + pub fn loader(&self) -> &Box { + self.load_image() + } + + pub fn loader_mut(&mut self) -> &mut Box { + self.load_image_mut() + } + + pub fn load_image(&self) -> &Box { + self.sleigh_proxy + .as_ref() + .unwrap() + .get_loader() + .as_ref() + .internal_ref() + } + + pub fn load_image_mut(&mut self) -> &mut Box { + self.sleigh_proxy + .as_mut() + .unwrap() + .get_loader_mut() + .as_mut() + .internal_mut() } } -pub fn arch(name: &str) -> Result<&str> { + +pub fn arch_spec(name: &str) -> Result<&str> { let content = *PRESET .get(&name.to_lowercase().as_str()) .ok_or(Error::ArchNotFound(name.to_string()))?; diff --git a/sleighcraft/tests/mod.rs b/sleighcraft/tests/mod.rs index afede42..04c9a17 100644 --- a/sleighcraft/tests/mod.rs +++ b/sleighcraft/tests/mod.rs @@ -1,129 +1,143 @@ +use sleighcraft::SleighSetting; use sleighcraft::prelude::*; -use sleighcraft::Mode::{MODE32, MODE64}; +use sleighcraft::Mode; -// #[test] -// fn test_custom_spec() { -// // let compiled = include_str!("../test/test.sla"); -// // let mut sleigh = Sleigh::from_spec(compiled).unwrap(); -// // let buf = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6]; -// // sleigh.decode(0, &buf, 1); -// // println!("{:?}", sleigh.pcode_emit) -// } +#[test] +fn test_custom_spec() { + let compiled = include_str!("test.sla"); + + let buf = [0x20, 0x30]; + + let loader = PlainLoadImage::from_buf(&buf, 0); + + let setting = SleighSetting::with_spec(compiled, Box::new(loader)); + let mut sleigh = Sleigh::with_setting(setting).unwrap(); + + sleigh.decode(0, None).unwrap(); + + let asm = sleigh + .asm_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + let pcode = sleigh + .pcode_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(asm.asms.len(), 1); + assert!(pcode.pcode_asms.len() > 0); +} #[test] fn test_x86() { - let mut sleigh_builder = SleighBuilder::default(); - let spec = arch("x86").unwrap(); let buf = [0x90, 0x32, 0x31]; - let mut loader = PlainLoadImage::from_buf(&buf, 0); - sleigh_builder.loader(&mut loader); - sleigh_builder.spec(spec); - let mut asm_emit = CollectingAssemblyEmit::default(); - let mut pcode_emit = CollectingPcodeEmit::default(); - sleigh_builder.asm_emit(&mut asm_emit); - sleigh_builder.pcode_emit(&mut pcode_emit); - let mut sleigh = sleigh_builder.try_build().unwrap(); - - sleigh.decode(0).unwrap(); - - println!("{:?}", asm_emit.asms); - println!("{:?}", pcode_emit.pcode_asms); + let loader = Box::new(PlainLoadImage::from_buf(&buf, 0)); + + let mut setting = SleighSetting::with_arch("x86", loader).unwrap(); + setting.mode(Mode::MODE32); + let mut sleigh = Sleigh::with_setting(setting).unwrap(); + + sleigh.decode(0, None).unwrap(); + + let asm = sleigh + .asm_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + let pcode = sleigh + .pcode_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(asm.asms.len() == 2); + assert!(pcode.pcode_asms.len() > 0); } #[test] fn test_x86_case_ignoring() { - let mut sleigh_builder = SleighBuilder::default(); - let spec = arch("x86").unwrap(); let buf = [0x90, 0x32, 0x31]; - let mut loader = PlainLoadImage::from_buf(&buf, 0); - sleigh_builder.loader(&mut loader); - sleigh_builder.spec(spec); - let mut asm_emit = CollectingAssemblyEmit::default(); - let mut pcode_emit = CollectingPcodeEmit::default(); - sleigh_builder.asm_emit(&mut asm_emit); - sleigh_builder.pcode_emit(&mut pcode_emit); - let mut sleigh = sleigh_builder.try_build().unwrap(); - - sleigh.decode(0).unwrap(); - - println!("{:?}", asm_emit.asms); - println!("{:?}", pcode_emit.pcode_asms); -} + let loader = Box::new(PlainLoadImage::from_buf(&buf, 0)); -#[test] -fn test_x86_32_bit() { - let mut sleigh_builder = SleighBuilder::default(); - let spec = arch("x86").unwrap(); - let buf = [0x90, 0x32, 0x31]; - let mut loader = PlainLoadImage::from_buf(&buf, 0); - sleigh_builder.loader(&mut loader); - sleigh_builder.spec(spec); - sleigh_builder.mode(MODE32); - let mut asm_emit = CollectingAssemblyEmit::default(); - let mut pcode_emit = CollectingPcodeEmit::default(); - sleigh_builder.asm_emit(&mut asm_emit); - sleigh_builder.pcode_emit(&mut pcode_emit); - let mut sleigh = sleigh_builder.try_build().unwrap(); - - sleigh.decode(0).unwrap(); - - println!("{:?}", asm_emit.asms); - println!("{:?}", pcode_emit.pcode_asms); + let mut setting = SleighSetting::with_arch("X86", loader).unwrap(); + setting.mode(Mode::MODE32); + let mut sleigh = Sleigh::with_setting(setting).unwrap(); + + sleigh.decode(0, None).unwrap(); + + let asm = sleigh + .asm_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + let pcode = sleigh + .pcode_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(asm.asms.len() == 2); + assert!(pcode.pcode_asms.len() > 0); } + #[test] fn test_x86_64_bit() { - let mut sleigh_builder = SleighBuilder::default(); - let spec = arch("x86-64").unwrap(); - let buf = [72, 49, 192]; - let mut loader = PlainLoadImage::from_buf(&buf, 0); - sleigh_builder.loader(&mut loader); - sleigh_builder.spec(spec); - sleigh_builder.mode(MODE64); - let mut asm_emit = CollectingAssemblyEmit::default(); - let mut pcode_emit = CollectingPcodeEmit::default(); - sleigh_builder.asm_emit(&mut asm_emit); - sleigh_builder.pcode_emit(&mut pcode_emit); - let mut sleigh = sleigh_builder.try_build().unwrap(); - - sleigh.decode(0).unwrap(); - - println!("{:?}", asm_emit.asms); - println!("{:?}", pcode_emit.pcode_asms); + let buf = [0x90, 0x58]; + let loader = Box::new(PlainLoadImage::from_buf(&buf, 0)); + let mut setting = SleighSetting::with_arch("X86-64", loader).unwrap(); + setting.mode(Mode::MODE64); + let mut sleigh = Sleigh::with_setting(setting).unwrap(); + + sleigh.decode(0, None).unwrap(); + + let asm = sleigh + .asm_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + let pcode = sleigh + .pcode_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(asm.asms.len() == 2); + assert!(pcode.pcode_asms.len() > 0); } + #[test] fn test_mips32le_bit() { - let mut sleigh_builder = SleighBuilder::default(); - let spec = arch("mips32le").unwrap(); /* 0x0 j 8 0x4 add $1, $2, $3 0x8 ori $1, $2, 0x64 */ let buf = [2, 0, 0, 8, 32, 8, 67, 0, 100, 0, 65, 52]; - let mut loader = PlainLoadImage::from_buf(&buf, 0); - sleigh_builder.loader(&mut loader); - sleigh_builder.spec(spec); - let mut asm_emit = CollectingAssemblyEmit::default(); - let mut pcode_emit = CollectingPcodeEmit::default(); - sleigh_builder.asm_emit(&mut asm_emit); - sleigh_builder.pcode_emit(&mut pcode_emit); - let mut sleigh = sleigh_builder.try_build().unwrap(); - - sleigh.decode(0).unwrap(); - - for asm in asm_emit.asms.iter() { - println!("{}:\t{}\t{}", asm.addr.offset, asm.mnemonic, asm.body); - } - - println!(); - - for pcode in pcode_emit.pcode_asms.iter() { - println!("address: {:?}", pcode.addr); - println!("opcode: {:?}", pcode.opcode); - println!("vars: {:?}", pcode.vars); - println!("out_var: {:?}", pcode.out_var); - println!(); - } -} + let loader = Box::new(PlainLoadImage::from_buf(&buf, 0)); + let mut sleigh = Sleigh::with_arch("mips32le", loader).unwrap(); + + sleigh.decode(0, None).unwrap(); + + let asm = sleigh + .asm_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + let pcode = sleigh + .pcode_emit() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(asm.asms.len() > 0); + assert!(pcode.pcode_asms.len() > 0); +} \ No newline at end of file