Python basic bindings for sugar and qemu (#302)

* InMemoryBytesCoverageSugar python binding

* InMemoryBytesCoverageSugar python binding

* python mod for qemu in libafl_sugar

* libafl_qemu python

* fix

* clippy fixes

* clippy

* added pyo3-build-config for MacOS builds

* gitignor

* python is not default

Co-authored-by: Dominik Maier <domenukk@gmail.com>

.gitignore

@@ -2,6 +2,8 @@ target
 out
 Cargo.lock
 
+.env
+
 *.o
 *.a
 *.so
@@ -13,6 +15,7 @@ Cargo.lock
 *.dSYM
 
 .cur_input
+.venv
 
 crashes
 
@@ -31,4 +34,4 @@ AFLplusplus
 *.log
 a
 
-forkserver_test
+forkserver_test

libafl_concolic/symcc_runtime/build.rs

@@ -200,7 +200,6 @@ fn write_symcc_rename_header(rename_header_path: &Path, cpp_bindings: &bindgen::
     )
     .unwrap();
 
-    #[allow(clippy::filter_map)]
     cpp_bindings
         .to_string()
         .lines()

libafl_frida/src/cmplog_rt.rs

@@ -40,6 +40,7 @@ impl CmpLogRuntime {
     }
 
     /// Generate the instrumentation blobs for the current arch.
+    #[allow(clippy::similar_names)]
     fn generate_instrumentation_blobs(&mut self) {
         macro_rules! blr_to_populate {
             ($ops:ident) => {dynasm!($ops

libafl_frida/src/helper.rs

@@ -46,6 +46,7 @@ enum CmplogOperandType {
     Mem(capstone::RegId, capstone::RegId, i32, u32),
 }
 
+#[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
 enum SpecialCmpLogCase {
     Tbz,
     Tbnz,

libafl_qemu/Cargo.toml

@@ -10,8 +10,9 @@ license = "MIT OR Apache-2.0"
 keywords = ["fuzzing", "qemu", "instrumentation"]
 edition = "2018"
 
-
+[features]
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+python = ["pyo3"]
+default = []
 
 [dependencies]
 libafl = { path = "../libafl", version = "0.6.1" }
@@ -22,10 +23,12 @@ num = "0.4"
 num_enum = "0.5.1"
 goblin = "0.4.2"
 libc = "0.2.97"
+pyo3 = { version = "0.14.3", features = ["extension-module"], optional = true }
 
 [build-dependencies]
 cc = { version = "1.0" }
 which = "4.1"
 
 [lib]
-crate-type = ["rlib", "cdylib"]
+name = "libafl_qemu"
+crate-type = ["cdylib", "rlib"]

libafl_qemu/src/lib.rs

@@ -32,3 +32,73 @@ pub fn filter_qemu_args() -> Vec<String> {
     }
     args
 }
+
+#[cfg(all(target_od = "linux", feature = "python"))]
+use pyo3::prelude::*;
+
+#[cfg(all(target_od = "linux", feature = "python"))]
+#[pymodule]
+#[pyo3(name = "libafl_qemu")]
+#[allow(clippy::items_after_statements)]
+fn python_module(_py: Python, m: &PyModule) -> PyResult<()> {
+    use core::mem::transmute;
+    use pyo3::exceptions::PyValueError;
+
+    #[pyfn(m)]
+    fn init(args: Vec<String>, env: Vec<(String, String)>) -> i32 {
+        emu::init(&args, &env)
+    }
+
+    #[pyfn(m)]
+    fn write_mem(addr: u64, buf: &[u8]) {
+        emu::write_mem(addr, buf)
+    }
+    #[pyfn(m)]
+    fn read_mem(addr: u64, size: usize) -> Vec<u8> {
+        let mut buf = vec![];
+        unsafe { buf.set_len(size) };
+        emu::read_mem(addr, &mut buf);
+        buf
+    }
+    #[pyfn(m)]
+    fn num_regs() -> i32 {
+        emu::num_regs()
+    }
+    #[pyfn(m)]
+    fn write_reg(reg: i32, val: u64) -> PyResult<()> {
+        emu::write_reg(reg, val).map_err(|e| PyValueError::new_err(e))
+    }
+    #[pyfn(m)]
+    fn read_reg(reg: i32) -> PyResult<u64> {
+        emu::read_reg(reg).map_err(|e| PyValueError::new_err(e))
+    }
+    #[pyfn(m)]
+    fn set_breakpoint(addr: u64) {
+        emu::set_breakpoint(addr)
+    }
+    #[pyfn(m)]
+    fn remove_breakpoint(addr: u64) {
+        emu::remove_breakpoint(addr)
+    }
+    #[pyfn(m)]
+    fn run() {
+        emu::run()
+    }
+    #[pyfn(m)]
+    fn g2h(addr: u64) -> u64 {
+        unsafe { transmute(emu::g2h::<*const u8>(addr)) }
+    }
+    #[pyfn(m)]
+    fn h2g(addr: u64) -> u64 {
+        emu::h2g(unsafe { transmute::<_, *const u8>(addr) })
+    }
+    #[pyfn(m)]
+    fn binary_path() -> String {
+        emu::binary_path().to_owned()
+    }
+    #[pyfn(m)]
+    fn load_addr() -> u64 {
+        emu::load_addr()
+    }
+    Ok(())
+}

libafl_sugar/Cargo.toml

@@ -9,11 +9,24 @@ readme = "../README.md"
 license = "MIT OR Apache-2.0"
 keywords = ["fuzzing"]
 edition = "2018"
+build = "build.rs"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[features]
+python = ["pyo3", "libafl_qemu/python", "pyo3-build-config"]
+default = []
+
+[build-dependencies]
+pyo3-build-config = { version = "0.14.2", optional = true }
 
 [dependencies]
 libafl = { path = "../libafl", version = "0.6.1" }
 libafl_targets = { path = "../libafl_targets", version = "0.6.1" }
 libafl_qemu = { path = "../libafl_qemu", version = "0.6.1" }
 typed-builder = "0.9.0" # Implement the builder pattern at compiletime
+pyo3 = { version = "0.14.5", features = ["extension-module"], optional = true }
+
+
+[lib]
+name = "libafl_sugar"
+crate-type = ["cdylib", "rlib"]
+

libafl_sugar/build.rs

@@ -0,0 +1,4 @@
+fn main() {
+    #[cfg(feature = "python")]
+    pyo3_build_config::add_extension_module_link_args();
+}

libafl_sugar/src/inmemory.rs

@@ -15,22 +15,22 @@ use libafl::{
         QueueCorpusScheduler,
     },
     events::EventConfig,
-    executors::{inprocess::InProcessExecutor, ExitKind, TimeoutExecutor},
+    executors::{inprocess::InProcessExecutor, ExitKind, ShadowExecutor, TimeoutExecutor},
     feedback_or, feedback_or_fast,
     feedbacks::{CrashFeedback, MapFeedbackState, MaxMapFeedback, TimeFeedback, TimeoutFeedback},
     fuzzer::{Fuzzer, StdFuzzer},
     generators::RandBytesGenerator,
     inputs::{BytesInput, HasTargetBytes},
     mutators::scheduled::{havoc_mutations, tokens_mutations, StdScheduledMutator},
-    mutators::token_mutations::Tokens,
+    mutators::token_mutations::{I2SRandReplace, Tokens},
     observers::{HitcountsMapObserver, StdMapObserver, TimeObserver},
-    stages::StdMutationalStage,
+    stages::{ShadowTracingStage, StdMutationalStage},
     state::{HasCorpus, HasMetadata, StdState},
     stats::MultiStats,
     Error,
 };
 
-use libafl_targets::{EDGES_MAP, MAX_EDGES_NUM};
+use libafl_targets::{CmpLogObserver, CMPLOG_MAP, EDGES_MAP, MAX_EDGES_NUM};
 
 use crate::{CORPUS_CACHE_SIZE, DEFAULT_TIMEOUT_SECS};
 
@@ -53,8 +53,8 @@ where
     #[builder(default = None, setter(strip_option))]
     tokens_file: Option<PathBuf>,
     /// Flag if use CmpLog
-    //#[builder(default = false)]
+    #[builder(default = false)]
-    //use_cmplog: bool,
+    use_cmplog: bool,
     #[builder(default = 1337_u16)]
     broker_port: u16,
     /// The list of cores to run on
@@ -107,6 +107,9 @@ where
             // Create an observation channel to keep track of the execution time
             let time_observer = TimeObserver::new("time");
 
+            let cmplog = unsafe { &mut CMPLOG_MAP };
+            let cmplog_observer = CmpLogObserver::new("cmplog", cmplog, true);
+
             // The state of the edges feedback.
             let feedback_state = MapFeedbackState::with_observer(&edges_observer);
 
@@ -161,15 +164,18 @@ where
             };
 
             // Create the executor for an in-process function with one observer for edge coverage and one for the execution time
-            let mut executor = TimeoutExecutor::new(
+            let mut executor = ShadowExecutor::new(
-                InProcessExecutor::new(
+                TimeoutExecutor::new(
-                    &mut harness,
+                    InProcessExecutor::new(
-                    tuple_list!(edges_observer, time_observer),
+                        &mut harness,
-                    &mut fuzzer,
+                        tuple_list!(edges_observer, time_observer),
-                    &mut state,
+                        &mut fuzzer,
-                    &mut mgr,
+                        &mut state,
-                )?,
+                        &mut mgr,
-                timeout,
+                    )?,
+                    timeout,
+                ),
+                tuple_list!(cmplog_observer),
             );
 
             // In case the corpus is empty (on first run), reset
@@ -204,24 +210,40 @@ where
                 }
             }
 
+            // Setup a tracing stage in which we log comparisons
+            let tracing = ShadowTracingStage::new(&mut executor);
+
+            // Setup a randomic Input2State stage
+            let i2s = StdMutationalStage::new(StdScheduledMutator::new(tuple_list!(
+                I2SRandReplace::new()
+            )));
+
             if self.tokens_file.is_some() {
                 // Setup a basic mutator
                 let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations()));
                 let mutational = StdMutationalStage::new(mutator);
 
                 // The order of the stages matter!
-                let mut stages = tuple_list!(mutational);
+                if self.use_cmplog {
-
+                    let mut stages = tuple_list!(tracing, i2s, mutational);
-                fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                    fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                } else {
+                    let mut stages = tuple_list!(mutational);
+                    fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                }
             } else {
                 // Setup a basic mutator
                 let mutator = StdScheduledMutator::new(havoc_mutations());
                 let mutational = StdMutationalStage::new(mutator);
 
                 // The order of the stages matter!
-                let mut stages = tuple_list!(mutational);
+                if self.use_cmplog {
-
+                    let mut stages = tuple_list!(tracing, i2s, mutational);
-                fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                    fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                } else {
+                    let mut stages = tuple_list!(mutational);
+                    fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
+                }
             }
 
             Ok(())
