diff --git a/src/args.rs b/src/args.rs index 7495c88..05daaaa 100644 --- a/src/args.rs +++ b/src/args.rs @@ -39,6 +39,9 @@ pub enum CliCommand { #[command(flatten)] redacter_args: Option, + + #[arg(long, help = "Filter by name using glob patterns such as *.txt", value_parser = CliCommand::parse_key_val::)] + mime_override: Vec<(mime::Mime, globset::Glob)>, }, #[command(about = "List files in the source")] Ls { @@ -57,6 +60,23 @@ pub enum CliCommand { }, } +impl CliCommand { + fn parse_key_val( + s: &str, + ) -> Result<(T, U), Box> + where + T: std::str::FromStr, + T::Err: std::error::Error + Send + Sync + 'static, + U: std::str::FromStr, + U::Err: std::error::Error + Send + Sync + 'static, + { + let pos = s + .find('=') + .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?; + Ok((s[..pos].parse()?, s[pos + 1..].parse()?)) + } +} + #[derive(ValueEnum, Debug, Clone)] pub enum RedacterType { GcpDlp, diff --git a/src/commands/copy_command.rs b/src/commands/copy_command.rs index 9588b3f..9bb93d8 100644 --- a/src/commands/copy_command.rs +++ b/src/commands/copy_command.rs @@ -1,7 +1,8 @@ use crate::errors::AppError; use crate::file_converters::FileConverters; use crate::filesystems::{ - DetectFileSystem, FileMatcher, FileMatcherResult, FileSystemConnection, FileSystemRef, + DetectFileSystem, FileMatcher, FileMatcherResult, FileMimeOverride, FileSystemConnection, + FileSystemRef, }; use crate::redacters::{ RedactSupportedOptions, Redacter, RedacterBaseOptions, RedacterOptions, Redacters, @@ -23,15 +24,21 @@ pub struct CopyCommandResult { #[derive(Debug, Clone)] pub struct CopyCommandOptions { pub file_matcher: FileMatcher, + pub file_mime_override: FileMimeOverride, } impl CopyCommandOptions { - pub fn new(filename_filter: Option, max_size_limit: Option) -> Self { + pub fn new( + filename_filter: Option, + max_size_limit: Option, + mime_override: Vec<(mime::Mime, globset::Glob)>, + ) -> Self { let filename_matcher = filename_filter .as_ref() .map(|filter| filter.compile_matcher()); CopyCommandOptions { file_matcher: FileMatcher::new(filename_matcher, max_size_limit), + file_mime_override: FileMimeOverride::new(mime_override), } } } @@ -119,7 +126,11 @@ pub async fn command_copy( } bar.println("Copying directory and listing source files..."); let source_files_result = source_fs.list_files(Some(&options.file_matcher)).await?; - let source_files = source_files_result.files; + let source_files: Vec = source_files_result + .files + .into_iter() + .map(|f| options.file_mime_override.override_for_file_ref(f)) + .collect(); let files_found = source_files.len(); let files_total_size: u64 = source_files .iter() diff --git a/src/filesystems/file_mime_override.rs b/src/filesystems/file_mime_override.rs new file mode 100644 index 0000000..6efafb4 --- /dev/null +++ b/src/filesystems/file_mime_override.rs @@ -0,0 +1,32 @@ +use crate::filesystems::FileSystemRef; +use rvstruct::ValueStruct; + +#[derive(Debug, Clone)] +pub struct FileMimeOverride { + mime_override: Vec<(mime::Mime, globset::GlobMatcher)>, +} + +impl FileMimeOverride { + pub fn new(mime_override: Vec<(mime::Mime, globset::Glob)>) -> Self { + Self { + mime_override: mime_override + .into_iter() + .map(|(set_mime, glob)| (set_mime, glob.compile_matcher())) + .collect(), + } + } + + pub fn override_for_file_ref(&self, file_ref: FileSystemRef) -> FileSystemRef { + match self + .mime_override + .iter() + .find(|(_, matcher)| matcher.is_match(file_ref.relative_path.value().as_str())) + { + Some((set_mime, _)) => FileSystemRef { + media_type: Some(set_mime.clone()), + ..file_ref + }, + None => file_ref, + } + } +} diff --git a/src/filesystems/mod.rs b/src/filesystems/mod.rs index 6d3f4cb..cee1aa6 100644 --- a/src/filesystems/mod.rs +++ b/src/filesystems/mod.rs @@ -9,15 +9,19 @@ use gcloud_sdk::prost::bytes::Bytes; use mime::Mime; use rvstruct::ValueStruct; +mod file_matcher; +pub use file_matcher::*; + +mod file_mime_override; +pub use file_mime_override::*; + mod aws_s3; mod gcs; mod local; mod zip; -mod file_matcher; use crate::filesystems::aws_s3::AwsS3FileSystem; use crate::reporter::AppReporter; -pub use file_matcher::*; #[derive(Debug, Clone, ValueStruct)] pub struct RelativeFilePath(pub String); diff --git a/src/main.rs b/src/main.rs index 10d70ed..66ae7e7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -70,8 +70,9 @@ async fn handle_args(cli: CliArgs, term: &Term) -> AppResult<()> { max_size_limit, filename_filter, redacter_args, + mime_override, } => { - let options = CopyCommandOptions::new(filename_filter, max_size_limit); + let options = CopyCommandOptions::new(filename_filter, max_size_limit, mime_override); let copy_result = command_copy( term, &source,