diff --git a/src/meta_plugin.rs b/src/meta_plugin.rs index 626a6da..bd83cad 100644 --- a/src/meta_plugin.rs +++ b/src/meta_plugin.rs @@ -1,12 +1,12 @@ use anyhow::Result; use std::io; use std::io::Write; +use rusqlite::Connection; pub mod program; pub mod digest; pub mod system; - use crate::meta_plugin::program::MetaPluginProgram; use crate::meta_plugin::digest::{DigestSha256MetaPlugin, ReadTimeMetaPlugin, ReadRateMetaPlugin}; use crate::meta_plugin::system::{CwdMetaPlugin, BinaryMetaPlugin, UidMetaPlugin, UserMetaPlugin, GidMetaPlugin, GroupMetaPlugin, ShellMetaPlugin, ShellPidMetaPlugin, KeepPidMetaPlugin, HostnameMetaPlugin, FullHostnameMetaPlugin}; @@ -57,6 +57,22 @@ pub trait MetaPlugin { fn program_info(&self) -> Option<(&str, Vec<&str>)> { None } + + // Initialize with database connection + fn initialize(&mut self, _conn: &Connection, _item_id: i64) -> Result<()> { + Ok(()) + } + + // Save metadata to database + fn save_meta(&self, conn: &Connection, item_id: i64, value: String) -> Result<()> { + let meta = crate::db::Meta { + id: item_id, + name: self.meta_name(), + value, + }; + crate::db::store_meta(conn, meta)?; + Ok(()) + } } pub fn get_meta_plugin(meta_plugin_type: MetaPluginType) -> Box { @@ -83,4 +99,3 @@ pub fn get_meta_plugin(meta_plugin_type: MetaPluginType) -> Box MetaPluginType::FullHostname => Box::new(FullHostnameMetaPlugin::new()), } } - diff --git a/src/meta_plugin/digest.rs b/src/meta_plugin/digest.rs index 59142ef..a08f463 100644 --- a/src/meta_plugin/digest.rs +++ b/src/meta_plugin/digest.rs @@ -3,6 +3,7 @@ use sha2::{Digest, Sha256}; use std::io; use std::io::Write; use std::time::Instant; +use rusqlite::Connection; use crate::meta_plugin::MetaPlugin; diff --git a/src/meta_plugin/program.rs b/src/meta_plugin/program.rs index 3ff2424..9f674a3 100644 --- a/src/meta_plugin/program.rs +++ b/src/meta_plugin/program.rs @@ -5,6 +5,7 @@ use std::io; use std::io::Write; use std::process::{Command, Stdio}; use which::which; +use rusqlite::Connection; use crate::meta_plugin::MetaPlugin; @@ -127,5 +128,4 @@ impl MetaPlugin for MetaPluginProgram { None } } - } diff --git a/src/meta_plugin/system.rs b/src/meta_plugin/system.rs index 9d59b65..d5d5ece 100644 --- a/src/meta_plugin/system.rs +++ b/src/meta_plugin/system.rs @@ -7,6 +7,7 @@ use std::io::Write; use std::env; use std::process; use uzers::{get_current_uid, get_current_gid, get_current_username, get_current_groupname}; +use rusqlite::Connection; use crate::common::is_binary::is_binary; use crate::meta_plugin::MetaPlugin; @@ -21,6 +22,8 @@ pub struct BinaryMetaPlugin { meta_name: String, buffer: Vec, max_buffer_size: usize, + item_id: Option, + conn: Option<*mut Connection>, } impl BinaryMetaPlugin { @@ -29,6 +32,8 @@ impl BinaryMetaPlugin { meta_name: "binary".to_string(), buffer: Vec::new(), max_buffer_size: 4096, // 4KB + item_id: None, + conn: None, } } @@ -60,6 +65,13 @@ impl MetaPlugin for BinaryMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + self.item_id = Some(item_id); + // Store raw pointer to connection - unsafe but necessary for this design + self.conn = Some(conn as *const _ as *mut Connection); + Ok(()) + } } impl CwdMetaPlugin { @@ -93,6 +105,15 @@ impl MetaPlugin for CwdMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let cwd = match env::current_dir() { + Ok(path) => path.to_string_lossy().to_string(), + Err(_) => "unknown".to_string(), + }; + self.save_meta(conn, item_id, cwd)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -128,6 +149,12 @@ impl MetaPlugin for UidMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let uid = get_current_uid().to_string(); + self.save_meta(conn, item_id, uid)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -166,6 +193,15 @@ impl MetaPlugin for UserMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let user = match get_current_username() { + Some(username) => username.to_string_lossy().to_string(), + None => "unknown".to_string(), + }; + self.save_meta(conn, item_id, user)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -201,6 +237,12 @@ impl MetaPlugin for GidMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let gid = get_current_gid().to_string(); + self.save_meta(conn, item_id, gid)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -239,6 +281,15 @@ impl MetaPlugin for GroupMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let group = match get_current_groupname() { + Some(groupname) => groupname.to_string_lossy().to_string(), + None => "unknown".to_string(), + }; + self.save_meta(conn, item_id, group)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -277,6 +328,15 @@ impl MetaPlugin for ShellMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let shell = match env::var("SHELL") { + Ok(shell) => shell, + Err(_) => "unknown".to_string(), + }; + self.save_meta(conn, item_id, shell)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -315,6 +375,15 @@ impl MetaPlugin for ShellPidMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let pid = match env::var("PPID") { + Ok(ppid) => ppid, + Err(_) => process::id().to_string(), + }; + self.save_meta(conn, item_id, pid)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -350,6 +419,12 @@ impl MetaPlugin for KeepPidMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let pid = process::id().to_string(); + self.save_meta(conn, item_id, pid)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -388,6 +463,15 @@ impl MetaPlugin for HostnameMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + let hostname = match gethostname().into_string() { + Ok(hostname) => hostname, + Err(_) => "unknown".to_string(), + }; + self.save_meta(conn, item_id, hostname)?; + Ok(()) + } } #[derive(Debug, Clone, Default)] @@ -444,5 +528,31 @@ impl MetaPlugin for FullHostnameMetaPlugin { fn meta_name(&mut self) -> String { self.meta_name.clone() } + + fn initialize(&mut self, conn: &Connection, item_id: i64) -> Result<()> { + // Try to get the FQDN through reverse DNS lookup + let hostname = match local_ip() { + Ok(my_local_ip) => { + match lookup_addr(&my_local_ip) { + Ok(hostname) => hostname, + Err(_) => { + // Fall back to regular hostname if reverse DNS fails + match gethostname().into_string() { + Ok(hostname) => hostname, + Err(_) => "unknown".to_string(), + } + } + } + } + Err(_) => { + // Fall back to regular hostname if we can't get local IP + match gethostname().into_string() { + Ok(hostname) => hostname, + Err(_) => "unknown".to_string(), + } + } + }; + self.save_meta(conn, item_id, hostname)?; + Ok(()) + } } - diff --git a/src/modes/save.rs b/src/modes/save.rs index ec18c0b..42380a4 100644 --- a/src/modes/save.rs +++ b/src/modes/save.rs @@ -150,6 +150,7 @@ fn process_input_stream( data_path: &std::path::PathBuf, item_id: i64, meta_plugins: &mut Vec>, + conn: &rusqlite::Connection, ) -> Result<(Box, crate::db::Item), anyhow::Error> { let mut item = crate::db::Item { id: Some(item_id), @@ -194,39 +195,26 @@ fn process_input_stream( } debug!("MAIN: Ending IO loop after {:?} bytes", item.size); + // Finalize meta plugins and save their metadata + for meta_plugin in meta_plugins.iter_mut() { + match meta_plugin.finalize() { + Ok(value) => { + if let Err(e) = meta_plugin.save_meta(conn, item_id, value) { + eprintln!("Warning: Failed to save meta value: {}", e); + } + } + Err(e) => { + eprintln!("Warning: Failed to finalize meta plugin: {}", e); + } + } + } + stdout.flush()?; item_out.flush()?; Ok((item_out, item)) } -fn finalize_meta_plugins( - conn: &rusqlite::Connection, - meta_plugins: &mut Vec>, - item: &crate::db::Item, -) -> Result<(), anyhow::Error> { - for meta_plugin in meta_plugins.iter_mut() { - let meta_name = meta_plugin.meta_name(); - - match meta_plugin.finalize() { - Ok(meta_value) => { - let meta = crate::db::Meta { - id: item.id.ok_or_else(|| anyhow!("Item missing ID"))?, - name: meta_name.clone(), - value: meta_value, - }; - if let Err(e) = crate::db::store_meta(conn, meta) { - eprintln!("Warning: Failed to store meta value for {}: {}", meta_name, e); - } - } - Err(e) => { - eprintln!("Warning: Failed to finalize meta plugin {}: {}", meta_name, e); - } - } - } - Ok(()) -} - pub fn mode_save( cmd: &mut Command, settings: &config::Settings, @@ -257,17 +245,24 @@ pub fn mode_save( crate::db::store_meta(conn, meta)?; } + // Initialize meta plugins with database connection + for meta_plugin in meta_plugins.iter_mut() { + if let Err(e) = meta_plugin.initialize(conn, item_id) { + eprintln!("Warning: Failed to initialize meta plugin: {}", e); + } + } + let (_item_out, processed_item) = process_input_stream( &compression_engine, &data_path, item_id, &mut meta_plugins, + conn, )?; item.size = processed_item.size; item.compression = compression_type.to_string(); - finalize_meta_plugins(conn, &mut meta_plugins, &item)?; crate::db::update_item(conn, item.clone())?; Ok(())