Added match feature to tmagic

Author: aahancoc

Cargo.toml

@@ -24,7 +24,7 @@ num_cpus = { version = "^1", optional = true }
 [features]
 cli = ["clap", "tabwriter", "scoped_threadpool", "walkdir", "num_cpus"]
 staticmime = []     # Use &'static str for output insead of String. (disabled for now)
-default = ["cli"]
+default = ["cli", "staticmime"]
 
 [lib]
 crate-type = ["lib"]

src/.main.rs.kate-swp

@@ -114,10 +114,7 @@ pub mod check {
             "all/allfiles" | "application/octet-stream" => return meta.is_file(),
             "inode/directory" => return meta.is_dir(),
             "text/plain" => return is_text_plain_from_filepath(filepath),
-            _ => {
-                println!("{}", mimetype);
-                panic!("This mime is not supported by the mod. (See can_check)")
-            }
+            _ => return false
         }
 
     }

src/basetype.rs

@@ -112,12 +112,6 @@ pub mod check {
 
     /// Test against all rules
     pub fn from_u8(file: &[u8], mimetype: &str) -> bool {
-		
-		/*// Get mimetype in case user provides alias
-		let mimetype = match super::ALIASES.get(mimetype) {
-			None => mimetype,
-			Some(x) => x
-		};*/
 
         // Get magic ruleset
         let graph = match super::ALLRULES.get(mimetype) {
@@ -166,6 +160,13 @@ pub mod check {
             Err(_) => return false
         };
         let mut b = Vec::<u8>::with_capacity(scanlen);
+        
+        // Fill up vector with something
+        for i in 0..scanlen {
+            let _ = i;
+            b.push(0);
+        }
+        
         match f.read_exact(&mut b) {
             Ok(_) => {},
             Err(_) => return false

src/fdo_magic/builtin.rs

@@ -187,6 +187,13 @@ pub mod check {
             Err(_) => return false
         };
         let mut b = Vec::<u8>::with_capacity(scanlen);
+        
+        // Fill up vector with something
+        for i in 0..scanlen {
+            let _ = i;
+            b.push(0);
+        }
+//         
         match f.read_exact(&mut b) {
             Ok(_) => {},
             Err(_) => return false

src/fdo_magic/sys.rs

@@ -318,7 +318,7 @@ fn get_alias(mimetype: &String) -> &String {
 fn match_u8_noalias(mimetype: &str, bytes: &[u8]) -> bool
 {
     match CHECKER_SUPPORT.get(mimetype) {
-        None => {panic!("{} not supported", mimetype);},
+        None => {false},
         Some(y) => (CHECKERS[*y].from_u8)(bytes, mimetype)
     }
 }
@@ -373,7 +373,7 @@ pub fn from_u8(bytes: &[u8]) -> Option<MIME>
 fn match_filepath_noalias(mimetype: &str, filepath: &str) -> bool
 {
     match CHECKER_SUPPORT.get(mimetype) {
-        None => {panic!("{} not supported", mimetype);},
+        None => {false},
         Some(y) => (CHECKERS[*y].from_filepath)(filepath, mimetype)
     }
 }

src/lib.rs

@@ -26,7 +26,7 @@ fn main() {
     use clap::{Arg, App};
 
     let args = App::new("TreeMagic")
-        .version("0.1")
+        .version("0.2.0")
         .about("Determines the MIME type of a file by traversing a filetype tree.")
         .arg(Arg::with_name("file")
             .required(true)
@@ -35,31 +35,40 @@ fn main() {
         )
         .arg(Arg::with_name("recursive")
             .short("r")
+            .long("recursive")
             .help("Search directories recursively")
         )
         .arg(Arg::with_name("match")
             .short("m")
             .long("match")
             .use_delimiter(true)
             .takes_value(true)
+            .require_equals(true)
             .help("Print only files that match given MIME")
         )
         .arg(Arg::with_name("ugly")
             .long("ugly")
             .help("Print results as they come in, at expense of tab alignment")
         )
         .get_matches();
+        
     let mut files: Vec<String> = args.values_of("file")
         .unwrap()
         .map(|x| x.to_string())
         .collect();
+    let is_ugly = args.is_present("ugly");
+    let is_recursive = args.is_present("recursive");
+    let check_against: Vec<String> = args.values_of("match")
+        .unwrap()
+        .map(|x| x.to_string())
+        .collect();
+    println!("{:?}", check_against);
 
     let mut tw = TabWriter::new(vec![]);
     let (tx, rx) = mpsc::channel();
 
     // Get recursive results if needed
-    if args.is_present("recursive") {
-        //println!("Recursive!");
+    if is_recursive {
         for dir in files.clone() {
             let entries = WalkDir::new(dir).into_iter().filter_map(|e| e.ok());
             for entry in entries {
@@ -70,29 +79,61 @@ fn main() {
                 );
             }
         }
-        //println!("{:#?}", files);
     }
+    let files = files;
 
     let mut pool = Pool::new(num_cpus::get() as u32);
-    let is_ugly = args.is_present("ugly");
-    
-    pool.scoped(|scope| {
-        for x in files {
-            let tx = tx.clone();
-            scope.execute(move || {
-                let result = tree_magic::from_filepath(x.as_str());
-                let result = result.unwrap_or(convmime!("inode/none"));
-                let result = format!("{}:\t{:?}", x, result);
-                if is_ugly {
-                    println!("{}", result);
-                } else {
-                    tx.send(result + "\n").unwrap_or_default();
-                }
-            });
-        }
-    });
+    // Acquire results for non-match
+    if check_against.is_empty(){
+        pool.scoped(|scope| {
+            for file in files {
+                let tx = tx.clone();
+                scope.execute(move || {
+                    let result = tree_magic::from_filepath(file.as_str());
+                    let result = result.unwrap_or(convmime!("inode/none"));
+                    let result = format!("{}:\t{:?}", file, result);
+                    if is_ugly {
+                        println!("{}", result);
+                    } else {
+                        tx.send(result + "\n").unwrap_or_default();
+                    }
+                });
+            }
+        });
+    // Acquire results for check against list of MIMES
+    } else {
+        pool.scoped(|scope| {
+            for file in files {
+                let tx = tx.clone();
+                let check_against = check_against.clone();
+                
+                scope.execute(move || {
+                    let mut result: Option<String> = None;
+                
+                    for mime in check_against {
+                        let out = tree_magic::match_filepath(mime.as_str(), file.as_str());
+                        if out {
+                            result = Some(mime);
+                            break;
+                        }
+                    }
+                    
+                    if result.is_none() { return; }
+                    
+                    let result = result.unwrap();
+                    let result = format!("{}:\t{:?}", file, result);
+                    if is_ugly {
+                        println!("{}", result);
+                    } else {
+                        tx.send(result + "\n").unwrap_or_default();
+                    }
+                });
+            }
+        });
+    }
     drop(tx);
 
+    // Pretty-print results
     if !is_ugly {
         let mut list: Vec<_> = rx.iter().collect();
         list.sort();