eris: move lua API to standalone module

Easier to work with now, yay.
This commit is contained in:
raf 2025-05-02 05:45:14 +03:00
commit 54f858aee9
Signed by: NotAShelf
GPG key ID: 29D95B64378DB4BF
5 changed files with 297 additions and 127 deletions

24
Cargo.lock generated
View file

@ -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"

View file

@ -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"

View file

@ -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

247
src/lua/mod.rs Normal file
View file

@ -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]");
}
}

View file

@ -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");