diff --git a/capi/include/yara_x.h b/capi/include/yara_x.h index 84646fc30..d504d35be 100644 --- a/capi/include/yara_x.h +++ b/capi/include/yara_x.h @@ -419,6 +419,9 @@ enum YRX_RESULT yrx_compiler_ban_module(struct YRX_COMPILER *compiler, enum YRX_RESULT yrx_compiler_new_namespace(struct YRX_COMPILER *compiler, const char *namespace_); +// Collects a string of a hashmap of all current loaded global vars +const char *yrx_compiler_get_globals(struct YRX_COMPILER *compiler); + // Defines a global variable of string type and sets its initial value. enum YRX_RESULT yrx_compiler_define_global_str(struct YRX_COMPILER *compiler, const char *ident, diff --git a/capi/src/compiler.rs b/capi/src/compiler.rs index 1df18a9b3..f57902776 100644 --- a/capi/src/compiler.rs +++ b/capi/src/compiler.rs @@ -1,4 +1,4 @@ -use std::ffi::{c_char, CStr}; +use std::ffi::{c_char, CStr, CString}; use std::mem; use std::mem::ManuallyDrop; @@ -387,6 +387,23 @@ unsafe fn yrx_compiler_define_global< } } +/// Collects a string of a hashmap of all current loaded global vars +#[no_mangle] +pub unsafe extern "C" fn yrx_compiler_get_globals( + compiler: *mut YRX_COMPILER, +) -> *const c_char { + let compiler = if let Some(compiler) = compiler.as_mut() { + compiler + } else { + return CString::new("Could not access the compiler").unwrap().into_raw(); + }; + + let globals = compiler.inner.show_globals(); + let json = serde_json::to_string(&globals).unwrap(); + + CString::new(json).unwrap().into_raw() +} + /// Defines a global variable of string type and sets its initial value. #[no_mangle] pub unsafe extern "C" fn yrx_compiler_define_global_str( diff --git a/go/compiler.go b/go/compiler.go index 4be9b6bdb..03b9d8b3d 100644 --- a/go/compiler.go +++ b/go/compiler.go @@ -557,6 +557,13 @@ func (c *Compiler) DefineGlobal(ident string, value interface{}) error { return nil } +// Returns a String of a hashmap of all of the currently loaded global variables in the compiler +func (c *Compiler) GetGlobals() string { + cStr := C.yrx_compiler_get_globals(c.cCompiler) + defer C.free(unsafe.Pointer(cStr)) + return C.GoString(cStr) +} + // Errors that occurred during the compilation, across multiple calls to // [Compiler.AddSource]. func (c *Compiler) Errors() []CompileError { diff --git a/go/compiler_test.go b/go/compiler_test.go index 359efc2ae..36f6ac0a5 100644 --- a/go/compiler_test.go +++ b/go/compiler_test.go @@ -6,6 +6,7 @@ import ( "testing" "os" "io/ioutil" + "encoding/json" "github.com/stretchr/testify/assert" ) @@ -106,6 +107,34 @@ func TestRelaxedReSyntax(t *testing.T) { assert.Len(t, scanResults.MatchingRules(), 1) } +func TestGetGlobals(t *testing.T) { + c, err := NewCompiler() + assert.NoError(t, err) + + x := map[string]interface{}{"a": map[string]interface{}{"a": "a"}, "b": "d"} + + c.DefineGlobal("A", "B") + c.DefineGlobal("B", 1.5) + c.DefineGlobal("C", x) + c.DefineGlobal("D", true) + + var globals map[string]interface{} + + // Unmarshal the JSON string into the map + err = json.Unmarshal([]byte(c.GetGlobals()), &globals) + assert.NoError(t, err) + + assert.Equal(t, globals["A"], "B") + assert.Equal(t, globals["B"], 1.5) + assert.Equal(t, globals["C"], x) + assert.Equal(t, globals["D"], true) + + c, err = NewCompiler() + assert.NoError(t, err) + + assert.Equal(t, c.GetGlobals(), "{}") +} + func TestConditionOptimization(t *testing.T) { _, err := Compile(` rule test { condition: true }`, diff --git a/lib/src/compiler/mod.rs b/lib/src/compiler/mod.rs index fda433e8f..260ac0845 100644 --- a/lib/src/compiler/mod.rs +++ b/lib/src/compiler/mod.rs @@ -682,6 +682,12 @@ impl<'a> Compiler<'a> { Ok(self) } + + /// Shows all current gloval variables of the compiler + pub fn show_globals(&mut self) -> serde_json::Value { + self.global_symbols.borrow_mut().show_globals() + } + /// Creates a new namespace. /// /// Further calls to [`Compiler::add_source`] will put the rules under the diff --git a/lib/src/symbols/mod.rs b/lib/src/symbols/mod.rs index 83a2e8290..a8000b754 100644 --- a/lib/src/symbols/mod.rs +++ b/lib/src/symbols/mod.rs @@ -210,6 +210,20 @@ impl SymbolTable { self.map.insert(ident.into(), symbol) } + /// Shows all currently loaded global vars inside the symbol table + pub fn show_globals(&self) -> serde_json::Value { + self.map + .iter() + .filter_map(|(k, v)| { + if let Symbol::Field { type_value, .. } = v { + Some((k.clone(), type_value.value_as_json())) + } else { + None + } + }) + .collect() + } + /// Returns true if the symbol table already contains a symbol with /// the given identifier. #[inline] diff --git a/lib/src/types/mod.rs b/lib/src/types/mod.rs index 912b4d229..5e852cc7e 100644 --- a/lib/src/types/mod.rs +++ b/lib/src/types/mod.rs @@ -7,6 +7,7 @@ use std::rc::Rc; use std::{mem, ptr}; use walrus::ir::InstrSeqType; use walrus::ValType; +use serde_json::{Map as JsonMap, Number as JsonNumber, Value as JsonValue}; use crate::modules::protos::yara::enum_value_options::Value as EnumValue; use crate::symbols::{Symbol, SymbolLookup, SymbolTable}; @@ -617,6 +618,79 @@ impl TypeValue { constraints: Some(constraints.into()), } } + + pub fn value_as_json(&self) -> JsonValue { + match self { + Self::Unknown => JsonValue::Null, + Self::Bool { value } => value.extract().cloned().map(JsonValue::Bool).unwrap_or(JsonValue::Null), + Self::Integer { value, .. } => { + if let Some(i) = value.extract().cloned() { + JsonValue::Number(JsonNumber::from(i)) + } else { + JsonValue::Null + } + } + Self::Float {value} => { + if let Some(f) = value.extract().cloned() { + JsonNumber::from_f64(f).map(JsonValue::Number).unwrap_or(JsonValue::Null) + } else { + JsonValue::Null + } + } + Self::String {value, ..} => { + if let Some(s) = value.extract().cloned() { + let s_str = String::from_utf8_lossy(s.as_slice()).into_owned(); + JsonValue::String(s_str) + } else { + JsonValue::Null + } + } + Self::Regexp(r) => { + if let Some(re) = r { + JsonValue::String(re.as_str().to_string()) + } else { + JsonValue::Null + } + } + Self::Struct(s) => { + let mut obj = JsonMap::new(); + for (key, field) in s.fields().iter() { + obj.insert(key.clone(), field.type_value.value_as_json()); + } + JsonValue::Object(obj) + } + Self::Array(a) => match a.as_ref() { + Array::Integers(items) => JsonValue::Array(items.iter().map(|i| JsonValue::Number(JsonNumber::from(*i))).collect()), + Array::Floats(items) => JsonValue::Array(items.iter().map(|f| JsonNumber::from_f64(*f).map(JsonValue::Number).unwrap_or(JsonValue::Null)).collect()), + Array::Bools(items) => JsonValue::Array(items.iter().map(|b| JsonValue::Bool(*b)).collect()), + Array::Strings(items) => JsonValue::Array(items.iter().map(|s| JsonValue::String(String::from_utf8_lossy(s.as_slice()).into_owned())).collect()), + Array::Structs(items) => JsonValue::Array(items.iter().map(|st| { + let mut obj = JsonMap::new(); + for (key, field) in st.fields().iter() { + obj.insert(key.clone(), field.type_value.value_as_json()); + } + JsonValue::Object(obj) + }).collect()), + } + Self::Map(m) => match m.as_ref() { + Map::IntegerKeys { map, .. } => { + let mut obj = JsonMap::new(); + for (k, v) in map.iter() { + obj.insert(k.to_string(), v.value_as_json()); + } + JsonValue::Object(obj) + } + Map::StringKeys { map, .. } => { + let mut obj = JsonMap::new(); + for (k, v) in map.iter() { + obj.insert(String::from_utf8_lossy(k.as_slice()).into_owned(), v.value_as_json()); + } + JsonValue::Object(obj) + } + } + Self::Func(_) => JsonValue::Null, + } + } } impl Display for TypeValue { diff --git a/lib/src/types/structure.rs b/lib/src/types/structure.rs index 0c5cdf72e..b38582073 100644 --- a/lib/src/types/structure.rs +++ b/lib/src/types/structure.rs @@ -170,6 +170,11 @@ impl Struct { self.protobuf_type_name.as_deref() } + /// Returns the fields in this structure. + pub fn fields(&self) -> &IndexMap { + &self.fields + } + /// Adds a new field to the structure. /// /// The field name may be a dot-separated sequence of field names, like diff --git a/py/src/lib.rs b/py/src/lib.rs index 5e784d702..440b90a33 100644 --- a/py/src/lib.rs +++ b/py/src/lib.rs @@ -15,6 +15,7 @@ matches = rules.scan(b'some dummy data') #![deny(missing_docs)] use std::borrow::Cow; +use std::collections::HashMap; use std::io::{Read, Write}; use std::marker::PhantomPinned; use std::ops::Deref; @@ -54,6 +55,22 @@ fn dict_to_json(dict: Bound) -> PyResult { .map_err(|err| PyValueError::new_err(err.to_string())) } +fn json_to_dict<'py>( + py: Python<'py>, + json: &serde_json::Value, +) -> PyResult> { + let json_str = serde_json::to_string(json) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + + static JSON_LOADS: OnceLock> = OnceLock::new(); + + let json_loads = JSON_LOADS.get_or_init(|| { + let json_mod = PyModule::import(py, "json").unwrap().unbind(); + json_mod.getattr(py, "loads").unwrap() + }); + json_loads.call1(py,(json_str,)) +} + #[derive(Debug, Clone, Display, EnumString, PartialEq)] #[strum(ascii_case_insensitive)] enum SupportedModules { @@ -521,6 +538,14 @@ impl Compiler { Ok(()) } + fn show_globals( + &mut self, + py: Python, + ) -> Py { + json_to_dict(py, &self.inner.show_globals()).unwrap() + } + + /// Creates a new namespace. /// /// Further calls to [`Compiler::add_source`] will put the rules under the diff --git a/py/tests/test_api.py b/py/tests/test_api.py index 0004b532a..d2828ace4 100644 --- a/py/tests/test_api.py +++ b/py/tests/test_api.py @@ -109,6 +109,19 @@ def test_dict_globals(): matching_rules = scanner.scan(b'').matching_rules assert len(matching_rules) == 1 +def test_show_globals(): + compiler = yara_x.Compiler() + compiler.define_global('some_dict', {"foo": "bar"}) + compiler.define_global("A", "B") + compiler.define_global("B", 1) + compiler.add_source('rule test {condition: some_dict.foo == "bar"}') + x = compiler.show_globals() + + assert(x['some_dict'] == {'foo': 'bar'}) + assert(x['A'] == "B") + assert(x['B'] == 1) + + def test_namespaces(): compiler = yara_x.Compiler() compiler.new_namespace('foo') diff --git a/py/yara_x.pyi b/py/yara_x.pyi index b40f669e4..ef36c8336 100644 --- a/py/yara_x.pyi +++ b/py/yara_x.pyi @@ -105,6 +105,11 @@ class Compiler: """ ... + def show_globals(self) -> dict: + r""" + Retrives a dict where the keys are the currently loaded global variable names and the values are their values + """ + def new_namespace(self, namespace: str) -> None: r""" Creates a new namespace.