feat: add compiler profile config in boreal-cli

vthib · Sep 28, 2024 · c3a89c2 · c3a89c2
1 parent 5e0cbda
commit c3a89c2
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 3 deletions.
diff --git a/boreal-cli/src/main.rs b/boreal-cli/src/main.rs
@@ -5,7 +5,7 @@ use std::process::ExitCode;
 use std::thread::JoinHandle;
 use std::time::Duration;
 
-use boreal::compiler::{CompilerBuilder, ExternalValue};
+use boreal::compiler::{CompilerBuilder, CompilerProfile, ExternalValue};
 use boreal::module::{Console, Value as ModuleValue};
 use boreal::scanner::{FragmentedScanMode, ScanError, ScanParams, ScanResult};
 use boreal::{statistics, Compiler, Metadata, MetadataValue, Scanner};
@@ -58,6 +58,13 @@ fn build_command() -> Command {
                 .value_parser(value_parser!(usize))
                 .help("Number of threads to use when scanning directories"),
         )
+        .arg(
+            Arg::new("profile")
+                .long("profile")
+                .value_name("speed|memory")
+                .value_parser(parse_compiler_profile)
+                .help("Profile to use when compiling rules"),
+        )
         .arg(
             Arg::new("rules_file")
                 .value_parser(value_parser!(PathBuf))
@@ -252,14 +259,21 @@ fn main() -> ExitCode {
     let mut scanner = {
         let rules_file: PathBuf = args.remove_one("rules_file").unwrap();
 
-        let no_console_logs = args.get_flag("no_console_logs");
+        let mut builder = CompilerBuilder::new();
+
         // Even if the console logs are disabled, add the module so that rules that use it
         // can still compile properly.
-        let builder = CompilerBuilder::new().add_module(Console::with_callback(move |log| {
+        let no_console_logs = args.get_flag("no_console_logs");
+        builder = builder.add_module(Console::with_callback(move |log| {
             if !no_console_logs {
                 println!("{log}");
             }
         }));
+
+        if let Some(profile) = args.get_one::<CompilerProfile>("profile") {
+            builder = builder.profile(*profile);
+        }
+
         let mut compiler = builder.build();
 
         compiler.set_params(
@@ -481,6 +495,14 @@ fn parse_fragmented_scan_mode(scan_mode: &str) -> Result<FragmentedScanMode, Str
     }
 }
 
+fn parse_compiler_profile(profile: &str) -> Result<CompilerProfile, String> {
+    match profile {
+        "speed" => Ok(CompilerProfile::Speed),
+        "memory" => Ok(CompilerProfile::Memory),
+        _ => Err("invalid value".to_string()),
+    }
+}
+
 #[derive(Clone, Debug)]
 struct ScanOptions {
     print_module_data: bool,

diff --git a/boreal-cli/tests/cli.rs b/boreal-cli/tests/cli.rs
@@ -813,6 +813,22 @@ fn test_invalid_fragmented_scan_mode() {
         .failure();
 }
 
+#[test]
+fn test_invalid_compiler_profile() {
+    cmd()
+        .arg("--profile")
+        .arg("bad_value")
+        .arg("rules.yar")
+        .arg("input")
+        .assert()
+        .stdout("")
+        .stderr(predicate::str::contains(
+            "invalid value 'bad_value' for \
+            '--profile <speed|memory>\': invalid value",
+        ))
+        .failure();
+}
+
 #[test]
 fn test_tags() {
     let rule_file = test_file(