mod.rs

use std::sync::LazyLock;

use protobuf::MessageDyn;
use protobuf::reflect::MessageDescriptor;
use rustc_hash::FxHashMap;

use thiserror::Error;

pub mod protos {
    #[cfg(feature = "generate-proto-code")]
    include!(concat!(env!("OUT_DIR"), "/protos/mod.rs"));

    #[cfg(not(feature = "generate-proto-code"))]
    include!("protos/generated/mod.rs");
}

#[cfg(test)]
mod tests;

#[allow(unused_imports)]
pub(crate) mod prelude {
    pub(crate) use crate::scanner::ScanContext;
    pub(crate) use crate::wasm::runtime::Caller;
    pub(crate) use crate::wasm::string::FixedLenString;
    pub(crate) use crate::wasm::string::RuntimeString;
    pub(crate) use crate::wasm::string::String as _;
    pub(crate) use crate::wasm::string::{Lowercase, Uppercase};
    pub(crate) use crate::wasm::*;
    pub(crate) use bstr::ByteSlice;
    #[cfg(not(feature = "inventory"))]
    pub(crate) use linkme::distributed_slice;
    pub(crate) use yara_x_macros::{module_export, module_main, wasm_export};
}

include!("modules.rs");

/// Enum describing errors occurred in modules.
#[derive(Error, Debug)]
#[non_exhaustive]
pub enum ModuleError {
    /// Invalid format of module metadata.
    #[error("invalid metadata: {err}")]
    MetadataError {
        /// The error that actually occurred.
        err: String,
    },
    /// Error occurred when processing the input data.
    #[error("internal error: {err}")]
    InternalError {
        /// The error that actually occurred.
        err: String,
    },
}

/// Signature of a module's main function.
type MainFn =
    fn(&[u8], Option<&[u8]>) -> Result<Box<dyn MessageDyn>, ModuleError>;

/// A structure describing a YARA module.
pub(crate) struct Module {
    /// Pointer to the module's main function.
    pub main_fn: Option<MainFn>,
    /// Name of the Rust module, if any, that contains code for this YARA
    /// module (e.g: "test_proto2").
    pub rust_module_name: Option<&'static str>,
    /// A [`MessageDescriptor`] that describes the module's structure. This
    /// corresponds to the protobuf message declared in the "root_message"
    /// for the YARA module. It allows iterating the fields declared by the
    /// module and obtaining their names and types.
    pub root_struct_descriptor: MessageDescriptor,
}

/// Macro that adds a module to the `BUILTIN_MODULES` map.
///
/// This macro is used by `add_modules.rs`, a file that is automatically
/// generated by `build.rs` based on the Protocol Buffers defined in the
/// `src/modules/protos` directory.
///
/// # Example
///
/// add_module!(modules, "test", test, "Test", test_mod, Some(test::main as
/// MainFn));
macro_rules! add_module {
    ($modules:expr, $name:literal, $proto:ident, $root_message:literal, $rust_module_name:expr, $main_fn:expr) => {{
        use std::stringify;
        let root_struct_descriptor = protos::$proto::file_descriptor()
            // message_by_full_name expects a dot (.) at the beginning
            // of the name.
            .message_by_full_name(format!(".{}", $root_message).as_str())
            .expect(format!(
                "`root_message` option in protobuf `{}` is wrong, message `{}` is not defined",
                stringify!($proto),
                $root_message
            ).as_str());

        $modules.insert(
            $name,
            Module {
                main_fn: $main_fn,
                rust_module_name: $rust_module_name,
                root_struct_descriptor,
            },
        );
    }};
}

/// `BUILTIN_MODULES` is a static, global map where keys are module names
/// and values are [`Module`] structures that describe a YARA module.
///
/// This table is populated with the modules defined by a `.proto` file in
/// `src/modules/protos`. Each `.proto` file that contains a statement like
/// the following one defines a YARA module:
///
/// ```protobuf
/// option (yara.module_options) = {
///   name : "foo"
///   root_message: "Foo"
///   rust_module: "foo"
/// };
/// ```
///
/// The `name` field is the module's name (i.e: the name used in `import`
/// statements), which is also the key in `BUILTIN_MODULES`. `root_message`
/// is the name of the message that describes the module's structure. This
/// is required because a `.proto` file can define more than one message.
///
/// `rust_module` is the name of the Rust module where functions exported
/// by the YARA module are defined. This field is optional, if not provided
/// the module is considered a data-only module.
pub(crate) static BUILTIN_MODULES: LazyLock<FxHashMap<&'static str, Module>> =
    LazyLock::new(|| {
        let mut modules = FxHashMap::default();
        // The `add_modules.rs` file is automatically generated at compile time
        // by `build.rs`. This is an example of how `add_modules.rs` looks like:
        //
        // {
        //  #[cfg(feature = "pe_module")]
        //  add_module!(modules, "pe", pe, "pe.PE", Some("pe"), Some(pe::__main__ as MainFn));
        //
        //  #[cfg(feature = "elf_module")]
        //  add_module!(modules, "elf", elf, "elf.ELF", Some("elf"), Some(elf::__main__ as MainFn));
        // }
        //
        // `add_modules.rs` will contain an `add_module!` statement for each
        // protobuf in `src/modules/protos` defining a YARA module.
        include!("add_modules.rs");
        modules
    });

pub mod mods {
    /*! Utility functions and structures that allow invoking YARA modules directly.

    The utility functions [`invoke`], [`invoke_dyn`] and [`invoke_all`]
    allow leveraging YARA modules for parsing some file formats independently
    of any YARA rule. With these functions you can pass arbitrary data to a
    YARA module and obtain the same data structure that is accessible to YARA
    rules and which you use in your rule conditions.

    This allows external projects to benefit from YARA's file-parsing
    capabilities for their own purposes.

    # Example

    ```rust
    # use yara_x;
    let pe_info = yara_x::mods::invoke::<yara_x::mods::PE>(&[]);
    ```
    */

    /// Data structures defined by the `crx` module.
    ///
    /// The main structure produced by the module is [`crx::Crx`]. The rest
    /// of them are used by one or more fields in the main structure.
    ///
    pub use super::protos::crx;
    /// Data structure returned by the `crx` module.
    pub use super::protos::crx::Crx;
    /// Data structures defined by the `dex` module.
    ///
    /// The main structure produced by the module is [`dex::Dex`]. The rest
    /// of them are used by one or more fields in the main structure.
    ///
    pub use super::protos::dex;
    /// Data structure returned by the `dex` module.
    pub use super::protos::dex::Dex;

    /// Data structures defined by the `dotnet` module.
    ///
    /// The main structure produced by the module is [`dotnet::Dotnet`]. The
    /// rest of them are used by one or more fields in the main structure.
    ///
    pub use super::protos::dotnet;
    /// Data structure returned by the `dotnet` module.
    pub use super::protos::dotnet::Dotnet;

    /// Data structures defined by the `elf` module.
    ///
    /// The main structure produced by the module is [`elf::ELF`]. The rest of
    /// them are used by one or more fields in the main structure.
    ///
    pub use super::protos::elf;
    /// Data structure returned by the `elf` module.
    pub use super::protos::elf::ELF;

    /// Data structures defined by the `lnk` module.
    ///
    /// The main structure produced by the module is [`lnk::Lnk`]. The rest of
    /// them are used by one or more fields in the main structure.
    ///
    pub use super::protos::lnk;
    /// Data structure returned by the `lnk` module.
    pub use super::protos::lnk::Lnk;

    /// Data structures defined by the `macho` module.
    ///
    /// The main structure produced by the module is [`macho::Macho`]. The rest
    /// of them are used by one or more fields in the main structure.
    ///
    pub use super::protos::macho;
    /// Data structure returned by the `macho` module.
    pub use super::protos::macho::Macho;

    /// Data structures defined by the `pe` module.
    ///
    /// The main structure produced by the module is [`pe::PE`]. The rest
    /// of them are used by one or more fields in the main structure.
    ///
    pub use super::protos::pe;
    /// Data structure returned by the `pe` module.
    pub use super::protos::pe::PE;

    /// A data structure containing the data returned by all modules.
    pub use super::protos::mods::Modules;

