Add Python Type Annotations (#34)

* Some initial changes

* More typing changes

Author: toolCHAINZ

crackers/src/config/meta.rs

@@ -4,9 +4,9 @@ use rand::random;
 use serde::{Deserialize, Serialize};
 use tracing::Level;
 
-#[derive(Copy, Clone, Debug, Deserialize, Serialize)]
+#[derive(Copy, Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
 #[serde(rename_all = "UPPERCASE")]
-#[cfg_attr(feature = "pyo3", pyclass)]
+#[cfg_attr(feature = "pyo3", pyclass(eq, eq_int))]
 pub enum CrackersLogLevel {
     Trace,
     Debug,

crackers/src/config/synthesis.rs

@@ -41,44 +41,4 @@ impl SynthesisConfig {
             combine_instructions,
         }
     }
-
-    #[getter]
-    fn strategy(&self) -> SynthesisSelectionStrategy {
-        self.strategy
-    }
-
-    #[setter]
-    fn set_strategy(&mut self, strategy: SynthesisSelectionStrategy) {
-        self.strategy = strategy;
-    }
-
-    #[getter]
-    fn max_candidates_per_slot(&self) -> usize {
-        self.max_candidates_per_slot
-    }
-
-    #[setter]
-    fn set_max_candidates_per_slot(&mut self, max_candidates_per_slot: usize) {
-        self.max_candidates_per_slot = max_candidates_per_slot
-    }
-
-    #[getter]
-    fn parallel(&self) -> usize {
-        self.parallel
-    }
-
-    #[setter]
-    fn set_parallel(&mut self, parallel: usize) {
-        self.parallel = parallel
-    }
-
-    #[getter]
-    fn combine_instructions(&self) -> bool {
-        self.combine_instructions
-    }
-
-    #[setter]
-    fn set_combine_instructions(&mut self, combine_instructions: bool) {
-        self.combine_instructions = combine_instructions
-    }
 }

crackers/src/gadget/library/builder.rs

@@ -1,10 +1,10 @@
-use std::collections::HashSet;
-
 use derive_builder::Builder;
 use jingle::sleigh::OpCode;
 #[cfg(feature = "pyo3")]
 use pyo3::pyclass;
+use pyo3::pymethods;
 use serde::{Deserialize, Serialize};
+use std::collections::HashSet;
 
 use crate::config::error::CrackersConfigError;
 use crate::config::object::load_sleigh;
@@ -67,3 +67,52 @@ fn default_blacklist() -> HashSet<OpCode> {
         OpCode::CPUI_MULTIEQUAL,
     ])
 }
+
+/**
+
+pub sample_size: Option<usize>,
+pub base_address: Option<u64>,
+*/
+
+#[pymethods]
+impl GadgetLibraryConfig {
+    #[getter]
+    pub fn get_max_gadget_length(&self) -> usize {
+        self.max_gadget_length
+    }
+
+    #[setter]
+    pub fn set_max_gadget_length(&mut self, l: usize) {
+        self.max_gadget_length = l;
+    }
+
+    #[getter]
+    pub fn get_path(&self) -> &str {
+        self.path.as_str()
+    }
+
+    #[setter]
+    pub fn set_path(&mut self, l: String) {
+        self.path = l;
+    }
+
+    #[getter]
+    pub fn get_sample_size(&self) -> Option<usize> {
+        self.sample_size
+    }
+
+    #[setter]
+    pub fn set_sample_size(&mut self, l: Option<usize>) {
+        self.sample_size = l;
+    }
+
+    #[getter]
+    pub fn get_base_address(&self) -> Option<u64> {
+        self.base_address
+    }
+
+    #[setter]
+    pub fn set_base_address(&mut self, l: Option<u64>) {
+        self.base_address = l;
+    }
+}

crackers_python/crackers/__init__.py

@@ -0,0 +1,6 @@
+import z3
+from .crackers import *
+
+__doc__ = crackers.__doc__
+if hasattr(crackers, "__all__"):
+    __all__ = crackers.__all__

crackers_python/crackers/__init__.pyi

