From ee22fe915fa46f2b7b1be2f68475b762b5deb29c Mon Sep 17 00:00:00 2001 From: Duncan Ogilvie Date: Sun, 4 Jun 2023 14:36:44 +0200 Subject: [PATCH] Add register manipulation functions --- Cargo.lock | 3 ++ Cargo.toml | 4 ++- icicle.pyi | 12 +++++++- src/lib.rs | 75 +++++++++++++++++++++++++++++++++++++++++------- tests/example.py | 4 +++ 5 files changed, 85 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6dbdaa5..e22973b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -361,9 +361,11 @@ version = "0.1.0" dependencies = [ "icicle-cpu", "icicle-vm", + "indexmap", "pcode", "pyo3", "pyo3-build-config", + "sleigh-runtime", "target-lexicon", ] @@ -547,6 +549,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3b1ac5b3731ba34fdaa9785f8d74d17448cd18f30cf19e0c7e7b1fdb5272109" dependencies = [ "cfg-if", + "indexmap", "indoc", "libc", "memoffset", diff --git a/Cargo.toml b/Cargo.toml index c385fb2..4de5802 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,9 @@ crate-type = ["cdylib"] icicle-cpu = { path = "icicle-emu/icicle-cpu" } icicle-vm = { path = "icicle-emu/icicle-vm" } pcode = { path = "icicle-emu/sleigh/pcode" } -pyo3 = { version = "0.18.3", features = ["extension-module"] } +sleigh-runtime = { path = "icicle-emu/sleigh/sleigh-runtime" } +indexmap = "1.9.3" +pyo3 = { version = "0.18.3", features = ["extension-module", "indexmap"] } target-lexicon = "0.12.7" [build-dependencies] diff --git a/icicle.pyi b/icicle.pyi index e22aab0..abe0c41 100644 --- a/icicle.pyi +++ b/icicle.pyi @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Dict, Tuple from enum import Enum class MemoryProtection(Enum): @@ -37,6 +37,16 @@ class Icicle: def mem_write(self, address: int, data: bytes) -> None: ... + def reg_list(self) -> Dict[str, Tuple[int, int]]: ... + + def reg_offset(self, name: str) -> int: ... + + def reg_size(self, name: str) -> int: ... + + def reg_read(self, name: str) -> int: ... + + def reg_write(self, name: str, value: int) -> None: ... + def architectures() -> List[str]: """ List all the available architectures. diff --git a/src/lib.rs b/src/lib.rs index 6b4b2fb..a3a4713 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,26 +1,25 @@ mod rawenv; use std::borrow::Cow; +use std::collections::HashMap; use std::fmt; use std::fmt::{format, Formatter}; use icicle_cpu::mem::{Mapping, MemError, perm}; use pyo3::prelude::*; use icicle_vm; use icicle_vm::{BuildError, Vm}; +use icicle_vm::linux::LinuxCpu; use pyo3::types::PyList; use pyo3::{create_exception, import_exception}; -use pyo3::exceptions::PyException; +use pyo3::exceptions::*; use target_lexicon; +use indexmap::IndexMap; +use sleigh_runtime::NamedRegister; // References: // - https://pyo3.rs/main/conversions/tables // - https://pyo3.rs/main/class -#[pyclass(unsendable, module = "icicle")] -struct Icicle { - vm: icicle_vm::Vm, -} - #[pyclass(module = "icicle")] #[derive(Clone)] enum MemoryProtection { @@ -89,13 +88,32 @@ fn convert_protection(protection: MemoryProtection) -> u8 { } } +#[pyclass(unsendable, module = "icicle")] +struct Icicle { + vm: icicle_vm::Vm, + regs: HashMap, +} + +fn reg_find<'a>(i: &'a Icicle, name: &str) -> PyResult<&'a NamedRegister> { + let sleigh = i.vm.cpu.sleigh(); + match sleigh.get_reg(name) { + None => { + i.regs.get(name.to_lowercase().as_str()) + .ok_or( + PyKeyError::new_err(format!("Register not found: {name}")) + ) + } + Some(r) => Ok(r), + } +} + #[pymethods] impl Icicle { #[new] fn new(architecture: String) -> PyResult { if architecture.split("-").count() != 1 { return Err( - pyo3::exceptions::PyException::new_err(format!("Bad architecture format: {architecture}")) + PyException::new_err(format!("Bad architecture format: {architecture}")) ); } // Setup the CPU state for the target triple @@ -104,7 +122,7 @@ impl Icicle { ); if cpu_config.triple.architecture == target_lexicon::Architecture::Unknown { return Err( - pyo3::exceptions::PyException::new_err(format!("Unknown architecture: {architecture}")) + PyException::new_err(format!("Unknown architecture: {architecture}")) ); } @@ -114,11 +132,20 @@ impl Icicle { let mut vm = icicle_vm::build(&cpu_config) .map_err(|e| { - pyo3::exceptions::PyException::new_err(format!("VM build error: {e}")) + PyException::new_err(format!("VM build error: {e}")) })?; + // Get the lowercase register map + let mut regs = HashMap::new(); + let sleigh = vm.cpu.sleigh(); + for reg in &sleigh.named_registers { + let name = sleigh.get_str(reg.name); + regs.insert(name.to_lowercase(), reg.clone()); + } + Ok(Icicle { vm, + regs, }) } @@ -170,7 +197,7 @@ impl Icicle { .map_err(|e| { raise_MemoryError( format!("Failed to read memory {address:X}[{size:X}]"), - e + e, ) })?; return Ok(Cow::Owned(buffer)); @@ -182,10 +209,36 @@ impl Icicle { .map_err(|e| { raise_MemoryError( format!("Failed to write memory {address:X}[{size:X}]"), - e + e, ) }) } + + fn reg_list(&self) -> PyResult> { + let mut result = IndexMap::new(); + let sleigh = self.vm.cpu.sleigh(); + for reg in &sleigh.named_registers { + let name = sleigh.get_str(reg.name); + result.insert(name.to_string(), (reg.offset, reg.var.size)); + } + return Ok(result); + } + + fn reg_offset(&self, name: &str) -> PyResult { + Ok(reg_find(self, name)?.offset) + } + + fn reg_size(&self, name: &str) -> PyResult { + Ok(reg_find(self, name)?.var.size) + } + + fn reg_read(&mut self, name: &str) -> PyResult { + Ok(self.vm.cpu.read_reg(reg_find(self, name)?.var)) + } + + fn reg_write(&mut self, name: &str, value: u64) -> PyResult<()> { + Ok(self.vm.cpu.write_reg(reg_find(self, name)?.var, value)) + } } #[pyfunction] diff --git a/tests/example.py b/tests/example.py index 93dbbcb..99329bb 100644 --- a/tests/example.py +++ b/tests/example.py @@ -7,6 +7,10 @@ def main(): print(f"Architectures: {icicle.architectures()}") vm = icicle.Icicle("i686") + for name, (offset, size) in vm.reg_list().items(): + value = vm.reg_read(name.lower()) + print(f"{name}[{size}] @ {offset} = {value}") + addr = 0x10000 vm.mem_map(addr, 0x1000, icicle.MemoryProtection.NoAccess) try: