Magic module

yara-x/Cargo.toml

@@ -24,6 +24,11 @@ test_proto3-module = []
 text-module = [
     "dep:lingua"
 ]
+# The Magic module allows you to identify the type of the file based on 
+# the output of file(1), the standard Unix command. 
+magic-module = [
+    "dep:magic"
+]
 
 # Features that are enabled by default.
 default = [
@@ -62,7 +67,7 @@ yara-x-parser = { workspace = true }
 yara-x-proto = { workspace = true }
 
 lingua = { version = "1.4.0", optional = true, default-features = false, features = ["english", "german", "french", "spanish"] }
-
+magic = { version = "0.13", optional = true, default-features = false }
 
 [build-dependencies]
 protobuf = { workspace = true }

yara-x/src/modules/magic/mod.rs

@@ -0,0 +1,47 @@
+use crate::modules::prelude::*;
+use crate::modules::protos::magic::*;
+
+#[cfg(test)]
+mod tests;
+
+thread_local! {
+    static MAGIC: magic::Cookie = {
+        let cookie = magic::Cookie::open(Default::default())
+            .expect("initialized libmagic");
+        cookie.load::<&str>(&[])
+            .expect("loaded libmagic database");
+        cookie
+    };
+}
+
+#[module_main]
+fn main(_ctx: &ScanContext) -> Magic {
+    // Nothing to do, but we have to return our protobuf
+    Magic::new()
+}
+
+#[module_export(name = "type")]
+fn file_type(ctx: &mut ScanContext) -> Option<RuntimeString> {
+    Some(RuntimeString::from_bytes(ctx, get_type(ctx.scanned_data())))
+}
+
+#[module_export(name = "mime_type")]
+fn mime_type(ctx: &mut ScanContext) -> Option<RuntimeString> {
+    Some(RuntimeString::from_bytes(ctx, get_mime_type(ctx.scanned_data())))
+}
+
+fn get_type(data: &[u8]) -> String {
+    MAGIC
+        .with(|magic| magic.set_flags(Default::default()))
+        .expect("set libmagic options");
+
+    MAGIC.with(|magic| magic.buffer(data)).expect("libmagic didn't break")
+}
+
+fn get_mime_type(data: &[u8]) -> String {
+    MAGIC
+        .with(|magic| magic.set_flags(magic::CookieFlags::MIME_TYPE))
+        .expect("set libmagic options");
+
+    MAGIC.with(|magic| magic.buffer(data)).expect("libmagic didn't break")
+}

yara-x/src/modules/magic/tests.rs

@@ -0,0 +1,39 @@
+use pretty_assertions::assert_eq;
+
+#[test]
+fn get_filetype() {
+    let data = b"Maestro\r";
+    let expected = "RISC OS music file";
+    assert_eq!(expected, crate::modules::magic::get_type(data))
+}
+
+#[test]
+fn get_mimetype() {
+    let expected = "text/plain";
+    assert_eq!(
+        expected,
+        crate::modules::magic::get_mime_type(expected.as_bytes())
+    )
+}
+
+#[test]
+fn e2e_test() {
+    let mut src = String::new();
+    src.push_str(r#"import "magic""#);
+    src.push_str(
+        r#"rule t {condition: magic.type() == "RISC OS music file"}"#,
+    );
+
+    let rules = crate::compiler::Compiler::new()
+        .add_source(src.as_str())
+        .unwrap()
+        .build()
+        .unwrap();
+
+    let data = b"Maestro\r";
+
+    let mut scanner = crate::scanner::Scanner::new(&rules);
+    let results = scanner.scan(data);
+
+    assert_eq!(results.num_matching_rules(), 1);
+}

yara-x/src/modules/modules.rs

@@ -1,6 +1,8 @@
 // File generated automatically by build.rs. Do not edit.
 #[cfg(feature = "text-module")]
 pub mod text;
+#[cfg(feature = "magic-module")]
+pub mod magic;
 #[cfg(feature = "test_proto2-module")]
 pub mod test_proto2;
 #[cfg(feature = "test_proto3-module")]

yara-x/src/modules/protos/magic.proto

@@ -0,0 +1,13 @@
+syntax = "proto2";
+
+import "yara.proto";
+
+option (yara.module_options) = {
+  name : "magic"
+  root_message : "Magic"
+  rust_module : "magic"
+};
+
+message Magic {
+  // This module contains only exported functions, and doesn't return any data
+}