@@ -0,0 +1,196 @@
+from typing import Any, Union, List, Optional, Iterable, Callable
+import z3  # Assuming Z3 is imported for type annotations
+
+
+class SpaceInfo:
+    # Placeholder for SpaceInfo class
+    ...
+
+class VarNode:
+    def __init__(self, space_index: int, offset: int, size: int) -> None: ...
+
+class RawVarNodeDisplay:
+    def __init__(self, offset: int, size: int, space_info: SpaceInfo) -> None: ...
+
+class VarNodeDisplay:
+    def __init__(self, raw: RawVarNodeDisplay = ..., register: tuple[str, VarNode] = ...) -> None: ...
+    # Represents the enum variants Raw and Register
+    raw: RawVarNodeDisplay
+    register: tuple[str, VarNode]
+
+
+class ResolvedIndirectVarNode:
+    def __init__(self, pointer: Any, pointer_space_info: SpaceInfo, access_size_bytes: int) -> None: ...
+
+    def pointer_bv(self) -> z3.BitVecRef: ...
+    def space_name(self) -> str: ...
+    def access_size(self) -> int: ...
+
+
+class ResolvedVarNode:
+    """
+    Represents the PythonResolvedVarNode enum with two variants:
+    - Direct: Contains a VarNodeDisplay
+    - Indirect: Contains a ResolvedIndirectVarNode
+    """
+    def __init__(self, value: Union[VarNodeDisplay, ResolvedIndirectVarNode]) -> None: ...
+    value: Union[VarNodeDisplay, ResolvedIndirectVarNode]
+
+
+class PcodeOperation:
+    pass
+
+
+class Instruction:
+    """
+    Represents a Python wrapper for a Ghidra instruction.
+    """
+    disassembly: str
+    def pcode(self) -> List[PcodeOperation]: ...
+
+class State:
+    def __init__(self, jingle: JingleContext) -> State: ...
+
+    def varnode(self, varnode: ResolvedVarNode) -> z3.BitVecRef: ...
+    def register(self, name: str) -> z3.BitVecRef: ...
+    def ram(self, offset: int, length: int) -> z3.BitVecRef: ...
+
+class ModeledInstruction:
+    original_state: State
+    final_state: State
+
+    def get_input_vns(self) -> Iterable[ResolvedVarNode]: ...
+    def get_output_vns(self) -> Iterable[ResolvedVarNode]: ...
+
+class ModeledBlock:
+    original_state: State
+    final_state: State
+    def get_input_vns(self) -> Iterable[ResolvedVarNode]: ...
+    def get_output_vns(self) -> Iterable[ResolvedVarNode]: ...
+
+class JingleContext:
+    def __init__(self, binary_path: str, ghidra: str) -> JingleContext: ...
+    def model_instruction_at(self, offset: int) -> ModeledInstruction: ...
+    def model_block_at(self, offset: int, max_instrs: int) -> ModeledBlock: ...
+
+class SleighContext:
+    """
+    Represents a Sleigh context in python.
+    """
+    def __init__(self, binary_path: str, ghidra: str) -> SleighContext: ...
+    base_address: int
+    def instruction_at(self, offset: int) -> Optional[Instruction]: ...
+    def make_jingle_context(self) -> JingleContext: ...
+
+
+"""
+Begin crackers types
+"""
+
+class CrackersLogLevel:
+    Debug = CrackersLogLevel.Debug
+    Error = CrackersLogLevel.Error
+    Info = CrackersLogLevel.Info
+    Trace = CrackersLogLevel.Trace
+    Warn = CrackersLogLevel.Warn
+
+class MetaConfig:
+    seed: int
+    log_level: CrackersLogLevel
+
+class SpecificationConfig:
+    path: str
+    max_instructions: int
+    base_address: Optional[int]
+
+class GadgetLibraryConfig:
+    max_gadget_length: int
+    path: str
+    sample_size: Optional[int]
+    base_address: Optional[int]
+
+class SleighConfig:
+    ghidra_path: str
+
+class SynthesisSelectionStrategy:
+    SatStrategy = SynthesisSelectionStrategy.SatStrategy
+    OptimizeStrategy = SynthesisSelectionStrategy.OptimizeStrategy
+
+class SynthesisConfig:
+    strategy: SynthesisSelectionStrategy
+    max_candidates_per_slot: int
+    parallel: int
+    combine_instructions: bool
+
+class MemoryEqualityConstraint:
+    space: str
+    address: int
+    size: int
+    value: int
+
+class StateEqualityConstraint:
+    register: Optional[dict[str, int]]
+    pointer: Optional[dict[str, str]]
+    memory: Optional[MemoryEqualityConstraint]
+
+class PointerRange:
+    min: int
+    max: int
+class PointerRangeConstraints:
+    read: Optional[list[PointerRange]]
+class ConstraintConfig:
+    precondition: Optional[StateEqualityConstraint]
+    postcondition: Optional[StateEqualityConstraint]
+    pointer: Optional[PointerRangeConstraints]
+
+
+class CrackersConfig:
+    meta: MetaConfig
+    spec: SpecificationConfig
+    library: GadgetLibraryConfig
+    sleigh: SleighConfig
+    synthesis: SynthesisConfig
+    constraint: ConstraintConfig
+
+    @classmethod
+    def from_toml(cls, path: str) -> CrackersConfig: ...
+
+    def resolve_config(self) -> SynthesisParams: ...
+
+class AssignmentModel:
+    def eval_bv(self, bv: z3.BitVecRef, model_completion: bool) -> z3.BitVecRef: ...
+    def initial_state(self) -> Optional[State]: ...
+    def final_state(self) -> Optional[State]: ...
+    def gadgets(self) -> list[ModeledBlock]: ...
+    def inputs(self) -> Iterable[ResolvedVarNode]: ...
+    def outputs(self) -> Iterable[ResolvedVarNode]: ...
+    def input_summary(self, model_completion: bool): Iterable[tuple[str, z3.BitVecRef]]
+    def output_summary(self, model_completion: bool): Iterable[tuple[str, z3.BitVecRef]]
+
+
+class PythonDecisionResult_AssignmentFound(DecisionResult):
+    _0: AssignmentModel
+    __match_args__ = ('_0',)
+
+
+class SelectionFailure:
+    indices: list[int]
+
+class PythonDecisionResult_Unsat(DecisionResult):
+    _0: SelectionFailure
+    pass
+
+class DecisionResult:
+    AssignmentFound : PythonDecisionResult_AssignmentFound
+    Unsat : PythonDecisionResult_Unsat
+
+DecisionResultType = Union[DecisionResult.AssignmentFound, DecisionResult.Unsat]
+
+type StateConstraintGenerator = Callable[[State, int], z3.BitVecRef]
+type TransitionConstraintGenerator = Callable[[ModeledBlock], z3.BitVecRef]
+class SynthesisParams:
+    def run(self) -> DecisionResultType: ...
+    def add_precondition(self, fn: StateConstraintGenerator): ...
+    def add_postcondition(self, fn: StateConstraintGenerator): ...
+    def add_transition_constraint(self, fn: TransitionConstraintGenerator): ...
+    pass

