1use super::Config;
2use crate::{Input, InputFile, InputFileType, Msg, Result, remove_xml_tags};
3use either::Either;
4use log::{error, info, warn};
5use quick_xml::{Reader, Writer};
6use rayon::prelude::*;
7use std::{
8 borrow::Cow,
9 fs::{self, File},
10 io::{BufRead, BufReader, BufWriter, Write},
11 iter,
12 path::{Component, Path, PathBuf},
13 sync::{Mutex, Once},
14};
15use walkdir::WalkDir;
16use zip::{ZipArchive, ZipWriter, write::SimpleFileOptions};
17
18impl Config {
19 pub fn run(&self) -> Result<()> {
21 let is_found_any = Once::new();
22 let it = Mutex::new(self.subdirs_src_iter());
23
24 let src_dests: Vec<(InputFile, Dest)> = (0..self.jobs)
25 .into_par_iter()
26 .map(|_| job_src_dests(self, &is_found_any, &it))
27 .collect::<std::result::Result<Vec<_>, String>>()
28 .map(|vecs| vecs.into_iter().flatten().collect())?;
29
30 if !src_dests.is_empty() {
31 force_overwrites(src_dests);
32 }
33
34 if !is_found_any.is_completed() {
35 if let Input::Dir(d) = &self.input {
36 return Err(
37 format!("{} '{}'", Msg::NotFoundAnyBookInDirectory, d.display()).into(),
38 );
39 }
40 }
41
42 Ok(())
43 }
44
45 fn subdirs_src_iter(&self) -> impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)> {
46 match &self.input {
47 Input::Dir(d) if self.recursive != 0 => {
48 let it = WalkDir::new(d)
49 .max_depth(self.recursive as usize)
50 .into_iter()
51 .filter_map(|e| e.ok())
52 .filter(|e| e.file_type().is_dir())
53 .filter(|e| !e.path().starts_with(&*self.output.dir))
54 .map(move |e| {
55 let subdirs: Vec<PathBuf> = e
56 .path()
57 .strip_prefix(d)
58 .unwrap_or(Path::new(""))
59 .components()
60 .filter_map(|x| match x {
61 Component::Normal(x) => Some(PathBuf::from(x)),
62 _ => None,
63 })
64 .collect();
65
66 (Some(subdirs), Input::Dir(e.into_path().into()))
67 })
68 .flat_map(|(subdirs, d)| iter::repeat(subdirs).zip(d.iter()));
69
70 Either::Left(it)
71 }
72 _ => Either::Right(iter::repeat(None).zip(self.input.iter())),
73 }
74 }
75}
76
77fn job_src_dests(
78 cfg: &Config,
79 is_found_any: &Once,
80 it: &Mutex<impl Iterator<Item = (Option<Vec<PathBuf>>, InputFile)>>,
81) -> std::result::Result<Vec<(InputFile, Dest)>, String> {
82 let mut src_dests: Vec<(InputFile, Dest)> = Vec::new();
83 let mut zip_owner: Option<ZipArchive<File>> = None;
84
85 loop {
86 let (src, dest) = {
87 let mut it = it.lock().ok();
88 match it.as_mut().and_then(|it| it.next()) {
89 Some((subdirs, src)) => {
90 is_found_any.call_once(|| {
91 info!("{}...", Msg::CleaningBooks);
92 });
93 let dest = Dest::new(cfg, subdirs, &src);
94 (src, dest)
95 }
96 None => break,
97 }
98 };
99 info!("{} '{}'...", Msg::Cleaning, src.path.display());
100
101 if !cfg.force && dest.path.exists() {
102 warn!(
103 "{} '{}'. {}",
104 Msg::FileIsAlreadyExists,
105 dest.path.display(),
106 Msg::Skipping,
107 );
108 continue;
109 }
110
111 match try_reader_writer(&mut zip_owner, &src, &dest)
112 .and_then(|(mut r, mut w)| remove_xml_tags(&mut r, &mut w, &cfg.tags))
113 {
114 Err(e) if cfg.exit_on_err => return Err(e.to_string()),
115 Err(e) => {
116 error!("{}. {}", e, Msg::Skipping);
117 continue;
118 }
119 Ok(()) => {
120 info!(
121 "{} '{}'",
122 Msg::SuccessCleanedAndSavedTo,
123 dest.path.display()
124 )
125 }
126 }
127
128 if cfg.force {
129 src_dests.push((src, dest));
130 }
131 }
132 Ok(src_dests)
133}
134
135fn force_overwrites(src_dests: Vec<(InputFile, Dest)>) {
136 info!("\n{}...", Msg::OverwritingBooks);
137
138 for (src, dest) in &src_dests {
139 info!("{} '{}'...", Msg::Overwriting, src.path.display());
140 match dest.force_overwrite(&src) {
141 Ok(()) => info!("{} '{}'", Msg::SuccessOverwritedFrom, dest.path.display()),
142 Err(e) => error!("{}: {}", Msg::Overwriting, e),
143 }
144 }
145
146 let mut dirs: Vec<PathBuf> = Vec::new();
147 for (_, dest) in src_dests {
148 for d in dest.created_dirs.into_iter().flat_map(|ds| ds.into_iter()) {
149 dirs.push(d);
150 }
151 }
152 dirs.par_sort_by(|a, b| b.as_os_str().len().cmp(&a.as_os_str().len()));
153
154 for d in dirs {
155 if let Err(e) = fs::remove_dir(&d) {
156 error!("{} '{}': {}", Msg::RemovingTempDirectory, d.display(), e);
157 }
158 }
159}
160
161fn try_reader_writer<'a>(
162 zip_owner: &'a mut Option<ZipArchive<File>>,
163 src: &'a InputFile,
164 dest: &'a Dest,
165) -> Result<(Reader<Box<dyn BufRead + 'a>>, Writer<Box<dyn Write + 'a>>)> {
166 let src_file = File::open(&src.path)?;
167 let dest_file = File::create(&dest.path)?;
168
169 let reader = match src.ty {
170 InputFileType::Fb2 => {
171 Reader::from_reader(Box::new(BufReader::new(src_file)) as Box<dyn BufRead>)
172 }
173 InputFileType::Fb2Zip => {
174 *zip_owner = Some(ZipArchive::new(src_file)?);
175 let zip = zip_owner.as_mut().unwrap();
176 let fb2_index = (0..zip.len())
177 .find(|&i| {
178 zip.by_index(i)
179 .ok()
180 .map(|f| {
181 let bytes = f.name().as_bytes();
182 let len = bytes.len();
183 len > 3 && bytes[len - 4..len].eq_ignore_ascii_case(b".fb2")
184 })
185 .unwrap_or(false)
186 })
187 .ok_or_else(|| {
188 format!("{} '{}'", Msg::NotFoundAnyFb2InArchive, src.path.display())
189 })?;
190
191 let fb2_file = zip.by_index(fb2_index)?;
192 Reader::from_reader(Box::new(BufReader::new(fb2_file)) as Box<dyn BufRead>)
193 }
194 };
195
196 let writer = match dest.ty {
197 InputFileType::Fb2 => Writer::new(Box::new(BufWriter::new(dest_file)) as Box<dyn Write>),
198 InputFileType::Fb2Zip => {
199 let mut zip_writer = ZipWriter::new(BufWriter::new(dest_file));
200 zip_writer.start_file(dest.zip_start_file(), SimpleFileOptions::default())?;
201 Writer::new(Box::new(zip_writer) as Box<dyn Write>)
202 }
203 };
204
205 Ok((reader, writer))
206}
207
208struct Dest {
209 created_dirs: Option<Vec<PathBuf>>,
210 ty: InputFileType,
211 stem: PathBuf,
212 path: PathBuf,
213}
214
215impl Dest {
216 fn new(cfg: &Config, subdirs: Option<Vec<PathBuf>>, src: &InputFile) -> Dest {
217 let ty = if cfg.zip {
218 InputFileType::Fb2Zip
219 } else if cfg.unzip {
220 InputFileType::Fb2
221 } else {
222 src.ty
223 };
224
225 let stem = match src.path.file_stem() {
226 Some(os) => {
227 let mut p = PathBuf::from(os);
228 if src.ty.is_fb2_zip() {
229 p.set_extension("");
230 }
231 p
232 }
233 _ => PathBuf::from("cleaned"),
234 };
235
236 let mut path = cfg.output.dir.clone().into_path_buf();
237 let mut created_dirs: Option<Vec<PathBuf>> = None;
238
239 if let Some(xs) = &subdirs {
240 for x in xs {
241 path.push(x);
242 if let Ok(()) = fs::create_dir(&path) {
243 created_dirs
244 .get_or_insert_with(|| Vec::with_capacity(xs.len()))
245 .push(path.clone())
246 }
247 }
248 }
249 path.push(&stem);
250 path.add_extension(ty.as_extension());
251
252 Dest {
253 created_dirs,
254 ty,
255 stem,
256 path,
257 }
258 }
259
260 fn force_overwrite(&self, src: &InputFile) -> Result<()> {
261 let force_path = self.force_path(src);
262
263 if let Err(_) = fs::rename(&self.path, &force_path) {
264 fs::copy(&self.path, &force_path)?;
265 if let Err(e) = fs::remove_file(&self.path) {
266 error!("{} '{}': {}", Msg::RemovingTempFile, self.path.display(), e);
267 }
268 }
269
270 if force_path != &*src.path {
271 if let Err(e) = fs::remove_file(&src.path) {
272 error!("{} '{}': {}", Msg::RemovingInputFile, src.path.display(), e);
273 }
274 }
275 Ok(())
276 }
277
278 fn force_path<'a>(&self, src: &'a InputFile) -> Cow<'a, Path> {
279 if self.ty == src.ty {
280 return Cow::Borrowed(&*src.path);
281 }
282
283 let mut p = match src.path.parent() {
284 Some(p) => PathBuf::from(p),
285 None => PathBuf::from("."),
286 };
287 p.push(&self.stem);
288 p.add_extension(self.ty.as_extension());
289 Cow::Owned(p)
290 }
291
292 fn zip_start_file(&self) -> String {
293 let mut s = self.stem.to_string_lossy().into_owned();
294 s.push_str(".fb2");
295 s
296 }
297}