2021-03-25 18:46:07 +00:00
|
|
|
use crate::inspectors::{mime_extension_lookup, BUF_SIZE};
|
|
|
|
use crate::mime_db::{MimeDb, XdgDb};
|
2021-02-28 09:47:18 +00:00
|
|
|
use crate::{extension_from_path, init_db, scan_directory, scan_from_walkdir};
|
|
|
|
|
2021-03-02 15:12:29 +00:00
|
|
|
use crate::parameters::{Parameters, ScanOpts};
|
2021-02-28 09:47:18 +00:00
|
|
|
use cfg_if::cfg_if;
|
|
|
|
use mime_guess::mime::{APPLICATION_OCTET_STREAM, APPLICATION_PDF, IMAGE_JPEG, IMAGE_PNG};
|
|
|
|
use mime_guess::Mime;
|
|
|
|
use smartstring::alias::String;
|
2021-03-25 18:46:07 +00:00
|
|
|
use std::borrow::Borrow;
|
2021-02-28 09:47:18 +00:00
|
|
|
use std::collections::HashMap;
|
2021-03-25 18:46:07 +00:00
|
|
|
use std::ffi::OsStr;
|
2021-02-28 09:47:18 +00:00
|
|
|
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! {
|
2021-03-02 15:12:29 +00:00
|
|
|
if #[cfg(any(all(unix, feature = "infer-backend"), all(not(unix), not(feature = "xdg-mime-backend"))))] {
|
2021-02-28 09:47:18 +00:00
|
|
|
fn get_mime_db() -> InferDb {
|
|
|
|
InferDb::init()
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
fn get_mime_db() -> XdgDb {
|
|
|
|
XdgDb::init()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn application_zip() -> Mime {
|
|
|
|
use std::str::FromStr;
|
|
|
|
Mime::from_str("application/zip").unwrap()
|
|
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn get_ext() {
|
2021-03-25 18:46:07 +00:00
|
|
|
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("")));
|
2021-02-28 09:47:18 +00:00
|
|
|
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);
|
|
|
|
}
|
|
|
|
|
2021-03-11 17:44:31 +00:00
|
|
|
let scan_opts = ScanOpts {
|
2021-03-02 15:12:29 +00:00
|
|
|
hidden: true,
|
2021-03-11 17:44:31 +00:00
|
|
|
extensionless: false,
|
2021-03-02 15:12:29 +00:00
|
|
|
};
|
|
|
|
|
2021-02-28 09:47:18 +00:00
|
|
|
let entries = scan_directory(
|
|
|
|
&dir.path().to_path_buf(),
|
2021-03-25 18:46:07 +00:00
|
|
|
&["jpg", "jpeg", "png", "pdf", "zip"],
|
2021-03-02 15:12:29 +00:00
|
|
|
&scan_opts,
|
2021-02-28 09:47:18 +00:00
|
|
|
)
|
|
|
|
.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");
|
2021-02-28 12:16:54 +00:00
|
|
|
|
2021-02-28 09:47:18 +00:00
|
|
|
if !result.valid {
|
|
|
|
// this should be "wrong.jpg", which is a misnamed png file
|
2021-02-28 12:16:54 +00:00
|
|
|
// 1. ensure extension is "png"
|
2021-03-25 18:46:07 +00:00
|
|
|
assert_eq!(extension_from_path(&*result.file).unwrap(), OsStr::new("jpg"));
|
2021-02-28 12:16:54 +00:00
|
|
|
// 2. ensure mime type detected is IMAGE_PNG
|
2021-02-28 09:47:18 +00:00
|
|
|
assert_eq!(result.mime, IMAGE_PNG);
|
2021-02-28 12:16:54 +00:00
|
|
|
// 3. ensure recommended extension is in the list of known extensions for PNG files
|
2021-02-28 14:06:05 +00:00
|
|
|
assert!(mime_extension_lookup(IMAGE_PNG)
|
|
|
|
.unwrap()
|
2021-03-25 18:46:07 +00:00
|
|
|
.contains(&result.recommended_extension().unwrap()));
|
2021-02-28 09:47:18 +00:00
|
|
|
continue;
|
|
|
|
}
|
|
|
|
|
2021-02-28 12:16:54 +00:00
|
|
|
// check if the recommended extension for this file is in the list of known extensions for its mimetype
|
2021-02-28 14:06:05 +00:00
|
|
|
assert!(mime_extension_lookup(result.mime.clone())
|
|
|
|
.unwrap()
|
2021-03-25 18:46:07 +00:00
|
|
|
.contains(&result.recommended_extension().unwrap()));
|
2021-02-28 12:16:54 +00:00
|
|
|
|
|
|
|
// make sure the guessed mimetype is correct based on the extension of the scanned file
|
2021-03-25 18:46:07 +00:00
|
|
|
let ext = extension_from_path(result.file);
|
|
|
|
assert!(ext.is_some());
|
2021-02-28 09:47:18 +00:00
|
|
|
assert_eq!(
|
|
|
|
result.mime,
|
2021-03-25 18:46:07 +00:00
|
|
|
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
|
2021-02-28 09:47:18 +00:00
|
|
|
}
|
|
|
|
);
|
|
|
|
}
|
|
|
|
}
|
2021-02-28 12:16:54 +00:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn argument_parsing() {
|
|
|
|
use clap::Clap;
|
|
|
|
|
|
|
|
// check if "jpg" is in the list of extensions to be considered when passing "-E images"
|
|
|
|
let args: Parameters = Parameters::parse_from(vec!["fif", "-E", "images"]);
|
|
|
|
assert!(args.extensions().contains(&"jpg"));
|
|
|
|
|
|
|
|
// make sure "scan_hidden" is false
|
|
|
|
assert!(!args.scan_hidden);
|
|
|
|
|
|
|
|
// exts should be none
|
|
|
|
assert!(args.exts.is_none());
|
2021-02-28 14:06:05 +00:00
|
|
|
}
|
2021-03-25 18:46:07 +00:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn rejects_bad_args() {
|
|
|
|
use clap::Clap;
|
|
|
|
assert!(Parameters::try_parse_from(vec!["fif", "-abcdefg", "-E", "-e"]).is_err());
|
|
|
|
}
|
|
|
|
|
|
|
|
#[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>())
|
|
|
|
}
|