fif/src/tests/mod.rs

use crate::inspectors::{mime_extension_lookup, BUF_SIZE};
use crate::{extension_from_path, init_db, scan_directory, scan_from_walkdir};

use crate::mime_db::MimeDb;
use crate::parameters::{Parameters, ScanOpts};
use crate::string_type::String;
use cfg_if::cfg_if;
use mime_guess::mime::{APPLICATION_OCTET_STREAM, APPLICATION_PDF, IMAGE_JPEG, IMAGE_PNG};
use mime_guess::Mime;

use std::borrow::Borrow;
use std::collections::HashMap;
use std::ffi::OsStr;
use std::path::Path;

const JPEG_BYTES: &[u8] = b"\xFF\xD8\xFF";
const PNG_BYTES: &[u8] = b"\x89\x50\x4E\x47\x0D\x0A\x1A\x0A";
const PDF_BYTES: &[u8] = b"%PDF-";
const ZIP_BYTES: &[u8] = b"PK\x03\x04";

cfg_if! {
	if #[cfg(any(all(unix, feature = "infer-backend"), all(not(unix), not(feature = "xdg-mime-backend"))))] {
		fn get_mime_db() -> crate::mime_db::InferDb {
			crate::mime_db::InferDb::init()
		}
	} else {
		fn get_mime_db() -> crate::mime_db::XdgDb {
			crate::mime_db::XdgDb::init()
		}
	}
}

fn application_zip() -> Mime {
	use std::str::FromStr;
	Mime::from_str("application/zip").unwrap()
}

#[test]
fn get_ext() {
	let mut ext_checks: HashMap<_, Option<&OsStr>> = HashMap::new();
	ext_checks.insert(Path::new("test.txt"), Some(OsStr::new("txt")));
	ext_checks.insert(Path::new("test.zip"), Some(OsStr::new("zip")));
	ext_checks.insert(Path::new("test.tar.gz"), Some(OsStr::new("gz")));
	ext_checks.insert(Path::new("test."), Some(OsStr::new("")));
	ext_checks.insert(Path::new("test"), None);
	ext_checks.insert(Path::new(".hidden"), None);

	for (path, ext) in ext_checks {
		assert_eq!(extension_from_path(path), ext)
	}
}

#[test]
fn detect_type() {
	let db = get_mime_db();
	assert_eq!(db.get_type(JPEG_BYTES), Some(IMAGE_JPEG));
	assert_eq!(db.get_type(PNG_BYTES), Some(IMAGE_PNG));
	assert_eq!(db.get_type(PDF_BYTES), Some(APPLICATION_PDF));
	assert_eq!(db.get_type(ZIP_BYTES), Some(application_zip()));
}

#[test]
fn recommend_ext() {
	assert!(mime_extension_lookup(IMAGE_JPEG)
		.unwrap()
		.contains(&String::from("jpg")));
	assert!(mime_extension_lookup(IMAGE_PNG).unwrap().contains(&String::from("png")));
	assert!(mime_extension_lookup(APPLICATION_PDF)
		.unwrap()
		.contains(&String::from("pdf")));
	assert!(mime_extension_lookup(application_zip())
		.unwrap()
		.contains(&String::from("zip")));
}