@@ -244,3 +266,61 @@ where
         }
     }
 }
+
+#[cfg(feature = "python")]
+pub mod pybind {
+    use crate::inmemory;
+    use pyo3::prelude::*;
+    use pyo3::types::PyBytes;
+    use std::path::PathBuf;
+
+    #[pyclass(unsendable)]
+    struct InMemoryBytesCoverageSugar {
+        input_dirs: Vec<PathBuf>,
+        output_dir: PathBuf,
+        broker_port: u16,
+        cores: Vec<usize>,
+    }
+
+    #[pymethods]
+    impl InMemoryBytesCoverageSugar {
+        #[new]
+        fn new(
+            input_dirs: Vec<PathBuf>,
+            output_dir: PathBuf,
+            broker_port: u16,
+            cores: Vec<usize>,
+        ) -> Self {
+            Self {
+                input_dirs,
+                output_dir,
+                broker_port,
+                cores,
+            }
+        }
+
+        #[allow(clippy::needless_pass_by_value)]
+        pub fn run(&self, harness: PyObject) {
+            inmemory::InMemoryBytesCoverageSugar::builder()
+                .input_dirs(&self.input_dirs)
+                .output_dir(self.output_dir.clone())
+                .broker_port(self.broker_port)
+                .cores(&self.cores)
+                .harness(|buf| {
+                    Python::with_gil(|py| -> PyResult<()> {
+                        let args = (PyBytes::new(py, buf),); // TODO avoid copy
+                        harness.call1(py, args)?;
+                        Ok(())
+                    })
+                    .unwrap();
+                })
+                .build()
+                .run();
+        }
+    }
+
+    pub fn register(_py: Python, m: &PyModule) -> PyResult<()> {
+        m.add_class::<InMemoryBytesCoverageSugar>()?;
+        Ok(())
+    }
+}

