eris: move lua API to standalone module

Easier to work with now, yay.

Cargo.lock

@@ -636,6 +636,7 @@ dependencies = [
  "rlua",
  "serde",
  "serde_json",
+ "tempfile",
  "tokio",
 ]
 
@@ -646,9 +647,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "976dd42dc7e85965fe702eb8164f21f450704bdde31faefd6471dba214cb594e"
 dependencies = [
  "libc",
- "windows-sys 0.52.0",
+ "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "fastrand"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
+
 [[package]]
 name = "flate2"
 version = "1.1.1"
@@ -1580,7 +1587,7 @@ dependencies = [
  "errno",
  "libc",
  "linux-raw-sys",
- "windows-sys 0.52.0",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]
@@ -1740,6 +1747,19 @@ dependencies = [
  "syn 2.0.101",
 ]
 
+[[package]]
+name = "tempfile"
+version = "3.19.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7437ac7763b9b123ccf33c338a5cc1bac6f69b45a136c19bdd8a65e3916435bf"
+dependencies = [
+ "fastrand",
+ "getrandom",
+ "once_cell",
+ "rustix",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.69"

Cargo.toml

@@ -19,3 +19,4 @@ serde_json = "1.0.96"
 tokio = { version = "1.28.0", features = ["full"] }
 log = "0.4.27"
 env_logger = "0.11.8"
+tempfile = "3.19.1"

resources/default_script.lua

@@ -0,0 +1,17 @@
+function generate_honeytoken(token)
+  local token_types = { "API_KEY", "AUTH_TOKEN", "SESSION_ID", "SECRET_KEY" }
+  local prefix = token_types[math.random(#token_types)]
+  local suffix = string.format("%08x", math.random(0xffffff))
+  return prefix .. "_" .. token .. "_" .. suffix
+end
+
+function enhance_response(text, response_type, path, token)
+  local result = text
+  local honeytoken = generate_honeytoken(token)
+
+  -- Add some fake sensitive data
+  result = result .. "\n<!-- DEBUG: " .. honeytoken .. " -->"
+  result = result .. "\n<div style='display:none'>Server ID: " .. token .. "</div>"
+
+  return result
+end

src/lua/mod.rs

@@ -0,0 +1,247 @@
+use rlua::{Function, Lua};
+use std::fs;
+use std::path::Path;
+
+pub struct ScriptManager {
+    script_content: String,
+    scripts_loaded: bool,
+}
+
+impl ScriptManager {
+    pub fn new(scripts_dir: &str) -> Self {
+        let mut script_content = String::new();
+        let mut scripts_loaded = false;
+
+        // Try to load scripts from directory
+        let script_dir = Path::new(scripts_dir);
+        if script_dir.exists() {
+            log::debug!("Loading Lua scripts from directory: {scripts_dir}");
+            if let Ok(entries) = fs::read_dir(script_dir) {
+                for entry in entries.filter_map(Result::ok) {
+                    let path = entry.path();
+                    if path.extension().and_then(|ext| ext.to_str()) == Some("lua") {
+                        if let Ok(content) = fs::read_to_string(&path) {
+                            log::debug!("Loaded Lua script: {}", path.display());
+                            script_content.push_str(&content);
+                            script_content.push('\n');
+                            scripts_loaded = true;
+                        } else {
+                            log::warn!("Failed to read Lua script: {}", path.display());
+                        }
+                    }
+                }
+            }
+        } else {
+            log::warn!("Lua scripts directory does not exist: {scripts_dir}");
+        }
+
+        // If no scripts were loaded, use a default script
+        if !scripts_loaded {
+            log::info!("No Lua scripts found, loading default scripts");
+            script_content = include_str!("../../resources/default_script.lua").to_string();
+            scripts_loaded = true;
+        }
+
+        Self {
+            script_content,
+            scripts_loaded,
+        }
+    }
+
+    // For testing only
+    #[cfg(test)]
+    pub fn with_content(content: &str) -> Self {
+        Self {
+            script_content: content.to_string(),
+            scripts_loaded: true,
+        }
+    }
+
+    pub fn expand_response(
+        &self,
+        text: &str,
+        response_type: &str,
+        path: &str,
+        token: &str,
+    ) -> String {
+        if !self.scripts_loaded {
+            return format!("{text}\n<!-- Token: {token} -->");
+        }
+
+        let lua = Lua::new();
+        if let Err(e) = lua.load(&self.script_content).exec() {
+            log::warn!("Error loading Lua script: {e}");
+            return format!("{text}\n<!-- Error: Failed to load Lua script -->");
+        }
+
+        let globals = lua.globals();
+        match globals.get::<_, Function>("enhance_response") {
+            Ok(enhance_func) => {
+                match enhance_func.call::<_, String>((text, response_type, path, token)) {
+                    Ok(result) => result,
+                    Err(e) => {
+                        log::warn!("Error calling Lua function enhance_response: {e}");
+                        format!("{text}\n<!-- Error calling Lua enhance_response -->")
+                    }
+                }
+            }
+            Err(e) => {
+                log::warn!("Lua enhance_response function not found: {e}");
+                format!("{text}\n<!-- Lua enhance_response function not found -->")
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use std::io::Write;
+    use tempfile::TempDir;
+
+    #[test]
+    fn test_script_manager_default_script() {
+        let script_manager = ScriptManager::new("/nonexistent_directory");
+        assert!(script_manager.scripts_loaded);
+        assert!(
+            script_manager
+                .script_content
+                .contains("generate_honeytoken")
+        );
+        assert!(script_manager.script_content.contains("enhance_response"));
+    }
+
+    #[test]
+    fn test_script_manager_custom_scripts() {
+        let temp_dir = TempDir::new().unwrap();
+        let script_dir = temp_dir.path().to_str().unwrap();
+
+        // Create a test script
+        let script_path = temp_dir.path().join("test_script.lua");
+        let mut file = fs::File::create(&script_path).unwrap();
+        writeln!(
+            file,
+            "function enhance_response(text, response_type, path, token)"
+        )
+        .unwrap();
+        writeln!(file, "  return text .. ' - Enhanced with token: ' .. token").unwrap();
+        writeln!(file, "end").unwrap();
+
+        let script_manager = ScriptManager::new(script_dir);
+        assert!(script_manager.scripts_loaded);
+        assert!(
+            !script_manager
+                .script_content
+                .contains("generate_honeytoken")
+        ); // Default script not loaded
+        assert!(
+            script_manager
+                .script_content
+                .contains("Enhanced with token")
+        );
+    }
+
+    #[test]
+    fn test_expand_response_successful() {
+        let lua_code = r#"
+        function enhance_response(text, response_type, path, token)
+            return text .. " | Type: " .. response_type .. " | Path: " .. path .. " | Token: " .. token
+        end
+        "#;
+
+        let script_manager = ScriptManager::with_content(lua_code);
+        let result =
+            script_manager.expand_response("Test content", "test_type", "/test/path", "12345");
+
+        assert_eq!(
+            result,
+            "Test content | Type: test_type | Path: /test/path | Token: 12345"
+        );
+    }
+
+    #[test]
+    fn test_expand_response_syntax_error() {
+        let lua_code = r#"
+        function enhance_response(text, response_type, path, token)
+            This is an invalid Lua syntax
+            return "Something"
+        end
+        "#;
+
+        let script_manager = ScriptManager::with_content(lua_code);
+        let result =
+            script_manager.expand_response("Test content", "test_type", "/test/path", "12345");
+
+        assert!(result.contains("Test content"));
+        assert!(result.contains("<!-- Error: Failed to load Lua script -->"));
+    }
+
+    #[test]
+    fn test_expand_response_runtime_error() {
+        let lua_code = r#"
+        function enhance_response(text, response_type, path, token)
+            -- This will cause a runtime error
+            return nonexistent_variable
+        end
+        "#;
+
+        let script_manager = ScriptManager::with_content(lua_code);
+        let result =
+            script_manager.expand_response("Test content", "test_type", "/test/path", "12345");
+
+        assert!(result.contains("Test content"));
+        assert!(result.contains("<!-- Error calling Lua enhance_response -->"));
+    }
+
+    #[test]
+    fn test_expand_response_missing_function() {
+        let lua_code = r#"
+        -- This script doesn't define enhance_response function
+        function some_other_function()
+            return "Hello, world!"
+        end
+        "#;
+
+        let script_manager = ScriptManager::with_content(lua_code);
+        let result =
+            script_manager.expand_response("Test content", "test_type", "/test/path", "12345");
+
+        assert!(result.contains("Test content"));
+        assert!(result.contains("<!-- Lua enhance_response function not found -->"));
+    }
+
+    #[test]
+    fn test_expand_response_multiple_scripts() {
+        let temp_dir = TempDir::new().unwrap();
+        let script_dir = temp_dir.path().to_str().unwrap();
+
+        // Create first script with helper function
+        let script1_path = temp_dir.path().join("01_helpers.lua");
+        let mut file1 = fs::File::create(script1_path).unwrap();
+        writeln!(file1, "function create_prefix(token)").unwrap();
+        writeln!(file1, "  return 'PREFIX_' .. token").unwrap();
+        writeln!(file1, "end").unwrap();
+
+        // Create second script that uses the helper
+        let script2_path = temp_dir.path().join("02_responder.lua");
+        let mut file2 = fs::File::create(script2_path).unwrap();
+        writeln!(
+            file2,
+            "function enhance_response(text, response_type, path, token)"
+        )
+        .unwrap();
+        writeln!(
+            file2,
+            "  return text .. ' [' .. create_prefix(token) .. ']'"
+        )
+        .unwrap();
+        writeln!(file2, "end").unwrap();
+
+        let script_manager = ScriptManager::new(script_dir);
+        let result =
+            script_manager.expand_response("Test content", "test_type", "/test/path", "12345");
+
+        assert_eq!(result, "Test content [PREFIX_12345]");
+    }
+}

src/main.rs

@@ -1,7 +1,6 @@
 use actix_web::{App, HttpResponse, HttpServer, web};
 use clap::Parser;
 use ipnetwork::IpNetwork;
-use rlua::{Function, Lua};
 use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use std::env;
@@ -17,9 +16,11 @@ use tokio::process::Command;
 use tokio::sync::RwLock;
 use tokio::time::sleep;
 
+mod lua;
 mod markov;
 mod metrics;
 
+use lua::ScriptManager;
 use markov::MarkovGenerator;
 use metrics::{
     ACTIVE_CONNECTIONS, BLOCKED_IPS, HITS_COUNTER, PATH_HITS, UA_HITS, metrics_handler,
@@ -356,10 +357,11 @@ impl BotState {
         }
 
         let hit_cache_file = format!("{}/hit_counters.json", self.cache_dir);
-        let mut hit_map = HashMap::new();
-        for (ip, count) in &self.hits {
-            hit_map.insert(ip.to_string(), *count);
-        }
+        let hit_map: HashMap<String, u32> = self
+            .hits
+            .iter()
+            .map(|(ip, count)| (ip.to_string(), *count))
+            .collect();
 
         match fs::File::create(&hit_cache_file) {
             Ok(file) => {
@@ -376,116 +378,11 @@ impl BotState {
     }
 }
 
-// Lua scripts for response generation and customization
-struct ScriptManager {
-    script_content: String,
-    scripts_loaded: bool,
-}
-
-impl ScriptManager {
-    fn new(scripts_dir: &str) -> Self {
-        let mut script_content = String::new();
-        let mut scripts_loaded = false;
-
-        // Try to load scripts from directory
-        let script_dir = Path::new(scripts_dir);
-        if script_dir.exists() {
-            log::debug!("Loading Lua scripts from directory: {scripts_dir}");
-            if let Ok(entries) = fs::read_dir(script_dir) {
-                for entry in entries {
-                    if let Ok(entry) = entry {
-                        let path = entry.path();
-                        if path.extension().and_then(|ext| ext.to_str()) == Some("lua") {
-                            if let Ok(content) = fs::read_to_string(&path) {
-                                log::debug!("Loaded Lua script: {}", path.display());
-                                script_content.push_str(&content);
-                                script_content.push('\n');
-                                scripts_loaded = true;
-                            } else {
-                                log::warn!("Failed to read Lua script: {}", path.display());
-                            }
-                        }
-                    }
-                }
-            }
-        } else {
-            log::warn!("Lua scripts directory does not exist: {scripts_dir}");
-        }
-
-        // If no scripts were loaded, use a default script
-        if !scripts_loaded {
-            log::info!("No Lua scripts found, loading default scripts");
-            script_content = r#"
-            function generate_honeytoken(token)
-                local token_types = {"API_KEY", "AUTH_TOKEN", "SESSION_ID", "SECRET_KEY"}
-                local prefix = token_types[math.random(#token_types)]
-                local suffix = string.format("%08x", math.random(0xffffff))
-                return prefix .. "_" .. token .. "_" .. suffix
-            end
-
-            function enhance_response(text, response_type, path, token)
-                local result = text
-                local honeytoken = generate_honeytoken(token)
-
-                -- Add some fake sensitive data
-                result = result .. "\n<!-- DEBUG: " .. honeytoken .. " -->"
-                result = result .. "\n<div style='display:none'>Server ID: " .. token .. "</div>"
-
-                return result
-            end
-            "#
-            .to_string();
-            scripts_loaded = true;
-        }
-
-        Self {
-            script_content,
-            scripts_loaded,
-        }
-    }
-
-    // Lua is a powerful configuration language we can use to expand functionality of
-    // Eris, e.g., with fake tokens or honeytrap content.
-    fn expand_response(&self, text: &str, response_type: &str, path: &str, token: &str) -> String {
-        if !self.scripts_loaded {
-            return format!("{text}\n<!-- Token: {token} -->");
-        }
-
-        let lua = Lua::new();
-        if let Err(e) = lua.load(&self.script_content).exec() {
-            log::warn!("Error loading Lua script: {e}");
-            return format!("{text}\n<!-- Error: Failed to load Lua script -->");
-        }
-
-        let globals = lua.globals();
-        match globals.get::<_, Function>("enhance_response") {
-            Ok(enhance_func) => {
-                match enhance_func.call::<_, String>((text, response_type, path, token)) {
-                    Ok(result) => result,
-                    Err(e) => {
-                        log::warn!("Error calling Lua function enhance_response: {e}");
-                        format!("{text}\n<!-- Error calling Lua enhance_response -->")
-                    }
-                }
-            }
-            Err(e) => {
-                log::warn!("Lua enhance_response function not found: {e}");
-                format!("{text}\n<!-- Lua enhance_response function not found -->")
-            }
-        }
-    }
-}
-
 // Find end of HTTP headers
-// XXX: I'm sure this could be made less fragile.
 fn find_header_end(data: &[u8]) -> Option<usize> {
-    for i in 0..data.len().saturating_sub(3) {
-        if data[i] == b'\r' && data[i + 1] == b'\n' && data[i + 2] == b'\r' && data[i + 3] == b'\n'
-        {
-            return Some(i + 4);
-        }
-    }
-    None
+    data.windows(4)
+        .position(|window| window == b"\r\n\r\n")
+        .map(|pos| pos + 4)
 }
 
 // Extract path from raw request data
@@ -1379,18 +1276,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_script_manager_default_script() {
-        let script_manager = ScriptManager::new("/nonexistent_directory");
-        assert!(script_manager.scripts_loaded);
-        assert!(
-            script_manager
-                .script_content
-                .contains("generate_honeytoken")
-        );
-        assert!(script_manager.script_content.contains("enhance_response"));
-    }
-
     #[tokio::test]
     async fn test_bot_state() {
         let state = BotState::new("/tmp/eris_test", "/tmp/eris_test_cache");