From 0261a8f8532e3f7b7bbd9cf946903d843f88c6bd Mon Sep 17 00:00:00 2001
From: "Victor M. Alvarez" <vmalvarez@virustotal.com>
Date: Fri, 16 Feb 2024 14:17:12 +0100
Subject: [PATCH] feat: add support for the `YRX_MOD_PROTO_DIRS` environment
 variable

With this environment variable you can specify additional directories where to find `.proto` files with module definitions. This variable can contain a single directory path or a comma-separated list of directory path.

The directories will be walked recursively looking for `.proto` files containing module definitions. Notice however that these modules are data-only, they can't contain functions because for adding functions to a YARA module the `.proto` file must be accompanied by Rust source files that implement the functions.
---
 lib/Cargo.toml |  1 +
 lib/build.rs   | 57 +++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/lib/Cargo.toml b/lib/Cargo.toml
index 663ec85e3..6b5bad82e 100644
--- a/lib/Cargo.toml
+++ b/lib/Cargo.toml
@@ -190,6 +190,7 @@ yara-x-proto = { workspace = true }
 lingua = { version = "1.6.0", optional = true, default-features = false, features = ["english", "german", "french", "spanish"] }
 
 [build-dependencies]
+globwalk = { workspace = true }
 protobuf = { workspace = true }
 protobuf-codegen = { workspace = true }
 protobuf-parse = { workspace = true }
diff --git a/lib/build.rs b/lib/build.rs
index 8a09de910..e196a7456 100644
--- a/lib/build.rs
+++ b/lib/build.rs
@@ -19,25 +19,54 @@ fn main() {
     proto_compiler
         .pure()
         .cargo_out_dir("protos")
-        //.out_dir("src/modules/protos")
         .include("../proto/src")
         .include("../proto-yaml/src")
-        .include("src/modules/protos")
         .input("../proto/src/yara.proto")
         .input("../proto-yaml/src/yaml.proto");
 
-    proto_parser
-        .include("../proto/src")
-        .include("../proto-yaml/src")
-        .include("src/modules/protos");
-
-    for entry in fs::read_dir("src/modules/protos").unwrap() {
-        let entry = entry.unwrap();
-        let path = entry.path();
-        if let Some(extension) = path.extension() {
-            if extension == "proto" {
-                proto_compiler.input(&path);
-                proto_parser.input(&path);
+    proto_parser.include("../proto/src").include("../proto-yaml/src");
+
+    // These are the directories where we are going to search for the .proto
+    // files that contain module definitions. Initially, the only directory
+    // is the one included in this repository. But the YRX_MOD_PROTOS_DIRS
+    // environment variable specify additional directories. These directories
+    // are walked recursively looking for .proto files.
+    let mut proto_dirs = vec!["src/modules/protos".to_string()];
+
+    // The YRX_MOD_PROTOS_DIRS environment variable can contain a
+    // comma-separated list of directories with additional module definitions
+    // in `.proto` files. Notice however that these modules are data-only,
+    // they can't contain functions because for adding functions to a YARA
+    // module the `.proto` file must be accompanied by Rust source files that
+    // implement the functions.
+    if let Ok(var) = env::var("YRX_MOD_PROTO_DIRS") {
+        proto_dirs.extend(
+            var.split(',').map(|s| s.to_string()).collect::<Vec<String>>(),
+        );
+    }
+
+    for dir in &proto_dirs {
+        let dir = fs::canonicalize(dir).unwrap();
+        proto_compiler.include(&dir);
+        proto_parser.include(&dir);
+
+        let walker = globwalk::GlobWalkerBuilder::from_patterns(dir, &["**"])
+            .follow_links(true)
+            .build()
+            .unwrap();
+
+        for entry in walker {
+            let entry = entry.unwrap();
+            let path = fs::canonicalize(entry.path()).unwrap();
+
+            if let Some(extension) = path.extension() {
+                if extension == "proto" {
+                    proto_compiler.input(&path);
+                    proto_parser.input(&path);
+                    let dir = path.with_file_name("");
+                    proto_compiler.include(&dir);
+                    proto_parser.include(&dir);
+                }
             }
         }
     }