    /// Invokes a YARA module with arbitrary data.
    ///
    /// <br>
    ///
    /// YARA modules typically parse specific file formats, returning structures
    /// that contain information about the file. These structures are used in YARA
    /// rules for expressing powerful and rich conditions. However, being able to
    /// access this information outside YARA rules can also be beneficial.
    ///
    /// <br>
    ///
    /// This function allows the direct invocation of a YARA module for parsing
    /// arbitrary data. It returns the structure produced by the module, which
    /// depends upon the invoked module. The result will be [`None`] if the
    /// module does not exist, or if it doesn't produce any information for
    /// the input data.
    ///
    /// `T` must be one of the structure types returned by a YARA module, which
    /// are defined in [`crate::mods`], like [`crate::mods::PE`], [`crate::mods::ELF`], etc.
    ///
    /// # Example
    /// ```rust
    /// # use yara_x;
    /// let elf_info = yara_x::mods::invoke::<yara_x::mods::ELF>(&[]);
    /// ```
    pub fn invoke<T: protobuf::MessageFull>(data: &[u8]) -> Option<Box<T>> {
        let module_output = invoke_dyn::<T>(data)?;
        Some(<dyn protobuf::MessageDyn>::downcast_box(module_output).unwrap())
    }

    /// Like [`invoke`], but allows passing metadata to the module.
    pub fn invoke_with_meta<T: protobuf::MessageFull>(
        data: &[u8],
        meta: Option<&[u8]>,
    ) -> Option<Box<T>> {
        let module_output = invoke_with_meta_dyn::<T>(data, meta)?;
        Some(<dyn protobuf::MessageDyn>::downcast_box(module_output).unwrap())
    }

    /// Invokes a YARA module with arbitrary data, returning a dynamic
    /// structure.
    ///
    /// This function is similar to [`invoke`] but its result is a dynamic-
    /// dispatch version of the structure returned by the YARA module.
    pub fn invoke_dyn<T: protobuf::MessageFull>(
        data: &[u8],
    ) -> Option<Box<dyn protobuf::MessageDyn>> {
        invoke_with_meta_dyn::<T>(data, None)
    }

    /// Like [`invoke_dyn`], but allows passing metadata to the module.
    pub fn invoke_with_meta_dyn<T: protobuf::MessageFull>(
        data: &[u8],
        meta: Option<&[u8]>,
    ) -> Option<Box<dyn protobuf::MessageDyn>> {
        let descriptor = T::descriptor();
        let proto_name = descriptor.full_name();
        let (_, module) =
            super::BUILTIN_MODULES.iter().find(|(_, module)| {
                module.root_struct_descriptor.full_name() == proto_name
            })?;

        module.main_fn?(data, meta).ok()
    }

    /// Invokes all YARA modules and returns the data produced by them.
    ///
    /// This function is similar to [`invoke`], but it returns the
    /// information produced by all modules at once.
    ///
    /// # Example
    /// ```rust
    /// # use yara_x;
    /// let modules_output = yara_x::mods::invoke_all(&[]);
    /// ```
    pub fn invoke_all(data: &[u8]) -> Box<Modules> {
        let mut info = Box::new(Modules::new());
        info.pe = protobuf::MessageField(invoke::<PE>(data));
        info.elf = protobuf::MessageField(invoke::<ELF>(data));
        info.dotnet = protobuf::MessageField(invoke::<Dotnet>(data));
        info.macho = protobuf::MessageField(invoke::<Macho>(data));
        info.lnk = protobuf::MessageField(invoke::<Lnk>(data));
        info.crx = protobuf::MessageField(invoke::<Crx>(data));
        info.dex = protobuf::MessageField(invoke::<Dex>(data));
        info
    }

