use anyhow::{Context, Result, anyhow}; use gethostname::gethostname; use is_terminal::IsTerminal; use std::collections::HashMap; use std::io::{self, Read, Write}; use clap::Command; use clap::error::ErrorKind; use log::debug; use rusqlite::Connection; use std::path::PathBuf; use crate::compression_engine::get_compression_engine; use crate::db::{self}; use crate::meta_plugin::{MetaPlugin, MetaPluginType, get_meta_plugin}; use crate::modes::common::{cmd_args_compression_type, cmd_args_digest_type, cmd_args_meta_plugin_types, get_meta_from_env, store_item_meta_value}; use chrono::Utc; pub fn mode_save( cmd: &mut Command, args: &crate::Args, ids: &mut Vec, tags: &mut Vec, conn: &mut Connection, data_path: PathBuf, ) -> Result<()> { if !ids.is_empty() { cmd.error( ErrorKind::InvalidValue, "ID given, you cannot supply IDs when using --save", ) .exit(); } if tags.is_empty() { tags.push("none".to_string()); } let digest_type = cmd_args_digest_type(cmd, &args); debug!("MAIN: Digest type: {:?}", digest_type); let compression_type = cmd_args_compression_type(cmd, &args); debug!("MAIN: Compression type: {:?}", compression_type); let compression_engine = get_compression_engine(compression_type.clone()).expect("Unable to get compression engine"); // Start with meta plugin types from command line let mut meta_plugin_types: Vec = cmd_args_meta_plugin_types(cmd, &args); debug!("MAIN: Meta plugin types: {:?}", meta_plugin_types); // Convert digest type to meta plugin type and add to the list if needed let digest_meta_plugin_type = match digest_type { crate::meta_plugin::MetaPluginType::DigestSha256 => Some(MetaPluginType::DigestSha256), crate::meta_plugin::MetaPluginType::DigestMd5 => Some(MetaPluginType::DigestMd5), _ => None, }; // Add digest meta plugin to the list if needed if let Some(digest_plugin_type) = digest_meta_plugin_type { if !meta_plugin_types.contains(&digest_plugin_type) { meta_plugin_types.push(digest_plugin_type); } } // Initialize meta_plugins with MetaPlugin instances for each MetaPluginType let mut meta_plugins: Vec> = meta_plugin_types .iter() .map(|meta_plugin_type| get_meta_plugin(meta_plugin_type.clone())) .collect(); // Check for unsupported meta plugins, warn the user, and remove them from the list let mut i = 0; meta_plugins.retain(|meta_plugin| { let is_supported = meta_plugin.is_supported(); if !is_supported { // We need to get the meta name for the warning message // Since we can't mutably borrow meta_plugin here, we create a temporary one let meta_plugin_type = meta_plugin_types[i].clone(); let mut temp_plugin = get_meta_plugin(meta_plugin_type); eprintln!("Warning: Meta plugin '{}' is enabled but not supported on this system", temp_plugin.meta_name()); } i += 1; is_supported }); let mut item = db::Item { id: None, ts: Utc::now(), size: None, compression: compression_type.to_string(), }; let id = db::insert_item(conn, item.clone())?; item.id = Some(id); debug!("MAIN: Added item {:?}", item.clone()); if !args.options.quiet { if std::io::stderr().is_terminal() { let mut t = term::stderr().unwrap(); t.reset().unwrap_or(()); t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "KEEP:").unwrap_or(()); t.reset().unwrap_or(()); write!(t, " New item ").unwrap_or(()); t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "{id}")?; t.reset().unwrap_or(()); write!(t, " tags: ")?; t.attr(term::Attr::Bold).unwrap_or(()); write!(t, "{}", tags.join(" "))?; t.reset().unwrap_or(()); writeln!(t)?; std::io::stderr().flush()?; } else { let mut t = std::io::stderr(); writeln!(t, "KEEP: New item: {} tags: {:?}", id, tags)?; } } db::set_item_tags(conn, item.clone(), tags)?; let mut item_meta: HashMap = get_meta_from_env(); if let Ok(hostname) = gethostname().into_string() { if !item_meta.contains_key("hostname") { item_meta.insert("hostname".to_string(), hostname); } } for item in args.item.meta.iter() { let item = item.clone(); item_meta.insert(item.key, item.value); } for kv in item_meta.iter() { let meta = db::Meta { id: item.id.unwrap(), name: kv.0.to_string(), value: kv.1.to_string(), }; db::store_meta(conn, meta)?; } let mut item_path = data_path.clone(); item_path.push(id.to_string()); let mut stdin = io::stdin().lock(); let mut stdout = io::stdout().lock(); let mut buffer = [0; libc::BUFSIZ as usize]; let mut item_out: Box = compression_engine .create(item_path.clone()) .context(anyhow!( "Unable to write file {:?} using compression {:?}", item_path, compression_type ))?; debug!("MAIN: Starting IO loop"); loop { let n = stdin.read(&mut buffer[..libc::BUFSIZ as usize])?; item.size = match item.size { None => Some(n as i64), Some(prev_n) => Some(prev_n + n as i64), }; if n == 0 { debug!("MAIN: EOF on STDIN"); break; } debug!("MAIN: Loop - {:?} bytes", item.size); stdout.write_all(&buffer[..n])?; item_out.write_all(&buffer[..n])?; for meta_plugin in meta_plugins.iter_mut() { meta_plugin.update(&buffer[..n]); } } debug!("MAIN: Ending IO loop after {:?} bytes", item.size); stdout.flush()?; item_out.flush()?; for meta_plugin in meta_plugins.iter_mut() { let meta_name = meta_plugin.meta_name(); match meta_plugin.finalize() { Ok(meta_value) => { if let Err(e) = store_item_meta_value(conn, item.clone(), meta_name.clone(), meta_value) { eprintln!("Warning: Failed to store meta value for {}: {}", meta_name, e); } } Err(e) => { eprintln!("Warning: Failed to finalize meta plugin {}: {}", meta_name, e); } } } db::update_item(conn, item.clone())?; Ok(()) }