diff --git a/Cargo.lock b/Cargo.lock index 7b5a43f4f..b42ae0603 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3346,6 +3346,12 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tinyzip" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6847a2bf223d67c916a25d5fb8d197fc1c7c3205421e05d50f3373e092034146" + [[package]] name = "tokio" version = "1.48.0" @@ -4416,6 +4422,7 @@ dependencies = [ "smallvec", "strum_macros", "thiserror 2.0.18", + "tinyzip", "uuid", "walrus", "wasm-bindgen", diff --git a/Cargo.toml b/Cargo.toml index 131036c34..c9f564988 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -112,6 +112,7 @@ yara-x-macros = { path = "macros", version = "1.18.0" } yara-x-parser = { path = "parser", version = "1.18.0" } yara-x-proto = { path = "proto", version = "1.18.0"} zip = { version = "8.2.0", default-features = false } +tinyzip = "0.4.0" simd-adler32 = "0.3.9" simd_cesu8 = "1.1.1" assert-call = "0.1.2" diff --git a/docs/ModuleDeveloperGuide.md b/docs/ModuleDeveloperGuide.md index 45ce0279b..6bbbf8320 100644 --- a/docs/ModuleDeveloperGuide.md +++ b/docs/ModuleDeveloperGuide.md @@ -313,7 +313,7 @@ So, let's create our `lib/src/modules/text.rs` file: use crate::mods::prelude::*; use crate::modules::protos::text::*; -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { let mut text_proto = Text::new(); // TODO: parse the data and populate text_proto. @@ -353,7 +353,7 @@ will be `crate::modules::protos::foobar` Next comes the module's main function and the module registration: ```rust -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { ... } @@ -361,8 +361,8 @@ register_module!("text", Text, main); ``` The module's main function is called for every file scanned by YARA. This -function receives a byte slice with the content of the file being scanned and an -optional byte slice with per-scan metadata, and it returns a `Result` containing the +function receives a mutable reference to a `ModuleContext` and a byte slice with +the content of the file being scanned. It returns a `Result` containing the `Text` structure that was generated from the `text.proto` file (or a `ModuleError`). Registering the module is as simple as calling the `register_module!` macro. @@ -379,7 +379,7 @@ use crate::modules::protos::text::*; use std::io; use std::io::BufRead; -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { // Create an empty instance of the Text protobuf. let mut text_proto = Text::new(); @@ -1093,11 +1093,11 @@ contract as in built-in modules—it receives the scanned data, populates the protobuf, and returns it: ```rust -use yara_x::errors::ModuleError; +use yara_x::mods::prelude::*; fn foobar_main( + _ctx: &mut ModuleContext, data: &[u8], - _meta: Option<&[u8]>, ) -> Result { let mut out = Foobar::new(); out.count = Some(data.len() as u64); diff --git a/examples/custom-module/src/lib.rs b/examples/custom-module/src/lib.rs index 96cb211fc..3629f3c63 100644 --- a/examples/custom-module/src/lib.rs +++ b/examples/custom-module/src/lib.rs @@ -19,8 +19,8 @@ pub mod proto { pub use proto::foobar::Foobar; fn foobar_main( + _ctx: &mut ModuleContext, data: &[u8], - _meta: Option<&[u8]>, ) -> Result { let mut out = Foobar::new(); out.count = Some(data.len() as u64); diff --git a/lib/Cargo.toml b/lib/Cargo.toml index a0190d4b5..7f4562e14 100644 --- a/lib/Cargo.toml +++ b/lib/Cargo.toml @@ -234,6 +234,9 @@ vt-module = [ "dep:psl", ] +# The `zip` module provides container file extraction capabilities. +zip-module = ["dep:tinyzip"] + # Enables all the default modules. default-modules = [ "console-module", @@ -252,6 +255,7 @@ default-modules = [ "test_proto2-module", "test_proto3-module", "vt-module", + "zip-module", ] # Features that are enabled by default. @@ -317,6 +321,7 @@ x509-parser = { workspace = true, optional = true } yara-x-macros = { workspace = true } yara-x-parser = { workspace = true, features = ["serde"] } zip = { workspace = true, optional = true, features = ["deflate"] } +tinyzip = { workspace = true, optional = true } simd-adler32 = { workspace = true, optional = true } simd_cesu8 = { workspace = true, optional = true } diff --git a/lib/src/compiler/context.rs b/lib/src/compiler/context.rs index 05cd2752a..752020cc1 100644 --- a/lib/src/compiler/context.rs +++ b/lib/src/compiler/context.rs @@ -127,8 +127,7 @@ impl<'src> CompileContext<'_, 'src> { // If the current symbol table is `None` it means that the // identifier is not a field or method of some structure. return if symbol_table.is_none() { - let module = crate::modules::registered_modules() - .find(|module| module.name() == ident.name); + let module = crate::modules::module_by_name(ident.name); // Build the error for the unknown identifier. let mut err = UnknownIdentifier::build( self.report_builder, diff --git a/lib/src/compiler/mod.rs b/lib/src/compiler/mod.rs index d7b070dd7..9ea8294dc 100644 --- a/lib/src/compiler/mod.rs +++ b/lib/src/compiler/mod.rs @@ -1913,8 +1913,7 @@ impl Compiler<'_> { fn c_import(&mut self, import: &Import) -> Result<(), CompileError> { let module_name = import.module_name; - let module = crate::modules::registered_modules() - .find(|m| m.name() == module_name); + let module = crate::modules::module_by_name(module_name); // Does a module with the given name actually exist? ... if module.is_none() { diff --git a/lib/src/modules/console.rs b/lib/src/modules/console.rs index 8697a8844..5df1c2e88 100644 --- a/lib/src/modules/console.rs +++ b/lib/src/modules/console.rs @@ -3,7 +3,7 @@ use std::borrow::Cow; use crate::mods::prelude::*; use crate::modules::protos::console::*; -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { // Nothing to do, but we have to return our protobuf Ok(Console::new()) } diff --git a/lib/src/modules/crx/mod.rs b/lib/src/modules/crx/mod.rs index 1d308b6c5..56c67ede0 100644 --- a/lib/src/modules/crx/mod.rs +++ b/lib/src/modules/crx/mod.rs @@ -18,7 +18,7 @@ thread_local!( const { RefCell::new(None) }; ); -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { PERMHASH_CACHE.with(|cache| *cache.borrow_mut() = None); match parser::Crx::parse(data) { Ok(crx) => Ok(crx.into()), diff --git a/lib/src/modules/cuckoo/mod.rs b/lib/src/modules/cuckoo/mod.rs index 3b43139c6..99176920b 100644 --- a/lib/src/modules/cuckoo/mod.rs +++ b/lib/src/modules/cuckoo/mod.rs @@ -22,8 +22,11 @@ fn set_local(value: schema::CuckooJson) { }); } -fn main(_data: &[u8], meta: Option<&[u8]>) -> Result { - let meta = match meta { +fn main( + _ctx: &mut ModuleContext, + _data: &[u8], +) -> Result { + let meta = match _ctx.get_module_metadata("cuckoo") { None | Some([]) => { set_local(schema::CuckooJson::default()); return Ok(Cuckoo::new()); diff --git a/lib/src/modules/dex/mod.rs b/lib/src/modules/dex/mod.rs index afd723661..2876aadb1 100644 --- a/lib/src/modules/dex/mod.rs +++ b/lib/src/modules/dex/mod.rs @@ -20,7 +20,7 @@ thread_local!( const { RefCell::new(None) }; ); -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { CHECKSUM_CACHE.with(|cache| *cache.borrow_mut() = None); SIGNATURE_CACHE.with(|cache| *cache.borrow_mut() = None); diff --git a/lib/src/modules/dotnet/mod.rs b/lib/src/modules/dotnet/mod.rs index 8581cc870..6366428f7 100644 --- a/lib/src/modules/dotnet/mod.rs +++ b/lib/src/modules/dotnet/mod.rs @@ -7,7 +7,7 @@ use crate::modules::protos::dotnet::*; pub mod parser; -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { match parser::Dotnet::parse(data) { Ok(dotnet) => Ok(dotnet.into()), Err(_) => { diff --git a/lib/src/modules/elf/mod.rs b/lib/src/modules/elf/mod.rs index 7319df84b..24a001d85 100644 --- a/lib/src/modules/elf/mod.rs +++ b/lib/src/modules/elf/mod.rs @@ -26,7 +26,7 @@ thread_local!( static TLSH_CACHE: RefCell> = const { RefCell::new(None) }; ); -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { IMPORT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None); TLSH_CACHE.with(|cache| *cache.borrow_mut() = None); diff --git a/lib/src/modules/hash/mod.rs b/lib/src/modules/hash/mod.rs index edc65e3cb..60a7f7bf7 100644 --- a/lib/src/modules/hash/mod.rs +++ b/lib/src/modules/hash/mod.rs @@ -29,7 +29,7 @@ thread_local!( RefCell::new(FxHashMap::default()); ); -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { // With every scanned file the cache must be cleared. SHA256_CACHE.with(|cache| cache.borrow_mut().clear()); SHA1_CACHE.with(|cache| cache.borrow_mut().clear()); diff --git a/lib/src/modules/lnk/mod.rs b/lib/src/modules/lnk/mod.rs index d7e2dcff1..def2894a5 100644 --- a/lib/src/modules/lnk/mod.rs +++ b/lib/src/modules/lnk/mod.rs @@ -15,7 +15,7 @@ use crate::mods::prelude::*; use crate::modules::protos::lnk::*; pub mod parser; -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { match parser::LnkParser::new().parse(data) { Ok(lnk) => Ok(lnk), Err(_) => { diff --git a/lib/src/modules/macho/mod.rs b/lib/src/modules/macho/mod.rs index a0ab4bd7b..cd3b8c23f 100644 --- a/lib/src/modules/macho/mod.rs +++ b/lib/src/modules/macho/mod.rs @@ -593,7 +593,7 @@ fn symhash(ctx: &mut ScanContext) -> Option>> { Some(Lowercase::>::new(digest)) } -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { DYLIB_MD5_CACHE.with(|cache| *cache.borrow_mut() = None); ENTITLEMENT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None); EXPORT_MD5_CACHE.with(|cache| *cache.borrow_mut() = None); diff --git a/lib/src/modules/magic/mod.rs b/lib/src/modules/magic/mod.rs index d67e60ae1..80d5e647f 100644 --- a/lib/src/modules/magic/mod.rs +++ b/lib/src/modules/magic/mod.rs @@ -27,7 +27,7 @@ thread_local! { static MIME_TYPE_CACHE: RefCell> = const { RefCell::new(None) }; } -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { // With every scanned file the cache must be cleared. TYPE_CACHE.set(None); MIME_TYPE_CACHE.set(None); diff --git a/lib/src/modules/math.rs b/lib/src/modules/math.rs index 61ef62313..51864a0d4 100644 --- a/lib/src/modules/math.rs +++ b/lib/src/modules/math.rs @@ -19,7 +19,7 @@ thread_local!( > = RefCell::new(FxHashMap::default()); ); -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { DISTRIBUTION_CACHE.with(|cache| cache.borrow_mut().clear()); // Nothing to do, but we have to return our protobuf diff --git a/lib/src/modules/mod.rs b/lib/src/modules/mod.rs index 0d7dab5bb..67c96d543 100644 --- a/lib/src/modules/mod.rs +++ b/lib/src/modules/mod.rs @@ -1,6 +1,6 @@ use protobuf::MessageDyn; use protobuf::reflect::MessageDescriptor; - +use rustc_hash::FxHashMap; use thiserror::Error; pub mod protos { @@ -36,6 +36,29 @@ pub enum ModuleError { }, } +/// Context passed to the main function of YARA modules. +#[derive(Default)] +pub struct ModuleContext<'a> { + module_metadata: FxHashMap<&'static str, &'a [u8]>, +} + +impl<'a> ModuleContext<'a> { + /// Set the metadata associated to the module with the given name. + pub fn set_module_metadata( + &mut self, + module_name: &'static str, + metadata: &'a [u8], + ) { + self.module_metadata.insert(module_name, metadata); + } + + /// Returns the metadata explicitly provided to the module with the given + /// name, if any. + pub fn get_module_metadata(&self, module_name: &str) -> Option<&[u8]> { + self.module_metadata.get(module_name).copied() + } +} + /// The trait implemented by all registered modules. pub trait RegisteredModule: Send + Sync { /// Name used for the module in `import` statements (e.g. `"my_module"`). @@ -47,10 +70,10 @@ pub trait RegisteredModule: Send + Sync { /// Main function called every time YARA scans some data, before /// evaluating the rules. Set to `None` for data-only modules. - fn main_fn( + fn main_fn<'a>( &self, - data: &[u8], - meta: Option<&[u8]>, + ctx: &mut ModuleContext<'a>, + data: &'a [u8], ) -> Option, ModuleError>>; /// Rust module path of the submodule inside the external crate that @@ -62,8 +85,8 @@ pub trait RegisteredModule: Send + Sync { fn rust_module_name(&self) -> Option<&'static str>; } -/// Main function in a YARA module. -pub type ModuleMainFn = fn(&[u8], Option<&[u8]>) -> Result; +pub type ModuleMainFn = + for<'a> fn(&mut ModuleContext<'a>, &'a [u8]) -> Result; /// Description of a YARA module, generic over the type `T` returned by the /// main function. @@ -93,13 +116,13 @@ where T::descriptor() } - fn main_fn( + fn main_fn<'a>( &self, - data: &[u8], - meta: Option<&[u8]>, + ctx: &mut ModuleContext<'a>, + data: &'a [u8], ) -> Option, ModuleError>> { self.main_fn.map(|f| { - f(data, meta).map(|ok| Box::new(ok) as Box) + f(ctx, data).map(|ok| Box::new(ok) as Box) }) } @@ -154,6 +177,14 @@ pub(crate) fn registered_modules() inventory::iter::<&'static dyn RegisteredModule>().copied() } +/// Returns a registered module given its name. +#[inline] +pub(crate) fn module_by_name( + name: &str, +) -> Option<&'static dyn RegisteredModule> { + registered_modules().find(|m| m.name() == name) +} + pub mod mods { /*! Utility functions and structures that allow invoking YARA modules directly. @@ -297,7 +328,13 @@ pub mod mods { let module = super::registered_modules() .find(|m| m.root_descriptor().full_name() == proto_name)?; - module.main_fn(data, meta)?.ok() + let mut ctx = super::ModuleContext::default(); + + if let Some(m) = meta { + ctx.module_metadata.insert(module.name(), m); + } + + module.main_fn(&mut ctx, data)?.ok() } /// Invokes all YARA modules and returns the data produced by them. @@ -333,8 +370,7 @@ pub mod mods { /// Returns the definition of the module with the given name. pub fn module_definition(name: &str) -> Option { use std::rc::Rc; - super::registered_modules() - .find(|m| m.name() == name) + super::module_by_name(name) .map(|m| reflect::Struct::new(Rc::::from(m))) } @@ -343,6 +379,7 @@ pub mod mods { #[allow(missing_docs)] pub mod prelude { pub use crate::modules::Module; + pub use crate::modules::ModuleContext; pub use crate::modules::ModuleError; pub use crate::modules::RegisteredModule; pub use crate::register_module; diff --git a/lib/src/modules/modules.rs b/lib/src/modules/modules.rs index d09bf4022..76a1d162a 100644 --- a/lib/src/modules/modules.rs +++ b/lib/src/modules/modules.rs @@ -32,4 +32,6 @@ mod test_proto3; #[cfg(feature = "time-module")] mod time; #[cfg(feature = "vt-module")] -mod vt; \ No newline at end of file +mod vt; +#[cfg(feature = "zip-module")] +mod zip; \ No newline at end of file diff --git a/lib/src/modules/pe/mod.rs b/lib/src/modules/pe/mod.rs index 142d1ad6d..a9385a54b 100644 --- a/lib/src/modules/pe/mod.rs +++ b/lib/src/modules/pe/mod.rs @@ -35,7 +35,7 @@ thread_local!( static CHECKSUM_CACHE: RefCell> = const { RefCell::new(None) }; ); -fn main(data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, data: &[u8]) -> Result { IMPHASH_CACHE.with(|cache| *cache.borrow_mut() = None); CHECKSUM_CACHE.with(|cache| *cache.borrow_mut() = None); diff --git a/lib/src/modules/protos/generated/mod.rs b/lib/src/modules/protos/generated/mod.rs index 3aa0693f7..8cee504b6 100644 --- a/lib/src/modules/protos/generated/mod.rs +++ b/lib/src/modules/protos/generated/mod.rs @@ -28,3 +28,4 @@ pub mod time; pub mod titan; pub mod vtnet; pub mod yara; +pub mod zip; diff --git a/lib/src/modules/protos/generated/zip.rs b/lib/src/modules/protos/generated/zip.rs new file mode 100644 index 000000000..8e9d495be --- /dev/null +++ b/lib/src/modules/protos/generated/zip.rs @@ -0,0 +1,687 @@ +// This file is generated by rust-protobuf 3.7.2. Do not edit +// .proto file is parsed by pure +// @generated + +// https://github.com/rust-lang/rust-clippy/issues/702 +#![allow(unknown_lints)] +#![allow(clippy::all)] + +#![allow(unused_attributes)] +#![cfg_attr(rustfmt, rustfmt::skip)] + +#![allow(dead_code)] +#![allow(missing_docs)] +#![allow(non_camel_case_types)] +#![allow(non_snake_case)] +#![allow(non_upper_case_globals)] +#![allow(trivial_casts)] +#![allow(unused_results)] +#![allow(unused_mut)] + +//! Generated file from `zip.proto` + +/// Generated files are compatible only with the same version +/// of protobuf runtime. +const _PROTOBUF_VERSION_CHECK: () = ::protobuf::VERSION_3_7_2; + +// @@protoc_insertion_point(message:zip.Entry) +#[derive(PartialEq,Clone,Default,Debug)] +pub struct Entry { + // message fields + // @@protoc_insertion_point(field:zip.Entry.file_path) + pub file_path: ::std::option::Option<::std::string::String>, + // @@protoc_insertion_point(field:zip.Entry.compression) + pub compression: ::std::option::Option<::protobuf::EnumOrUnknown>, + // @@protoc_insertion_point(field:zip.Entry.uncompressed_size) + pub uncompressed_size: ::std::option::Option, + // @@protoc_insertion_point(field:zip.Entry.compressed_size) + pub compressed_size: ::std::option::Option, + // special fields + // @@protoc_insertion_point(special_field:zip.Entry.special_fields) + pub special_fields: ::protobuf::SpecialFields, +} + +impl<'a> ::std::default::Default for &'a Entry { + fn default() -> &'a Entry { + ::default_instance() + } +} + +impl Entry { + pub fn new() -> Entry { + ::std::default::Default::default() + } + + // optional string file_path = 1; + + pub fn file_path(&self) -> &str { + match self.file_path.as_ref() { + Some(v) => v, + None => "", + } + } + + pub fn clear_file_path(&mut self) { + self.file_path = ::std::option::Option::None; + } + + pub fn has_file_path(&self) -> bool { + self.file_path.is_some() + } + + // Param is passed by value, moved + pub fn set_file_path(&mut self, v: ::std::string::String) { + self.file_path = ::std::option::Option::Some(v); + } + + // Mutable pointer to the field. + // If field is not initialized, it is initialized with default value first. + pub fn mut_file_path(&mut self) -> &mut ::std::string::String { + if self.file_path.is_none() { + self.file_path = ::std::option::Option::Some(::std::string::String::new()); + } + self.file_path.as_mut().unwrap() + } + + // Take field + pub fn take_file_path(&mut self) -> ::std::string::String { + self.file_path.take().unwrap_or_else(|| ::std::string::String::new()) + } + + // optional .zip.Compression compression = 2; + + pub fn compression(&self) -> Compression { + match self.compression { + Some(e) => e.enum_value_or(Compression::STORED), + None => Compression::STORED, + } + } + + pub fn clear_compression(&mut self) { + self.compression = ::std::option::Option::None; + } + + pub fn has_compression(&self) -> bool { + self.compression.is_some() + } + + // Param is passed by value, moved + pub fn set_compression(&mut self, v: Compression) { + self.compression = ::std::option::Option::Some(::protobuf::EnumOrUnknown::new(v)); + } + + // optional uint64 uncompressed_size = 3; + + pub fn uncompressed_size(&self) -> u64 { + self.uncompressed_size.unwrap_or(0) + } + + pub fn clear_uncompressed_size(&mut self) { + self.uncompressed_size = ::std::option::Option::None; + } + + pub fn has_uncompressed_size(&self) -> bool { + self.uncompressed_size.is_some() + } + + // Param is passed by value, moved + pub fn set_uncompressed_size(&mut self, v: u64) { + self.uncompressed_size = ::std::option::Option::Some(v); + } + + // optional uint64 compressed_size = 4; + + pub fn compressed_size(&self) -> u64 { + self.compressed_size.unwrap_or(0) + } + + pub fn clear_compressed_size(&mut self) { + self.compressed_size = ::std::option::Option::None; + } + + pub fn has_compressed_size(&self) -> bool { + self.compressed_size.is_some() + } + + // Param is passed by value, moved + pub fn set_compressed_size(&mut self, v: u64) { + self.compressed_size = ::std::option::Option::Some(v); + } + + fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData { + let mut fields = ::std::vec::Vec::with_capacity(4); + let mut oneofs = ::std::vec::Vec::with_capacity(0); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "file_path", + |m: &Entry| { &m.file_path }, + |m: &mut Entry| { &mut m.file_path }, + )); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "compression", + |m: &Entry| { &m.compression }, + |m: &mut Entry| { &mut m.compression }, + )); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "uncompressed_size", + |m: &Entry| { &m.uncompressed_size }, + |m: &mut Entry| { &mut m.uncompressed_size }, + )); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "compressed_size", + |m: &Entry| { &m.compressed_size }, + |m: &mut Entry| { &mut m.compressed_size }, + )); + ::protobuf::reflect::GeneratedMessageDescriptorData::new_2::( + "Entry", + fields, + oneofs, + ) + } +} + +impl ::protobuf::Message for Entry { + const NAME: &'static str = "Entry"; + + fn is_initialized(&self) -> bool { + true + } + + fn merge_from(&mut self, is: &mut ::protobuf::CodedInputStream<'_>) -> ::protobuf::Result<()> { + while let Some(tag) = is.read_raw_tag_or_eof()? { + match tag { + 10 => { + self.file_path = ::std::option::Option::Some(is.read_string()?); + }, + 16 => { + self.compression = ::std::option::Option::Some(is.read_enum_or_unknown()?); + }, + 24 => { + self.uncompressed_size = ::std::option::Option::Some(is.read_uint64()?); + }, + 32 => { + self.compressed_size = ::std::option::Option::Some(is.read_uint64()?); + }, + tag => { + ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?; + }, + }; + } + ::std::result::Result::Ok(()) + } + + // Compute sizes of nested messages + #[allow(unused_variables)] + fn compute_size(&self) -> u64 { + let mut my_size = 0; + if let Some(v) = self.file_path.as_ref() { + my_size += ::protobuf::rt::string_size(1, &v); + } + if let Some(v) = self.compression { + my_size += ::protobuf::rt::int32_size(2, v.value()); + } + if let Some(v) = self.uncompressed_size { + my_size += ::protobuf::rt::uint64_size(3, v); + } + if let Some(v) = self.compressed_size { + my_size += ::protobuf::rt::uint64_size(4, v); + } + my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields()); + self.special_fields.cached_size().set(my_size as u32); + my_size + } + + fn write_to_with_cached_sizes(&self, os: &mut ::protobuf::CodedOutputStream<'_>) -> ::protobuf::Result<()> { + if let Some(v) = self.file_path.as_ref() { + os.write_string(1, v)?; + } + if let Some(v) = self.compression { + os.write_enum(2, ::protobuf::EnumOrUnknown::value(&v))?; + } + if let Some(v) = self.uncompressed_size { + os.write_uint64(3, v)?; + } + if let Some(v) = self.compressed_size { + os.write_uint64(4, v)?; + } + os.write_unknown_fields(self.special_fields.unknown_fields())?; + ::std::result::Result::Ok(()) + } + + fn special_fields(&self) -> &::protobuf::SpecialFields { + &self.special_fields + } + + fn mut_special_fields(&mut self) -> &mut ::protobuf::SpecialFields { + &mut self.special_fields + } + + fn new() -> Entry { + Entry::new() + } + + fn clear(&mut self) { + self.file_path = ::std::option::Option::None; + self.compression = ::std::option::Option::None; + self.uncompressed_size = ::std::option::Option::None; + self.compressed_size = ::std::option::Option::None; + self.special_fields.clear(); + } + + fn default_instance() -> &'static Entry { + static instance: Entry = Entry { + file_path: ::std::option::Option::None, + compression: ::std::option::Option::None, + uncompressed_size: ::std::option::Option::None, + compressed_size: ::std::option::Option::None, + special_fields: ::protobuf::SpecialFields::new(), + }; + &instance + } +} + +impl ::protobuf::MessageFull for Entry { + fn descriptor() -> ::protobuf::reflect::MessageDescriptor { + static descriptor: ::protobuf::rt::Lazy<::protobuf::reflect::MessageDescriptor> = ::protobuf::rt::Lazy::new(); + descriptor.get(|| file_descriptor().message_by_package_relative_name("Entry").unwrap()).clone() + } +} + +impl ::std::fmt::Display for Entry { + fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result { + ::protobuf::text_format::fmt(self, f) + } +} + +impl ::protobuf::reflect::ProtobufValue for Entry { + type RuntimeType = ::protobuf::reflect::rt::RuntimeTypeMessage; +} + +// @@protoc_insertion_point(message:zip.Zip) +#[derive(PartialEq,Clone,Default,Debug)] +pub struct Zip { + // message fields + // @@protoc_insertion_point(field:zip.Zip.extracted_file_count) + pub extracted_file_count: ::std::option::Option, + // @@protoc_insertion_point(field:zip.Zip.is_zip) + pub is_zip: ::std::option::Option, + // @@protoc_insertion_point(field:zip.Zip.entries) + pub entries: ::std::vec::Vec, + // special fields + // @@protoc_insertion_point(special_field:zip.Zip.special_fields) + pub special_fields: ::protobuf::SpecialFields, +} + +impl<'a> ::std::default::Default for &'a Zip { + fn default() -> &'a Zip { + ::default_instance() + } +} + +impl Zip { + pub fn new() -> Zip { + ::std::default::Default::default() + } + + // optional int64 extracted_file_count = 1; + + pub fn extracted_file_count(&self) -> i64 { + self.extracted_file_count.unwrap_or(0) + } + + pub fn clear_extracted_file_count(&mut self) { + self.extracted_file_count = ::std::option::Option::None; + } + + pub fn has_extracted_file_count(&self) -> bool { + self.extracted_file_count.is_some() + } + + // Param is passed by value, moved + pub fn set_extracted_file_count(&mut self, v: i64) { + self.extracted_file_count = ::std::option::Option::Some(v); + } + + // optional bool is_zip = 2; + + pub fn is_zip(&self) -> bool { + self.is_zip.unwrap_or(false) + } + + pub fn clear_is_zip(&mut self) { + self.is_zip = ::std::option::Option::None; + } + + pub fn has_is_zip(&self) -> bool { + self.is_zip.is_some() + } + + // Param is passed by value, moved + pub fn set_is_zip(&mut self, v: bool) { + self.is_zip = ::std::option::Option::Some(v); + } + + fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData { + let mut fields = ::std::vec::Vec::with_capacity(3); + let mut oneofs = ::std::vec::Vec::with_capacity(0); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "extracted_file_count", + |m: &Zip| { &m.extracted_file_count }, + |m: &mut Zip| { &mut m.extracted_file_count }, + )); + fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>( + "is_zip", + |m: &Zip| { &m.is_zip }, + |m: &mut Zip| { &mut m.is_zip }, + )); + fields.push(::protobuf::reflect::rt::v2::make_vec_simpler_accessor::<_, _>( + "entries", + |m: &Zip| { &m.entries }, + |m: &mut Zip| { &mut m.entries }, + )); + ::protobuf::reflect::GeneratedMessageDescriptorData::new_2::( + "Zip", + fields, + oneofs, + ) + } +} + +impl ::protobuf::Message for Zip { + const NAME: &'static str = "Zip"; + + fn is_initialized(&self) -> bool { + true + } + + fn merge_from(&mut self, is: &mut ::protobuf::CodedInputStream<'_>) -> ::protobuf::Result<()> { + while let Some(tag) = is.read_raw_tag_or_eof()? { + match tag { + 8 => { + self.extracted_file_count = ::std::option::Option::Some(is.read_int64()?); + }, + 16 => { + self.is_zip = ::std::option::Option::Some(is.read_bool()?); + }, + 26 => { + self.entries.push(is.read_message()?); + }, + tag => { + ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?; + }, + }; + } + ::std::result::Result::Ok(()) + } + + // Compute sizes of nested messages + #[allow(unused_variables)] + fn compute_size(&self) -> u64 { + let mut my_size = 0; + if let Some(v) = self.extracted_file_count { + my_size += ::protobuf::rt::int64_size(1, v); + } + if let Some(v) = self.is_zip { + my_size += 1 + 1; + } + for value in &self.entries { + let len = value.compute_size(); + my_size += 1 + ::protobuf::rt::compute_raw_varint64_size(len) + len; + }; + my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields()); + self.special_fields.cached_size().set(my_size as u32); + my_size + } + + fn write_to_with_cached_sizes(&self, os: &mut ::protobuf::CodedOutputStream<'_>) -> ::protobuf::Result<()> { + if let Some(v) = self.extracted_file_count { + os.write_int64(1, v)?; + } + if let Some(v) = self.is_zip { + os.write_bool(2, v)?; + } + for v in &self.entries { + ::protobuf::rt::write_message_field_with_cached_size(3, v, os)?; + }; + os.write_unknown_fields(self.special_fields.unknown_fields())?; + ::std::result::Result::Ok(()) + } + + fn special_fields(&self) -> &::protobuf::SpecialFields { + &self.special_fields + } + + fn mut_special_fields(&mut self) -> &mut ::protobuf::SpecialFields { + &mut self.special_fields + } + + fn new() -> Zip { + Zip::new() + } + + fn clear(&mut self) { + self.extracted_file_count = ::std::option::Option::None; + self.is_zip = ::std::option::Option::None; + self.entries.clear(); + self.special_fields.clear(); + } + + fn default_instance() -> &'static Zip { + static instance: Zip = Zip { + extracted_file_count: ::std::option::Option::None, + is_zip: ::std::option::Option::None, + entries: ::std::vec::Vec::new(), + special_fields: ::protobuf::SpecialFields::new(), + }; + &instance + } +} + +impl ::protobuf::MessageFull for Zip { + fn descriptor() -> ::protobuf::reflect::MessageDescriptor { + static descriptor: ::protobuf::rt::Lazy<::protobuf::reflect::MessageDescriptor> = ::protobuf::rt::Lazy::new(); + descriptor.get(|| file_descriptor().message_by_package_relative_name("Zip").unwrap()).clone() + } +} + +impl ::std::fmt::Display for Zip { + fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result { + ::protobuf::text_format::fmt(self, f) + } +} + +impl ::protobuf::reflect::ProtobufValue for Zip { + type RuntimeType = ::protobuf::reflect::rt::RuntimeTypeMessage; +} + +#[derive(Clone,Copy,PartialEq,Eq,Debug,Hash)] +// @@protoc_insertion_point(enum:zip.Compression) +pub enum Compression { + // @@protoc_insertion_point(enum_value:zip.Compression.STORED) + STORED = 0, + // @@protoc_insertion_point(enum_value:zip.Compression.SHRUNK) + SHRUNK = 1, + // @@protoc_insertion_point(enum_value:zip.Compression.REDUCED_1) + REDUCED_1 = 2, + // @@protoc_insertion_point(enum_value:zip.Compression.REDUCED_2) + REDUCED_2 = 3, + // @@protoc_insertion_point(enum_value:zip.Compression.REDUCED_3) + REDUCED_3 = 4, + // @@protoc_insertion_point(enum_value:zip.Compression.REDUCED_4) + REDUCED_4 = 5, + // @@protoc_insertion_point(enum_value:zip.Compression.IMPLODED) + IMPLODED = 6, + // @@protoc_insertion_point(enum_value:zip.Compression.DEFLATED) + DEFLATED = 8, + // @@protoc_insertion_point(enum_value:zip.Compression.DEFLATE64) + DEFLATE64 = 9, + // @@protoc_insertion_point(enum_value:zip.Compression.PKWARE_IMPLODING) + PKWARE_IMPLODING = 10, + // @@protoc_insertion_point(enum_value:zip.Compression.BZIP2) + BZIP2 = 12, + // @@protoc_insertion_point(enum_value:zip.Compression.LZMA) + LZMA = 14, + // @@protoc_insertion_point(enum_value:zip.Compression.ZSTD) + ZSTD = 93, + // @@protoc_insertion_point(enum_value:zip.Compression.XZ) + XZ = 95, + // @@protoc_insertion_point(enum_value:zip.Compression.UNKNOWN) + UNKNOWN = 255, +} + +impl ::protobuf::Enum for Compression { + const NAME: &'static str = "Compression"; + + fn value(&self) -> i32 { + *self as i32 + } + + fn from_i32(value: i32) -> ::std::option::Option { + match value { + 0 => ::std::option::Option::Some(Compression::STORED), + 1 => ::std::option::Option::Some(Compression::SHRUNK), + 2 => ::std::option::Option::Some(Compression::REDUCED_1), + 3 => ::std::option::Option::Some(Compression::REDUCED_2), + 4 => ::std::option::Option::Some(Compression::REDUCED_3), + 5 => ::std::option::Option::Some(Compression::REDUCED_4), + 6 => ::std::option::Option::Some(Compression::IMPLODED), + 8 => ::std::option::Option::Some(Compression::DEFLATED), + 9 => ::std::option::Option::Some(Compression::DEFLATE64), + 10 => ::std::option::Option::Some(Compression::PKWARE_IMPLODING), + 12 => ::std::option::Option::Some(Compression::BZIP2), + 14 => ::std::option::Option::Some(Compression::LZMA), + 93 => ::std::option::Option::Some(Compression::ZSTD), + 95 => ::std::option::Option::Some(Compression::XZ), + 255 => ::std::option::Option::Some(Compression::UNKNOWN), + _ => ::std::option::Option::None + } + } + + fn from_str(str: &str) -> ::std::option::Option { + match str { + "STORED" => ::std::option::Option::Some(Compression::STORED), + "SHRUNK" => ::std::option::Option::Some(Compression::SHRUNK), + "REDUCED_1" => ::std::option::Option::Some(Compression::REDUCED_1), + "REDUCED_2" => ::std::option::Option::Some(Compression::REDUCED_2), + "REDUCED_3" => ::std::option::Option::Some(Compression::REDUCED_3), + "REDUCED_4" => ::std::option::Option::Some(Compression::REDUCED_4), + "IMPLODED" => ::std::option::Option::Some(Compression::IMPLODED), + "DEFLATED" => ::std::option::Option::Some(Compression::DEFLATED), + "DEFLATE64" => ::std::option::Option::Some(Compression::DEFLATE64), + "PKWARE_IMPLODING" => ::std::option::Option::Some(Compression::PKWARE_IMPLODING), + "BZIP2" => ::std::option::Option::Some(Compression::BZIP2), + "LZMA" => ::std::option::Option::Some(Compression::LZMA), + "ZSTD" => ::std::option::Option::Some(Compression::ZSTD), + "XZ" => ::std::option::Option::Some(Compression::XZ), + "UNKNOWN" => ::std::option::Option::Some(Compression::UNKNOWN), + _ => ::std::option::Option::None + } + } + + const VALUES: &'static [Compression] = &[ + Compression::STORED, + Compression::SHRUNK, + Compression::REDUCED_1, + Compression::REDUCED_2, + Compression::REDUCED_3, + Compression::REDUCED_4, + Compression::IMPLODED, + Compression::DEFLATED, + Compression::DEFLATE64, + Compression::PKWARE_IMPLODING, + Compression::BZIP2, + Compression::LZMA, + Compression::ZSTD, + Compression::XZ, + Compression::UNKNOWN, + ]; +} + +impl ::protobuf::EnumFull for Compression { + fn enum_descriptor() -> ::protobuf::reflect::EnumDescriptor { + static descriptor: ::protobuf::rt::Lazy<::protobuf::reflect::EnumDescriptor> = ::protobuf::rt::Lazy::new(); + descriptor.get(|| file_descriptor().enum_by_package_relative_name("Compression").unwrap()).clone() + } + + fn descriptor(&self) -> ::protobuf::reflect::EnumValueDescriptor { + let index = match self { + Compression::STORED => 0, + Compression::SHRUNK => 1, + Compression::REDUCED_1 => 2, + Compression::REDUCED_2 => 3, + Compression::REDUCED_3 => 4, + Compression::REDUCED_4 => 5, + Compression::IMPLODED => 6, + Compression::DEFLATED => 7, + Compression::DEFLATE64 => 8, + Compression::PKWARE_IMPLODING => 9, + Compression::BZIP2 => 10, + Compression::LZMA => 11, + Compression::ZSTD => 12, + Compression::XZ => 13, + Compression::UNKNOWN => 14, + }; + Self::enum_descriptor().value_by_index(index) + } +} + +impl ::std::default::Default for Compression { + fn default() -> Self { + Compression::STORED + } +} + +impl Compression { + fn generated_enum_descriptor_data() -> ::protobuf::reflect::GeneratedEnumDescriptorData { + ::protobuf::reflect::GeneratedEnumDescriptorData::new::("Compression") + } +} + +static file_descriptor_proto_data: &'static [u8] = b"\ + \n\tzip.proto\x12\x03zip\x1a\nyara.proto\"\xae\x01\n\x05Entry\x12\x1b\n\ + \tfile_path\x18\x01\x20\x01(\tR\x08filePath\x122\n\x0bcompression\x18\ + \x02\x20\x01(\x0e2\x10.zip.CompressionR\x0bcompression\x12+\n\x11uncompr\ + essed_size\x18\x03\x20\x01(\x04R\x10uncompressedSize\x12'\n\x0fcompresse\ + d_size\x18\x04\x20\x01(\x04R\x0ecompressedSize\"t\n\x03Zip\x120\n\x14ext\ + racted_file_count\x18\x01\x20\x01(\x03R\x12extractedFileCount\x12\x15\n\ + \x06is_zip\x18\x02\x20\x01(\x08R\x05isZip\x12$\n\x07entries\x18\x03\x20\ + \x03(\x0b2\n.zip.EntryR\x07entries*\xdf\x01\n\x0bCompression\x12\n\n\x06\ + STORED\x10\0\x12\n\n\x06SHRUNK\x10\x01\x12\r\n\tREDUCED_1\x10\x02\x12\r\ + \n\tREDUCED_2\x10\x03\x12\r\n\tREDUCED_3\x10\x04\x12\r\n\tREDUCED_4\x10\ + \x05\x12\x0c\n\x08IMPLODED\x10\x06\x12\x0c\n\x08DEFLATED\x10\x08\x12\r\n\ + \tDEFLATE64\x10\t\x12\x14\n\x10PKWARE_IMPLODING\x10\n\x12\t\n\x05BZIP2\ + \x10\x0c\x12\x08\n\x04LZMA\x10\x0e\x12\x08\n\x04ZSTD\x10]\x12\x06\n\x02X\ + Z\x10_\x12\x0c\n\x07UNKNOWN\x10\xff\x01\x1a\x06\x92\x93\x19\x02\x10\x01B\ + \x1e\xfa\x92\x19\x1a\n\x03zip\x12\x07zip.Zip\x1a\nzip-moduleb\x06proto2\ +"; + +/// `FileDescriptorProto` object which was a source for this generated file +fn file_descriptor_proto() -> &'static ::protobuf::descriptor::FileDescriptorProto { + static file_descriptor_proto_lazy: ::protobuf::rt::Lazy<::protobuf::descriptor::FileDescriptorProto> = ::protobuf::rt::Lazy::new(); + file_descriptor_proto_lazy.get(|| { + ::protobuf::Message::parse_from_bytes(file_descriptor_proto_data).unwrap() + }) +} + +/// `FileDescriptor` object which allows dynamic access to files +pub fn file_descriptor() -> &'static ::protobuf::reflect::FileDescriptor { + static generated_file_descriptor_lazy: ::protobuf::rt::Lazy<::protobuf::reflect::GeneratedFileDescriptor> = ::protobuf::rt::Lazy::new(); + static file_descriptor: ::protobuf::rt::Lazy<::protobuf::reflect::FileDescriptor> = ::protobuf::rt::Lazy::new(); + file_descriptor.get(|| { + let generated_file_descriptor = generated_file_descriptor_lazy.get(|| { + let mut deps = ::std::vec::Vec::with_capacity(1); + deps.push(super::yara::file_descriptor().clone()); + let mut messages = ::std::vec::Vec::with_capacity(2); + messages.push(Entry::generated_message_descriptor_data()); + messages.push(Zip::generated_message_descriptor_data()); + let mut enums = ::std::vec::Vec::with_capacity(1); + enums.push(Compression::generated_enum_descriptor_data()); + ::protobuf::reflect::GeneratedFileDescriptor::new_generated( + file_descriptor_proto(), + deps, + messages, + enums, + ) + }); + ::protobuf::reflect::FileDescriptor::new_generated_2(generated_file_descriptor) + }) +} diff --git a/lib/src/modules/protos/zip.proto b/lib/src/modules/protos/zip.proto new file mode 100644 index 000000000..0a29e899d --- /dev/null +++ b/lib/src/modules/protos/zip.proto @@ -0,0 +1,43 @@ +syntax = "proto2"; +import "yara.proto"; + +package zip; + +option (yara.module_options) = { + name : "zip" + root_message: "zip.Zip" + cargo_feature: "zip-module" +}; + +enum Compression { + option (yara.enum_options).inline = true; + STORED = 0; + SHRUNK = 1; + REDUCED_1 = 2; + REDUCED_2 = 3; + REDUCED_3 = 4; + REDUCED_4 = 5; + IMPLODED = 6; + DEFLATED = 8; + DEFLATE64 = 9; + PKWARE_IMPLODING = 10; + BZIP2 = 12; + LZMA = 14; + ZSTD = 93; + XZ = 95; + UNKNOWN = 255; +} + +message Entry { + optional string file_path = 1; + optional Compression compression = 2; + optional uint64 uncompressed_size = 3; + optional uint64 compressed_size = 4; +} + +message Zip { + // Number of files successfully extracted from the archive. + optional int64 extracted_file_count = 1; + optional bool is_zip = 2; + repeated Entry entries = 3; +} diff --git a/lib/src/modules/string.rs b/lib/src/modules/string.rs index 9ae287d6f..bf2235bfc 100644 --- a/lib/src/modules/string.rs +++ b/lib/src/modules/string.rs @@ -1,7 +1,7 @@ use crate::mods::prelude::*; use crate::modules::protos::string::*; -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { // Nothing to do, but we have to return our protobuf Ok(String::new()) } diff --git a/lib/src/modules/test_proto2/mod.rs b/lib/src/modules/test_proto2/mod.rs index 5383f57be..41c5ca815 100644 --- a/lib/src/modules/test_proto2/mod.rs +++ b/lib/src/modules/test_proto2/mod.rs @@ -95,7 +95,10 @@ fn to_int(ctx: &ScanContext, string: RuntimeString) -> Option { string.parse::().ok() } -fn main(data: &[u8], meta: Option<&[u8]>) -> Result { +fn main( + ctx: &mut ModuleContext, + data: &[u8], +) -> Result { let mut test = TestProto2::new(); test.set_int32_zero(0); @@ -187,7 +190,7 @@ fn main(data: &[u8], meta: Option<&[u8]>) -> Result { test.set_timestamp(1748591440); - test.metadata = meta.map(Vec::from); + test.metadata = ctx.get_module_metadata("test_proto2").map(Vec::from); Ok(test) } diff --git a/lib/src/modules/test_proto3/mod.rs b/lib/src/modules/test_proto3/mod.rs index 2cbb3e4bc..9e7867685 100644 --- a/lib/src/modules/test_proto3/mod.rs +++ b/lib/src/modules/test_proto3/mod.rs @@ -1,10 +1,7 @@ use crate::mods::prelude::*; use crate::modules::protos::test_proto3::TestProto3; -fn main( - _data: &[u8], - _meta: Option<&[u8]>, -) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { let mut test = TestProto3::new(); test.int32_zero = 0; diff --git a/lib/src/modules/time.rs b/lib/src/modules/time.rs index 4ce5eeca1..a96a51400 100644 --- a/lib/src/modules/time.rs +++ b/lib/src/modules/time.rs @@ -1,7 +1,7 @@ use crate::mods::prelude::*; use crate::modules::protos::time::*; -fn main(_data: &[u8], _meta: Option<&[u8]>) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { // Nothing to do, but we have to return our protobuf Ok(Time::new()) } diff --git a/lib/src/modules/vt/mod.rs b/lib/src/modules/vt/mod.rs index 4e1197f5a..96603b10a 100644 --- a/lib/src/modules/vt/mod.rs +++ b/lib/src/modules/vt/mod.rs @@ -49,10 +49,7 @@ static SUBDOMAIN: LazyLock = LazyLock::new(|| { Struct::enum_value_i64(&Permutation::SUBDOMAIN.descriptor()).unwrap() }); -fn main( - _data: &[u8], - _meta: Option<&[u8]>, -) -> Result { +fn main(_ctx: &mut ModuleContext, _data: &[u8]) -> Result { Ok(LiveHuntData::new()) } diff --git a/lib/src/modules/zip/mod.rs b/lib/src/modules/zip/mod.rs new file mode 100644 index 000000000..747302751 --- /dev/null +++ b/lib/src/modules/zip/mod.rs @@ -0,0 +1,68 @@ +use protobuf::Enum; +use tinyzip::Archive; + +use crate::mods::prelude::*; +use crate::modules::ModuleError; +use crate::modules::protos::zip::{Compression, Entry, Zip}; +use crate::register_module; + +pub fn main( + _ctx: &mut ModuleContext, + data: &[u8], +) -> Result { + let mut zip = Zip::new(); + + let archive = match Archive::open(data) { + Ok(arch) => arch, + Err(_) => { + zip.set_is_zip(false); + return Ok(zip); + } + }; + + zip.set_is_zip(true); + + let mut entries = Vec::new(); + let max_entries = 100000; // Guardrail: prevent DoS with huge entry counts + + let mut path_buf = vec![0u8; 65536]; + + for entry in archive.entries().filter_map(|entry| entry.ok()) { + if entries.len() >= max_entries { + break; + } + + let path_bytes = match entry.read_path(&mut path_buf) { + Ok(b) => b, + Err(_) => continue, + }; + + let file_path = String::from_utf8_lossy(path_bytes).to_string(); + + let mut proto_entry = Entry::new(); + + proto_entry.set_file_path(file_path); + proto_entry.set_uncompressed_size(entry.uncompressed_size()); + proto_entry.set_compressed_size(entry.compressed_size()); + + let compression = match entry.compression() { + Ok(tinyzip::Compression::Stored) => Compression::STORED, + Ok(tinyzip::Compression::Deflated) => Compression::DEFLATED, + Err(tinyzip::Error::UnsupportedCompression(raw)) => { + Compression::from_i32(raw as i32) + .unwrap_or(Compression::UNKNOWN) + } + Err(_) => Compression::UNKNOWN, + }; + + proto_entry.compression = Some(compression.into()); + + entries.push(proto_entry); + } + + zip.entries = entries; + + Ok(zip) +} + +register_module!("zip", Zip, main); diff --git a/lib/src/modules/zip/tests/testdata/stored_and_deflated.in.zip b/lib/src/modules/zip/tests/testdata/stored_and_deflated.in.zip new file mode 100644 index 000000000..696e1f934 Binary files /dev/null and b/lib/src/modules/zip/tests/testdata/stored_and_deflated.in.zip differ diff --git a/lib/src/modules/zip/tests/testdata/stored_and_deflated.out b/lib/src/modules/zip/tests/testdata/stored_and_deflated.out new file mode 100644 index 000000000..68d92d294 --- /dev/null +++ b/lib/src/modules/zip/tests/testdata/stored_and_deflated.out @@ -0,0 +1,10 @@ +is_zip: true +entries: + - file_path: "deflated_file.txt" + compression: DEFLATED + uncompressed_size: 21 + compressed_size: 23 + - file_path: "stored_file.txt" + compression: STORED + uncompressed_size: 19 + compressed_size: 19 \ No newline at end of file diff --git a/lib/src/scanner/mod.rs b/lib/src/scanner/mod.rs index 3d5d346a7..9de2fb89c 100644 --- a/lib/src/scanner/mod.rs +++ b/lib/src/scanner/mod.rs @@ -25,7 +25,9 @@ use thiserror::Error; use crate::Variable; use crate::compiler::{RuleId, Rules}; use crate::models::Rule; -use crate::modules::{ModuleError, RegisteredModule}; +use crate::modules::{ + ModuleContext, ModuleError, RegisteredModule, module_by_name, +}; pub(crate) use crate::scanner::context::RuntimeObject; pub(crate) use crate::scanner::context::RuntimeObjectHandle; pub(crate) use crate::scanner::context::ScanContext; @@ -487,23 +489,14 @@ impl<'r> Scanner<'r> { // Try to find the module by name first, if not found, then try // to find a module where the fully-qualified name for its protobuf // message matches the `name` arguments. - let descriptor = crate::modules::registered_modules() - .find_map(|module| { - if module.name() == name { - Some(module.root_descriptor()) - } else { - None - } - }) + let descriptor = module_by_name(name) + .map(|module| module.root_descriptor()) .or_else(|| { - crate::modules::registered_modules().find_map(|module| { - let descriptor = module.root_descriptor(); - if descriptor.full_name() == name { - Some(descriptor) - } else { - None - } - }) + crate::modules::registered_modules() + .find(|module| { + module.root_descriptor().full_name() == name + }) + .map(|module| module.root_descriptor()) }); if descriptor.is_none() { @@ -620,13 +613,33 @@ impl<'r> Scanner<'r> { // Set the global variable `filesize` to the size of the scanned data. ctx.set_filesize(data.len() as i64); + // Create the context that will be passed to the main function of each + // module. + let mut mod_ctx = ModuleContext::default(); + + // For each item in options.module_metadata, check if the module name + // is actually a registered module, and then add the corresponding + // metadata to mod_ctx. + for (module, meta) in options + .map(|options| options.module_metadata) + .into_iter() + .flatten() + .filter_map(|(name, meta)| Some((module_by_name(name)?, meta))) + { + mod_ctx.set_module_metadata(module.name(), meta); + } + // Indicate that the scanner is currently scanning the given data. ctx.scan_state = ScanState::ScanningData(data); + let data = match &ctx.scan_state { + ScanState::ScanningData(data) => data.as_ref(), + _ => unreachable!(), + }; + for module_name in ctx.compiled_rules.imports() { // Look up the module in the module registry. - let module = crate::modules::registered_modules() - .find(|module| module.name() == module_name) + let module = module_by_name(module_name) .unwrap_or_else(|| panic!("module `{module_name}` not found")); let module_root_descriptor = module.root_descriptor(); @@ -642,14 +655,7 @@ impl<'r> Scanner<'r> { { module_output = Some(output); } else { - let meta: Option<&'opts [u8]> = - options.as_ref().and_then(|options| { - options.module_metadata.get(module_name).copied() - }); - - if let Some(main_res) = - module.main_fn(ctx.scanned_data().unwrap(), meta) - { + if let Some(main_res) = module.main_fn(&mut mod_ctx, data) { module_output = Some(main_res.map_err(|err| { ScanError::ModuleError { module: module_name.to_string(), @@ -869,14 +875,8 @@ impl<'a, 'r> ScanResults<'a, 'r> { &self, module_name: &str, ) -> Option<&'a dyn MessageDyn> { - let module_descriptor = crate::modules::registered_modules() - .find_map(|m| { - if m.name() == module_name { - Some(m.root_descriptor()) - } else { - None - } - })?; + let module_descriptor = + module_by_name(module_name).map(|m| m.root_descriptor())?; let module_output = self .ctx .module_outputs diff --git a/ls/editors/code/package-lock.json b/ls/editors/code/package-lock.json index 055795a1f..301141697 100644 --- a/ls/editors/code/package-lock.json +++ b/ls/editors/code/package-lock.json @@ -5012,9 +5012,9 @@ "license": "MIT" }, "node_modules/undici": { - "version": "7.24.1", - "resolved": "https://registry.npmjs.org/undici/-/undici-7.24.1.tgz", - "integrity": "sha512-5xoBibbmnjlcR3jdqtY2Lnx7WbrD/tHlT01TmvqZUFVc9Q1w4+j5hbnapTqbcXITMH1ovjq/W7BkqBilHiVAaA==", + "version": "7.28.0", + "resolved": "https://registry.npmjs.org/undici/-/undici-7.28.0.tgz", + "integrity": "sha512-cRZYrTDwWznlnRiPjggAGxZXanty6M8RV1ff8Wm4LWXBp7/IG8v5DnOm74DtUBp9OONpK75YlPnIjQqX0dBDtA==", "dev": true, "license": "MIT", "engines": { diff --git a/ls/src/tests/testdata/completion10.response.json b/ls/src/tests/testdata/completion10.response.json index 5041b83aa..b10a8786c 100644 --- a/ls/src/tests/testdata/completion10.response.json +++ b/ls/src/tests/testdata/completion10.response.json @@ -78,5 +78,10 @@ "kind": 9, "label": "vt", "preselect": true + }, + { + "kind": 9, + "label": "zip", + "preselect": true } ] \ No newline at end of file diff --git a/ls/src/tests/testdata/completion11.response.json b/ls/src/tests/testdata/completion11.response.json index 5041b83aa..b10a8786c 100644 --- a/ls/src/tests/testdata/completion11.response.json +++ b/ls/src/tests/testdata/completion11.response.json @@ -78,5 +78,10 @@ "kind": 9, "label": "vt", "preselect": true + }, + { + "kind": 9, + "label": "zip", + "preselect": true } ] \ No newline at end of file diff --git a/ls/src/tests/testdata/completion14.response.json b/ls/src/tests/testdata/completion14.response.json index a78c66506..dc386473a 100644 --- a/ls/src/tests/testdata/completion14.response.json +++ b/ls/src/tests/testdata/completion14.response.json @@ -394,5 +394,24 @@ ], "kind": 9, "label": "vt" + }, + { + "additionalTextEdits": [ + { + "newText": "import \"zip\"\n", + "range": { + "end": { + "character": 0, + "line": 0 + }, + "start": { + "character": 0, + "line": 0 + } + } + } + ], + "kind": 9, + "label": "zip" } ] \ No newline at end of file diff --git a/ls/src/tests/testdata/completion15.response.json b/ls/src/tests/testdata/completion15.response.json index 6d4f0e6ea..786644415 100644 --- a/ls/src/tests/testdata/completion15.response.json +++ b/ls/src/tests/testdata/completion15.response.json @@ -373,5 +373,24 @@ ], "kind": 9, "label": "vt" + }, + { + "additionalTextEdits": [ + { + "newText": "import \"zip\"\n", + "range": { + "end": { + "character": 0, + "line": 0 + }, + "start": { + "character": 0, + "line": 0 + } + } + } + ], + "kind": 9, + "label": "zip" } ] \ No newline at end of file diff --git a/ls/src/tests/testdata/completion2.response.json b/ls/src/tests/testdata/completion2.response.json index ab40f27af..4503a49c4 100644 --- a/ls/src/tests/testdata/completion2.response.json +++ b/ls/src/tests/testdata/completion2.response.json @@ -395,5 +395,24 @@ ], "kind": 9, "label": "vt" + }, + { + "additionalTextEdits": [ + { + "newText": "import \"zip\"\n", + "range": { + "end": { + "character": 0, + "line": 0 + }, + "start": { + "character": 0, + "line": 0 + } + } + } + ], + "kind": 9, + "label": "zip" } ] \ No newline at end of file