#[test]
fn simple_directory() {
	use std::fs::File;
	use std::io::Write;
	use tempfile::tempdir;

	let mut files = HashMap::new();
	files.insert("test.jpg", JPEG_BYTES);
	files.insert("test.jpeg", JPEG_BYTES);
	files.insert("test.png", PNG_BYTES);
	files.insert("test.pdf", PDF_BYTES);
	files.insert("test.zip", ZIP_BYTES);
	files.insert("wrong.jpg", PNG_BYTES);

	let dir = tempdir().expect("Failed to create temporary directory.");

	for (name, bytes) in &files {
		let mut file = File::create(dir.path().join(name)).expect(&*format!("Failed to create file: {}", name));

		file
			.write_all(bytes)
			.expect(&*format!("Failed to write to file: {}", name));
		drop(file);
	}

	let scan_opts = ScanOpts {
		hidden: true,
		extensionless: false,
		follow_symlinks: false,
	};

	let entries = scan_directory(
		&dir.path().to_path_buf(),
		Some(&vec!["jpg", "jpeg", "png", "pdf", "zip"]),
		&scan_opts,
	)
	.expect("Directory scan failed.");

	assert_eq!(entries.len(), files.len());

	// initialise global mime DB
	init_db();

	let results = scan_from_walkdir(&entries);
	for result in results {
		let result = result.expect("Error while scanning file");

		if !result.valid {
			// this should be "wrong.jpg", which is a misnamed png file
			// 1. ensure extension is "png"
			assert_eq!(extension_from_path(&*result.file).unwrap(), OsStr::new("jpg"));
			// 2. ensure mime type detected is IMAGE_PNG
			assert_eq!(result.mime, IMAGE_PNG);
			// 3. ensure recommended extension is in the list of known extensions for PNG files
			assert!(mime_extension_lookup(IMAGE_PNG)
				.unwrap()
				.contains(&result.recommended_extension().unwrap()));
			continue;
		}

		// check if the recommended extension for this file is in the list of known extensions for its mimetype
		assert!(mime_extension_lookup(result.mime.clone())
			.unwrap()
			.contains(&result.recommended_extension().unwrap()));

		// make sure the guessed mimetype is correct based on the extension of the scanned file
		let ext = extension_from_path(result.file);
		assert!(ext.is_some());
		assert_eq!(
			result.mime,
			match ext.unwrap().to_string_lossy().borrow() {
				"jpg" | "jpeg" => IMAGE_JPEG,
				"png" => IMAGE_PNG,
				"pdf" => APPLICATION_PDF,
				"zip" => application_zip(),
				_ => APPLICATION_OCTET_STREAM, // general "fallback" type
			}
		);
	}
}

#[test]
fn argument_parsing() {
	use clap::Clap;

	// pass `-f`, which enables following symlinks, and `-E images`, which scans files with image extensions
	let args: Parameters = Parameters::parse_from(vec!["fif", "-f", "-E", "images"]);

	// check if "jpg" is in the list of extensions to be scanned
	assert!(args
		.extensions()
		.expect("args.extensions() should contain the `images` set!")
		.contains(&"jpg"));

	// make sure "scan_hidden" is false
	assert!(!args.scan_hidden);

	// exts should be none
	assert!(args.exts.is_none());

	// get the ScanOpts, and make sure they match expectations
	assert_eq!(
		args.get_scan_opts(),
		ScanOpts {
			hidden: false,
			extensionless: false,
			follow_symlinks: true
		}
	)
}

#[test]
fn rejects_bad_args() {
	use clap::Clap;
	let tests = [
		// Non-existent flags:
		vec!["fif", "-abcdefghijklmnopqrstuvwxyz"],
		// `-E` without specifying a set:
		vec!["fif", "-E"],
		// `-E` with an invalid set:
		vec!["fif", "-E", "pebis"],
		// `-E` and `-e`:
		vec!["fif", "-E", "media", "-e", "jpg"],
		// `-e` with nothing but commas:
		vec!["fif", "-e", ",,,,,"],
	];

	for test in &tests {
		assert!(Parameters::try_parse_from(test).is_err(), "Failed to reject {:?}", test);
	}
}

#[test]
fn identify_random_bytes() {
	let db = get_mime_db();
	let rng = fastrand::Rng::new();
	let mut bytes: Vec<u8>;
	let mut results: HashMap<Mime, i32> = HashMap::new();

	for _ in 1..500 {
		bytes = std::iter::repeat_with(|| rng.u8(..)).take(BUF_SIZE * 2).collect();
		if let Some(detected_type) = db.get_type(&*bytes) {
			*results.entry(detected_type).or_insert(0) += 1;
		}
	}

	for (mime, count) in &results {
		println!("{}:\t{} counts", mime, count);
	}
	println!("No type found:\t{} counts", 500 - results.values().sum::<i32>())
}