crackers_python/crackers/py.typed

@@ -15,7 +15,7 @@ use pyo3::{Py, PyAny, PyResult, pyclass, pymethods};
 use std::rc::Rc;
 use z3::ast::BV;
 
-#[pyclass(unsendable)]
+#[pyclass(unsendable, name = "AssignmentModel")]
 #[derive(Clone)]
 pub struct PythonAssignmentModel {
     pub inner: Rc<AssignmentModel<'static, ModeledBlock<'static>>>,

crackers_python/src/decision/assignment_model/mod.rs

@@ -4,7 +4,7 @@ use pyo3::pyclass;
 
 pub mod assignment_model;
 
-#[pyclass(unsendable)]
+#[pyclass(unsendable, name = "DecisionResult")]
 pub enum PythonDecisionResult {
     AssignmentFound(PythonAssignmentModel),
     Unsat(SelectionFailure),

crackers_python/src/decision/mod.rs

@@ -3,6 +3,8 @@ mod decision;
 mod synthesis;
 
 use crate::config::PythonCrackersConfig;
+use crate::decision::PythonDecisionResult;
+use crate::synthesis::PythonSynthesisParams;
 use ::crackers::config::constraint::{
     ConstraintConfig, MemoryEqualityConstraint, PointerRange, PointerRangeConstraints,
     StateEqualityConstraint,
@@ -13,14 +15,38 @@ use ::crackers::config::specification::SpecificationConfig;
 use ::crackers::config::synthesis::SynthesisConfig;
 use ::crackers::gadget::library::builder::GadgetLibraryConfig;
 use ::crackers::synthesis::builder::SynthesisSelectionStrategy;
+use ::jingle::python::instruction::PythonInstruction;
+use ::jingle::python::jingle_context::PythonJingleContext;
+use ::jingle::python::modeled_block::PythonModeledBlock;
+use ::jingle::python::modeled_instruction::PythonModeledInstruction;
+use ::jingle::python::sleigh_context::LoadedSleighContextWrapper;
+use ::jingle::python::state::PythonState;
+use ::jingle::sleigh::{IndirectVarNode, PcodeOperation, VarNode};
 use pyo3::prelude::*;
-use std::ffi::CString;
 
+#[pymodule]
+#[pyo3(submodule)]
+fn jingle(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<VarNode>()?;
+    m.add_class::<IndirectVarNode>()?;
+    m.add_class::<PcodeOperation>()?;
+    m.add_class::<PythonInstruction>()?;
+    m.add_class::<LoadedSleighContextWrapper>()?;
+    m.add_class::<PythonJingleContext>()?;
+    m.add_class::<PythonState>()?;
+    m.add_class::<PythonModeledInstruction>()?;
+    m.add_class::<PythonModeledBlock>()?;
+    Ok(())
+}
 /// A Python module implemented in Rust.
 #[pymodule]
 fn crackers(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.py().run(&CString::new("import z3")?, None, None)?;
+    let j = PyModule::new(m.py(), "jingle")?;
+    jingle(&j)?;
+    m.add_submodule(&j)?;
     m.add_class::<PythonCrackersConfig>()?;
+    m.add_class::<PythonDecisionResult>()?;
+    m.add_class::<PythonSynthesisParams>()?;
     m.add_class::<MetaConfig>()?;
     m.add_class::<SpecificationConfig>()?;
     m.add_class::<SleighConfig>()?;