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/>.

use std::io::{stdout, BufWriter};
use std::path::{Path, PathBuf};

use clap::Clap;
use log::{debug, info, trace, warn};
use once_cell::sync::OnceCell;
#[cfg(feature = "multi-threaded")]
use rayon::prelude::*;
use smartstring::alias::String;
use walkdir::{DirEntry, WalkDir};

use crate::findings::Findings;
use crate::formats::{Format, Script};
use crate::mimedb::MimeDb;
use crate::parameters::OutputFormat;
use crate::scanerror::ScanError;

mod findings;
mod formats;
mod inspectors;
mod mimedb;
mod parameters;
mod scanerror;
mod extensionset;

#[cfg(feature = "infer-backend")]
static MIMEDB: OnceCell<mimedb::InferDb> = OnceCell::new();

#[cfg(feature = "xdg-mime-backend")]
static MIMEDB: OnceCell<mimedb::XdgDb> = OnceCell::new();

// TODO: test if this actually works on a windows machine
#[cfg(windows)]
fn is_hidden(entry: &DirEntry) -> bool {
	use std::os::windows::prelude::*;
	std::fs::metadata(entry) // try to get metadata for file
		.map_or(
			false,                             // if getting metadata/attributes fails, assume it's not hidden
			|f| f.file_attributes() & 0x2 > 0, // flag for hidden - https://docs.microsoft.com/windows/win32/fileio/file-attribute-constants
		)
}

#[cfg(not(windows))]
fn is_hidden(entry: &DirEntry) -> bool {
	entry
		.file_name()
		.to_str()
		.map_or(false, |f| f.starts_with('.') && f != ".")
}

fn wanted_file(args: &parameters::Parameters, exts: &[&str], entry: &DirEntry) -> bool {
	if !args.scan_hidden && is_hidden(entry) {
		// skip hidden files and directories. this check is performed first because it's very lightweight.
		return false;
	}

	if entry.file_type().is_dir() {
		// always allow directories - there's no point doing file extension matching on something that isn't a file.
		return true;
	}

	let ext = extension_from_path(entry.path());

	if ext.is_none() {
		return false;
	} // don't scan files without extensions. TODO - this should be configurable

	exts.contains(&ext.unwrap().to_lowercase().as_str())
}

fn extension_from_path(path: &Path) -> Option<String> {
	path.extension(). // Get the path's extension
		map(|e| String::from(e.to_string_lossy())) // Convert from OsStr to String
}

fn scan_file(entry: &DirEntry) -> Result<Findings, (ScanError, PathBuf)> {
	// try to determine mimetype for this entry
	let result = inspectors::mime_type(MIMEDB.get().unwrap(), entry.path());

	if result.is_err() {
		// an error occurred while trying to read the file
		// error!("{}: {}", entry.path().to_string_lossy(), error);
		return Err((ScanError::File, entry.path().to_path_buf()));
	}

	let result = result.unwrap();
	if result.is_none() {
		// the file was read successfully, but we were unable to determine its mimetype
		// warn!("Couldn't determine mimetype for {}", entry.path().to_string_lossy());
		return Err((ScanError::Mime, entry.path().to_path_buf()));
	}

	let result = result.unwrap();

	// set of known extensions for the given mimetype
	let known_exts = inspectors::mime_extension_lookup(result.clone());
	// file extension for this particular file
	let entry_ext = extension_from_path(entry.path());

	let valid = match known_exts {
		// there is a known set of extensions for this mimetype, and the file has an extension
		Some(e) if entry_ext.is_some() => e.contains(&entry_ext.unwrap().to_lowercase().into()),
		// there is a known set of extensions for this mimetype, but the file has no extension
		Some(_) => false,
		// there is no known set of extensions for this mimetype -- assume it's correct
		None => true,
	};

	Ok(Findings {
		file: entry.path().to_path_buf(),
		valid, // make this a function
		mime: result,
	})
}

fn scan_from_walkdir(entries: Vec<DirEntry>) -> Vec<Result<Findings, (ScanError, PathBuf)>> {
	#[cfg(feature = "multi-threaded")]
	{
		// rather than using a standard par_iter, split the entries into chunks of 32 first.
		// this allows each spawned thread to handle 32 files before before closing, rather than creating a new thread for
		// each file. this leads to a pretty substantial speedup that i'm pretty substantially happy about 0u0
		entries
			.par_chunks(32) // split into chunks of 32
			.flat_map(|chunk| {
				chunk // return Vec<...> instead of Chunk<Vec<...>>
					.iter() // iter over the chunk, which is a slice of DirEntry structs
					.map(|entry| scan_file(entry))
					.collect::<Vec<_>>()
			})
			.collect()
	}

	#[cfg(not(feature = "multi-threaded"))]
	{
		entries.iter().map(|entry: &DirEntry| scan_file(entry)).collect()
	}
}

fn main() {
	let args = parameters::Parameters::parse();
	let mut builder = env_logger::Builder::from_default_env();
	builder
		// .format(|buf, r| writeln!(buf, "{} - {}", r.level(), r.args()))
		.format_module_path(false) // don't include module in logs, as it's not necessary
		.format_timestamp(None) // don't include timestamps (unnecessary, and the feature flag is disabled anyway)
		// .target(env_logger::Target::Stdout) // log to stdout rather than stderr
		.init();

	#[cfg(feature = "infer-backend")]
	MIMEDB
		.set(mimedb::InferDb::init())
		.or(Err("Failed to initialise MIMEDB"))
		.unwrap();

	#[cfg(feature = "xdg-mime-backend")]
	MIMEDB
		.set(mimedb::XdgDb::init())
		.or(Err("Failed to initialise MIMEDB"))
		.unwrap();

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

	let extensions: Vec<&str> = if let Some(exts) = &args.exts {
		exts
			.iter()
			.map(|s| s.as_str())
			.collect()
	} else if let Some(exts) = &args.ext_set {
		exts.extensions().to_vec()
	} else {
		unreachable!()
	};

	debug!("Checking files with extensions: {:?}", extensions);

	let stepper = WalkDir::new(&args.dirs).into_iter();
	let entries: Vec<DirEntry> = stepper
		.filter_entry(|e| wanted_file(&args, &extensions, e)) // filter out unwanted files
		.filter_map(|e| e.ok()) // ignore anything that fails, e.g. files we don't have read access on
		.filter(|e| !e.file_type().is_dir()) // remove directories from the final list
		.collect();

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

	let results = scan_from_walkdir(entries);

	for result in &results {
		match result {
			Ok(r) => {
				if !r.valid {
					info!(
						"{:?} should have file extension {}",
						r.file,
						r.recommended_extension().unwrap()
					)
				} else {
					trace!("{:?} is totally fine", r.file)
				}
			}
			Err(f) => warn!("{:#?}: Error 0uo - {}", f.1, f.0),
		}
	}

	match args.output_format {
		OutputFormat::Script => {
			let s = Script::new();
			s.write_all(&results, &mut BufWriter::new(stdout().lock()))
				.expect("failed to output");
		}
		OutputFormat::Text => todo!(),
	}

	debug!("Done");
}