    /// Iterator over built-in module names.
    ///
    /// See the "debug modules" command.
    pub fn module_names() -> impl Iterator<Item = &'static str> {
        use itertools::Itertools;
        super::BUILTIN_MODULES.keys().sorted_by_key(|k| **k).copied()
    }

    /// Returns the definition of the module with the given name.
    pub fn module_definition(name: &str) -> Option<reflect::Struct> {
        use crate::types;
        use std::rc::Rc;
        super::BUILTIN_MODULES
            .get(name)
            .map(|m| reflect::Struct::new(Rc::<types::Struct>::from(m)))
    }

    /// Types that allow for module introspection.
    ///
    /// This API is unstable and not ready for public use.
    #[doc(hidden)]
    pub mod reflect {
        use std::borrow::Cow;
        use std::rc::Rc;

        use crate::types;
        use crate::types::{Map, TypeValue};

        /// Describes a structure or module.
        #[derive(Clone, Debug, PartialEq)]
        pub struct Struct {
            inner: Rc<types::Struct>,
        }

        impl Struct {
            pub(super) fn new(inner: Rc<types::Struct>) -> Self {
                Self { inner }
            }

            /// Returns an iterator over the fields defined in the structure.
            ///
            /// The fields are sorted by name.
            pub fn fields(&self) -> impl Iterator<Item = Field<'_>> + '_ {
                self.inner
                    .fields()
                    .map(|(name, field)| Field::new(name, field))
            }
        }

        /// Describes a function.
        #[derive(Clone, Debug, PartialEq)]
        pub struct Func {
            /// All the existing signatures for this function. A function
            /// can have multiple signatures that differ in their arguments
            /// or return type.
            pub signatures: Vec<FuncSignature>,
        }

        impl From<Rc<types::Func>> for Func {
            fn from(func: Rc<types::Func>) -> Self {
                let mut signatures =
                    Vec::with_capacity(func.signatures().len());

                for signature in func.signatures() {
                    signatures.push(FuncSignature {
                        args: signature.args.iter().map(Type::from).collect(),
                        ret: Type::from(&signature.result),
                        description: signature.description.clone(),
                    });
                }

                Func { signatures }
            }
        }

        /// Describes a function signature.
        #[derive(Clone, Debug, PartialEq)]
        pub struct FuncSignature {
            /// The types of the function arguments.
            pub args: Vec<Type>,
            /// The return type for the function.
            pub ret: Type,
            /// Function's documentation description.
            pub description: Option<Cow<'static, str>>,
        }

        /// Describes a field within a structure or module.
        #[derive(Clone)]
        pub struct Field<'a> {
            name: &'a str,
            struct_field: &'a types::StructField,
        }

        impl<'a> Field<'a> {
            fn new(
                name: &'a str,
                struct_field: &'a types::StructField,
            ) -> Self {
                Self { name, struct_field }
            }

            /// Returns the name of the field.
            pub fn name(&self) -> &'a str {
                self.name
            }

            /// Returns the type of the field.
            pub fn ty(&self) -> Type {
                Type::from(&self.struct_field.type_value)
            }
        }

        /// The type of field, function argument or return value.
        #[derive(Clone, Debug, PartialEq)]
        pub enum Type {
            /// An integer.
            Integer,
            /// A float.
            Float,
            /// A boolean.
            Bool,
            /// A string.
            String,
            /// A regular expression
            Regexp,
            /// A structure.
            Struct(Struct),
            /// An array.
            Array(Box<Type>),
            /// A map.
            Map(Box<Type>, Box<Type>),
            /// A function.
            Func(Func),
        }

        impl From<&TypeValue> for Type {
            fn from(type_value: &TypeValue) -> Self {
                match type_value {
                    TypeValue::Bool { .. } => Type::Bool,
                    TypeValue::Float { .. } => Type::Float,
                    TypeValue::Integer { .. } => Type::Integer,
                    TypeValue::String { .. } => Type::String,
                    TypeValue::Regexp(_) => Type::Regexp,
                    TypeValue::Struct(s) => {
                        Type::Struct(Struct::new(s.clone()))
                    }
                    TypeValue::Array(a) => {
                        Type::Array(Box::new(Type::from(&a.deputy())))
                    }
                    TypeValue::Map(m) => {
                        let key_kind = match **m {
                            Map::IntegerKeys { .. } => Type::Integer,
                            Map::StringKeys { .. } => Type::String,
                        };
                        Type::Map(
                            Box::new(key_kind),
                            Box::new(Type::from(&m.deputy())),
                        )
                    }
                    TypeValue::Func(func) => Type::Func(func.clone().into()),
                    TypeValue::Unknown => unreachable!(),
                }
            }
        }
    }
}

#[cfg(feature = "crypto")]
pub(crate) mod utils;