libafl_sugar/src/lib.rs

@@ -1,4 +1,4 @@
-//! Sugar API to simplify the life of the naibe user of `LibAFL`
+//! Sugar API to simplify the life of the naive user of `LibAFL`
 
 pub mod inmemory;
 pub use inmemory::InMemoryBytesCoverageSugar;
@@ -10,3 +10,18 @@ pub use qemu::QemuBytesCoverageSugar;
 
 pub const DEFAULT_TIMEOUT_SECS: u64 = 1200;
 pub const CORPUS_CACHE_SIZE: usize = 4096;
+
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
+
+#[cfg(feature = "python")]
+#[pymodule]
+#[pyo3(name = "libafl_sugar")]
+fn python_module(py: Python, m: &PyModule) -> PyResult<()> {
+    inmemory::pybind::register(py, m)?;
+    #[cfg(target_os = "linux")]
+    {
+        qemu::pybind::register(py, m)?;
+    }
+    Ok(())
+}

libafl_sugar/src/qemu.rs

@@ -264,3 +264,61 @@ where
         launcher.build().launch().expect("Launcher failed");
     }
 }
+
+#[cfg(feature = "python")]
+pub mod pybind {
+    use crate::qemu;
+    use pyo3::prelude::*;
+    use pyo3::types::PyBytes;
+    use std::path::PathBuf;
+
+    #[pyclass(unsendable)]
+    struct QemuBytesCoverageSugar {
+        input_dirs: Vec<PathBuf>,
+        output_dir: PathBuf,
+        broker_port: u16,
+        cores: Vec<usize>,
+    }
+
+    #[pymethods]
+    impl QemuBytesCoverageSugar {
+        #[new]
+        fn new(
+            input_dirs: Vec<PathBuf>,
+            output_dir: PathBuf,
+            broker_port: u16,
+            cores: Vec<usize>,
+        ) -> Self {
+            Self {
+                input_dirs,
+                output_dir,
+                broker_port,
+                cores,
+            }
+        }
+
+        #[allow(clippy::needless_pass_by_value)]
+        pub fn run(&self, harness: PyObject) {
+            qemu::QemuBytesCoverageSugar::builder()
+                .input_dirs(&self.input_dirs)
+                .output_dir(self.output_dir.clone())
+                .broker_port(self.broker_port)
+                .cores(&self.cores)
+                .harness(|buf| {
+                    Python::with_gil(|py| -> PyResult<()> {
+                        let args = (PyBytes::new(py, buf),); // TODO avoid copy
+                        harness.call1(py, args)?;
+                        Ok(())
+                    })
+                    .unwrap();
+                })
+                .build()
+                .run();
+        }
+    }
+
+    pub fn register(_py: Python, m: &PyModule) -> PyResult<()> {
+        m.add_class::<QemuBytesCoverageSugar>()?;
+        Ok(())
+    }
+}