fif/src/main.rs

// fif - a command-line tool for detecting and optionally correcting files with incorrect extensions.
// Copyright (C) 2021 Lynnesbian
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#![forbid(unsafe_code)]
#![warn(trivial_casts, unused_lifetimes, unused_qualifications)]

use std::io::{stdin, stdout, BufWriter, Write};
use std::process::exit;

use cfg_if::cfg_if;
use clap::Clap;
use fif::files::{scan_directory, scan_from_walkdir};
use fif::formats::Format;
use fif::parameters::{OutputFormat, Prompt};
use fif::utils::{os_name, CLAP_LONG_VERSION};
use fif::{formats, parameters};
use itertools::Itertools;
use log::{debug, error, info, trace, warn, Level};

#[cfg(test)]
mod tests;

#[doc(hidden)]
#[allow(clippy::cognitive_complexity, clippy::too_many_lines)]
fn main() {
	let args: parameters::Parameters = parameters::Parameters::parse();

	let mut builder = env_logger::Builder::new();
	builder
		.filter_level(args.get_verbosity()) // set default log level
		.parse_default_env() // set log level from RUST_LOG
		.parse_env("FIF_LOG") // set log level from FIF_LOG
		.format(|buf, r| {
			let mut style = buf.default_level_style(r.level());
			// use bold for warnings and errors
			style.set_bold(r.level() <= Level::Warn);
			// only use the first character of the log level name
			let abbreviation = style.value(r.level().to_string().chars().next().unwrap());
			// e.g. [D] Debug message
			writeln!(buf, "[{}] {}", abbreviation, r.args())
		})
		.init();

	trace!(
		"fif {}, running on {} {}",
		CLAP_LONG_VERSION.as_str(),
		std::env::consts::ARCH,
		os_name()
	);

	debug!("Iterating directory: {:?}", args.dir);

	let extensions = args.extensions();
	let excludes = args.excluded_extensions();

	if let Some(extensions) = &extensions {
		debug!("Checking files with extensions: {:?}", extensions);
	} else if let Some(excludes) = &excludes {
		debug!("Skipping files with extensions: {:?}", excludes);
	} else {
		debug!("Checking files regardless of extensions");
	}

	let entries = match scan_directory(&args.dir, extensions.as_ref(), excludes.as_ref(), &args.get_scan_opts()) {
		// no need to log anything for fatal errors - fif will already have printed something obvious like
		// "[ERROR] /fake/path: No such file or directory (os error 2)". we can assume that if this has happened, the dir
		// given as input doesn't exist or is otherwise unreadable.
		None => exit(exitcode::NOINPUT),
		Some(e) => e,
	};

	if entries.is_empty() {
		warn!("No files matching requested options found.");
		exit(exitcode::OK);
	}

	trace!("Found {} items to check", entries.len());

	cfg_if! {
		if #[cfg(feature = "multi-threaded")] {
			let use_threads = args.jobs != 1;

			if use_threads {
				// 0 is a special case - it should be understood to mean "all available host CPUs"
				let jobs = if args.jobs == 0 { num_cpus::get() } else { args.jobs };

				// set up the global thread pool with the requested number of threads
				rayon::ThreadPoolBuilder::new().num_threads(jobs).build_global().unwrap();
				trace!("Multithreading enabled, using {} threads", jobs);
			} else {
				trace!("Multithreading disabled at runtime");
			}

		} else { // `multi-threading` feature disabled
			let use_threads = false;
			trace!("Multithreading disabled at compile time");
		}
	}

	let (findings, errors) = scan_from_walkdir(&entries, args.canonical_paths, use_threads);
	trace!("Scanning complete");

	if findings.is_empty() && errors.is_empty() {
		info!("All files have valid extensions!");
		exit(exitcode::OK);
	}

	// remove files that already have the correct extension, then sort - first by whether or not they have a
	// recommended_extension() (with None before Some(ext)), then by filename
	let findings = findings
		.into_iter()
		.filter(|f| !f.valid)
		.sorted_unstable()
		.collect_vec();
	// sort errors (File errors before Mime errors), then log a warning for each error
	let errors = errors
		.into_iter()
		.sorted_unstable()
		.map(|e| {
			warn!("{}", &e);
			e
		})
		.collect_vec();

	if args.fix {
		fn ask(message: &str) -> bool {
			let mut buf = String::with_capacity(1);
			print!("{} [y/N] ", message);

			// flush stdout to ensure message is displayed
			stdout().flush().expect("Failed to flush stdout");

			if let Err(e) = stdin().read_line(&mut buf) {
				// something went wrong while reading input - just exit
				error!("{}", e);
				exit(exitcode::IOERR)
			}
			buf.starts_with('y') || buf.starts_with('Y')
		}

		let prompt = args.prompt.unwrap_or(Prompt::Error);

		for f in findings {
			if let Some(rename_to) = f.recommended_path() {
				let will_rename = {
					if !args.overwrite && rename_to.exists() {
						// handles: --prompt never, --prompt error, --prompt always
						// user didn't specify --overwrite, and the destination exists
						info!("Not renaming {:#?}: Target {:#?} exists", f.file, rename_to);
						false
					} else if prompt == Prompt::Never {
						// handles: --prompt never --overwrite
						// user specified --prompt never in conjunction with --overwrite, so always rename
						true
					} else if prompt == Prompt::Error || ask(&*format!("Rename {:#?} to {:#?}?", &f.file, &rename_to)) {
						// handles: --prompt error --overwrite, --prompt always --overwrite [y]
						// if the target exists, prompt before renaming; otherwise, just rename
						!rename_to.exists() || ask(&*format!("Destination {:#?} already exists, overwrite?", rename_to))
					} else {
						// handles: --prompt always --overwrite [n]
						// user was prompted and replied "no"
						false
					}
				};

				if !will_rename {
					continue;
				}

				loop {
					match std::fs::rename(&f.file, &rename_to) {
						Ok(_) => {
							info!("Renamed {:#?} -> {:#?}", f.file, rename_to);
							break;
						}
						Err(e) => {
							warn!("Couldn't rename {:#?} to {:#?}: {:#?}", f.file, rename_to, e);
							// if the user passed --prompt never, continue to the next file
							// otherwise, prompt user to retry move, retrying until the rename succeeds or they respond "N"
							if prompt == Prompt::Never || !ask(&*format!("Error while renaming file: {:#?}. Try again?", e)) {
								break;
							}
						}
					}
				}
			} else {
				// no recommended name :c
				info!("No known extension for file {:#?} of type {}", f.file, f.mime);
			}
		}
	} else {
		let mut buffered_stdout = BufWriter::new(stdout());

		let result = match args.output_format {
			// i want to simplify this to something like formats::write_all(args.output_format, ...)
			OutputFormat::Sh => formats::Shell.write_all(&mut buffered_stdout, &findings, &errors),
			OutputFormat::PowerShell => formats::PowerShell.write_all(&mut buffered_stdout, &findings, &errors),
			#[cfg(feature = "json")]
			OutputFormat::Json => formats::Json.write_all(&mut buffered_stdout, &findings, &errors),
			OutputFormat::Text => formats::Text.write_all(&mut buffered_stdout, &findings, &errors),
		};

		if result.is_err() {
			error!("Failed to write to stdout.");
			exit(exitcode::IOERR);
		}

		if buffered_stdout.flush().is_err() {
			error!("Failed to flush stdout.");
			exit(exitcode::IOERR);
		}
	}

	debug!("Done");
}