pars_core/operation/
grep.rs

1use std::collections::hash_map::DefaultHasher;
2use std::hash::{Hash, Hasher};
3use std::io::Write;
4use std::path::Path;
5
6use anyhow::Result;
7use colored::{Color, Colorize};
8use regex::Regex;
9use secrecy::ExposeSecret;
10use walkdir::WalkDir;
11
12use crate::config::cli::PrintConfig;
13use crate::pgp::PGPClient;
14use crate::util::fs_util::{get_dir_gpg_id_content, path_to_str};
15use crate::util::tree::string_to_color_opt;
16
17#[derive(Default)]
18pub struct GrepPrintConfig {
19    pub grep_pass_color: Option<Color>,
20    pub grep_match_color: Option<Color>,
21}
22
23impl<CFG: AsRef<PrintConfig>> From<CFG> for GrepPrintConfig {
24    fn from(config: CFG) -> Self {
25        Self {
26            grep_pass_color: string_to_color_opt(&config.as_ref().grep_pass_color),
27            grep_match_color: string_to_color_opt(&config.as_ref().grep_match_color),
28        }
29    }
30}
31
32pub fn grep_stream<O>(
33    pgp_executable: &str,
34    root: &Path,
35    search_str: &str,
36    print_cfg: &GrepPrintConfig,
37    out_stream: &mut O,
38) -> Result<()>
39where
40    O: Write,
41{
42    let mut cache: Vec<(u64, Vec<String>, PGPClient)> = Vec::new();
43    let search_regex = Regex::new(&regex::escape(search_str))?;
44
45    for entry in WalkDir::new(root) {
46        let entry = entry?;
47        if entry.file_type().is_file() && entry.path().extension().unwrap_or_default() == "gpg" {
48            let relative_path = entry.path().strip_prefix(root)?;
49            let relative_path_str = path_to_str(relative_path)?;
50
51            let mut keys_fpr = get_dir_gpg_id_content(root, entry.path())?;
52            keys_fpr.sort();
53
54            let mut hasher = DefaultHasher::new();
55            keys_fpr.hash(&mut hasher);
56            let key_hash = hasher.finish();
57
58            let client: &mut PGPClient = if let Some((_, _, client)) = cache
59                .iter_mut()
60                .find(|(h, cached_keys, _)| *h == key_hash && *cached_keys == keys_fpr)
61            {
62                client
63            } else {
64                let new_client = PGPClient::new(
65                    pgp_executable,
66                    &keys_fpr.iter().map(String::as_str).collect::<Vec<_>>(),
67                )?;
68                cache.push((key_hash, keys_fpr.clone(), new_client));
69                &mut cache.last_mut().unwrap().2
70            };
71
72            let decrypted = client.decrypt_stdin(root, path_to_str(entry.path())?)?;
73            let mut has_matches = false;
74
75            for line in decrypted.expose_secret().lines() {
76                if line.contains(search_str) {
77                    if !has_matches {
78                        if let Some(color) = print_cfg.grep_pass_color {
79                            writeln!(
80                                out_stream,
81                                "{}:",
82                                &relative_path_str[..relative_path_str.len() - 4].color(color)
83                            )?;
84                        } else {
85                            writeln!(
86                                out_stream,
87                                "{}:",
88                                &relative_path_str[..relative_path_str.len() - 4]
89                            )?;
90                        }
91                        has_matches = true;
92                    }
93
94                    let output_line = if let Some(color) = print_cfg.grep_match_color {
95                        search_regex
96                            .replace_all(line, |caps: &regex::Captures| {
97                                caps[0].color(color).to_string()
98                            })
99                            .to_string()
100                    } else {
101                        line.to_string()
102                    };
103                    writeln!(out_stream, "{}", output_line)?;
104                }
105            }
106        }
107    }
108
109    Ok(())
110}
111
112pub fn grep(
113    pgp_executable: &str,
114    root: &Path,
115    search_str: &str,
116    print_cfg: &GrepPrintConfig,
117) -> Result<Vec<String>> {
118    let mut results = Vec::new();
119    let mut cache: Vec<(u64, Vec<String>, PGPClient)> = Vec::new();
120    let search_regex = Regex::new(&regex::escape(search_str))?;
121
122    for entry in WalkDir::new(root) {
123        let entry = entry?;
124        if entry.file_type().is_file() && entry.path().extension().unwrap_or_default() == "gpg" {
125            let relative_path = entry.path().strip_prefix(root)?;
126            let relative_path_str = path_to_str(relative_path)?;
127
128            let mut keys_fpr = get_dir_gpg_id_content(root, entry.path())?;
129            keys_fpr.sort();
130
131            let mut hasher = DefaultHasher::new();
132            keys_fpr.hash(&mut hasher);
133            let key_hash = hasher.finish();
134
135            let client: &mut PGPClient = if let Some((_, _, client)) = cache
136                .iter_mut()
137                .find(|(h, cached_keys, _)| *h == key_hash && *cached_keys == keys_fpr)
138            {
139                client
140            } else {
141                let new_client = PGPClient::new(
142                    pgp_executable,
143                    &keys_fpr.iter().map(String::as_str).collect::<Vec<_>>(),
144                )?;
145                cache.push((key_hash, keys_fpr.clone(), new_client));
146                &mut cache.last_mut().unwrap().2
147            };
148
149            let decrypted = client.decrypt_stdin(root, path_to_str(entry.path())?)?;
150            let matching_lines: Vec<String> = decrypted
151                .expose_secret()
152                .lines()
153                .filter(|line| line.contains(search_str))
154                .map(|line| {
155                    if let Some(color) = print_cfg.grep_match_color {
156                        search_regex
157                            .replace_all(line, |caps: &regex::Captures| {
158                                caps[0].color(color).to_string()
159                            })
160                            .to_string()
161                    } else {
162                        line.to_string()
163                    }
164                })
165                .collect();
166
167            if !matching_lines.is_empty() {
168                if let Some(color) = print_cfg.grep_pass_color {
169                    results.push(format!(
170                        "{}:",
171                        &relative_path_str[..relative_path_str.len() - 4].color(color)
172                    ));
173                } else {
174                    results.push(format!("{}:", &relative_path_str[..relative_path_str.len() - 4]));
175                }
176                results.extend(matching_lines);
177            }
178        }
179    }
180
181    Ok(results)
182}
183
184#[cfg(test)]
185mod tests {
186    use std::path::{self, PathBuf};
187
188    use pretty_assertions::assert_eq;
189    use serial_test::serial;
190    use tempfile::TempDir;
191
192    use super::*;
193    use crate::pgp::key_management::key_gen_batch;
194    use crate::util::defer::cleanup;
195    use crate::util::test_util::*;
196
197    fn setup_test_environment() -> (String, String, TempDir, PathBuf) {
198        let executable = get_test_executable();
199        let email = get_test_email();
200        let (_tmp_dir, root) = gen_unique_temp_dir();
201
202        let file1_content = "INF\n2112112";
203        let file2_content = "Overlord\nNAN";
204
205        let structure: &[(Option<&str>, &[&str])] =
206            &[(Some("dir1"), &[][..]), (Some("dir2"), &[][..])];
207        create_dir_structure(&root, structure);
208
209        key_gen_batch(&executable, &gpg_key_gen_example_batch()).unwrap();
210        let test_client = PGPClient::new(executable.clone(), &[&email]).unwrap();
211        test_client.key_edit_batch(&gpg_key_edit_example_batch()).unwrap();
212
213        test_client.encrypt(file1_content, root.join("dir1/01.gpg").to_str().unwrap()).unwrap();
214        test_client.encrypt(file2_content, root.join("dir2/10.gpg").to_str().unwrap()).unwrap();
215        write_gpg_id(&root, &test_client.get_keys_fpr());
216        (executable, email, _tmp_dir, root)
217    }
218
219    #[test]
220    #[serial]
221    #[ignore = "need run interactively"]
222    fn grep_content_match() {
223        let (executable, email, _tmp_dir, root) = setup_test_environment();
224
225        cleanup!(
226            {
227                let results = grep(&executable, &root, "211", &GrepPrintConfig::default()).unwrap();
228                assert_eq!(results, vec![&format!("dir1{}01:", path::MAIN_SEPARATOR), "2112112"]);
229            },
230            {
231                clean_up_test_key(&executable, &[&email]).unwrap();
232            }
233        );
234    }
235
236    #[test]
237    #[serial]
238    #[ignore = "need run interactively"]
239    fn grep_filename_match() {
240        let (executable, email, _tmp_dir, root) = setup_test_environment();
241
242        cleanup!(
243            {
244                let results =
245                    grep(&executable, &root, "Overlord", &GrepPrintConfig::default()).unwrap();
246                assert_eq!(results, vec![&format!("dir2{}10:", path::MAIN_SEPARATOR), "Overlord"]);
247
248                let results = grep(&executable, &root, "01", &GrepPrintConfig::default()).unwrap();
249                assert_eq!(results, Vec::<String>::new());
250            },
251            {
252                clean_up_test_key(&executable, &[&email]).unwrap();
253            }
254        );
255    }
256
257    #[test]
258    #[serial]
259    #[ignore = "need run interactively"]
260    fn grep_no_matches() {
261        let (executable, email, _tmp_dir, root) = setup_test_environment();
262
263        cleanup!(
264            {
265                let results =
266                    grep(&executable, &root, "nonexistent", &GrepPrintConfig::default()).unwrap();
267                assert!(results.is_empty());
268            },
269            {
270                clean_up_test_key(&executable, &[&email]).unwrap();
271            }
272        );
273    }
274}
pars_core/operation/grep.rs

pars_core/operation/
grep.rs
