use anyhow::{anyhow, Context, Result}; use is_terminal::IsTerminal; use std::collections::HashMap; use std::io::{self, Read}; use std::str::FromStr; use clap::error::ErrorKind; use clap::Command; use log::debug; use std::path::PathBuf; use rusqlite::Connection; use crate::db::{self}; use crate::digest_engine::{get_digest_engine, DigestType}; use crate::compression_engine::{get_compression_engine, CompressionType}; use crate::modes::common::get_meta_from_env; 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()); } // Get the digest type to use let digest_name = args .item .digest .clone() .unwrap_or(DigestType::Sha256.to_string()); let compression_name = args .item .compression .unwrap_or(crate::compression_engine::default_compression_type().to_string()); use gethostname::gethostname; use std::io::Write; let compression_type_opt = CompressionType::from_str(&compression_name); if compression_type_opt.is_err() { cmd.error( ErrorKind::InvalidValue, format!("Unknown compression type: {}", compression_name), ) .exit(); } let compression_type = compression_type_opt.unwrap(); debug!("MAIN: Compression type: {}", compression_type); // Create a new digest engine let digest_type = DigestType::from_str(&digest_name).unwrap_or(DigestType::Sha256); let mut digest_engine = get_digest_engine(digest_type.clone()); let mut item = db::Item { id: None, ts: Utc::now(), size: None, compression: compression_type.to_string(), digest_type: digest_type.to_string(), digest_value: Some(String::new()), }; 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 compression_engine = get_compression_engine(compression_type.clone()) .expect("Unable to get compression engine"); 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])?; if n == 0 { debug!("MAIN: EOF on STDIN"); break; } stdout.write_all(&buffer[..n])?; item_out.write_all(&buffer[..n])?; digest_engine.update(&buffer[..n]); item.size = match item.size { None => Some(n as i64), Some(prev_n) => Some(prev_n + n as i64), }; } debug!("MAIN: Ending IO loop"); stdout.flush()?; item_out.flush()?; // Finalize the digest and log the result let digest = digest_engine.finalize()?; debug!("DIGEST: {}", digest); item.digest_value = Some(digest); db::update_item(conn, item.clone())?; Ok(()) }