diff --git a/Cargo.lock b/Cargo.lock index 3ab4b04..3e3eb40 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 2292db4..c02e277 100644 --- a/Cargo.toml +++ b/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" diff --git a/resources/default_script.lua b/resources/default_script.lua new file mode 100644 index 0000000..90f0966 --- /dev/null +++ b/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" + result = result .. "\n
Server ID: " .. token .. "
" + + return result +end diff --git a/src/lua/mod.rs b/src/lua/mod.rs new file mode 100644 index 0000000..3dbf80b --- /dev/null +++ b/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"); + } + + 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"); + } + + 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") + } + } + } + Err(e) => { + log::warn!("Lua enhance_response function not found: {e}"); + format!("{text}\n") + } + } + } +} + +#[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("")); + } + + #[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("")); + } + + #[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("")); + } + + #[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]"); + } +} diff --git a/src/main.rs b/src/main.rs index b384a9b..154b374 100644 --- a/src/main.rs +++ b/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 = 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" - result = result .. "\n
Server ID: " .. token .. "
" - - 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"); - } - - 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"); - } - - 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") - } - } - } - Err(e) => { - log::warn!("Lua enhance_response function not found: {e}"); - format!("{text}\n") - } - } - } -} - // Find end of HTTP headers -// XXX: I'm sure this could be made less fragile. fn find_header_end(data: &[u8]) -> Option { - 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");