antivirus.rs

#![feature(test)]

extern crate md5;
extern crate walkdir;
extern crate test;

use walkdir::WalkDir;

use std::path::Path;
use std::io::{self, Read};
use std::fs::File;

static MALWARE: &str = "C:\\Users\\ngc0202\\Downloads\\Android_Antivirus\\Malware\\";

fn md5_file(path: &Path) -> io::Result<[u8; 16]> {
	let mut v = Vec::new();
	File::open(path)?.read_to_end(&mut v)?;
	Ok(*md5::compute(&v))
}

fn md5_apks2(root_path: &str) -> Vec<(String, [u8; 16])> {
	let mut hashes = Vec::new();
	for entry in WalkDir::new(root_path).into_iter() {
		if let Ok(entry) = entry {
			let name = entry.file_name().to_owned().into_string().unwrap();
			if name.ends_with("apk") {
				let hash = md5_file(entry.path()).unwrap();
				hashes.push((name, hash));
			}
		}
	}

	hashes
}

fn md5_apks(root_path: &str) -> Vec<(String, [u8; 16])> {
	WalkDir::new(root_path).into_iter()
		.filter_map(|r| r.ok().map(|e| e.path().to_path_buf()))
		.filter(|p| p.extension().map(|ex| ex == "apk").unwrap_or(false))
		.map(|path| (path.file_name().unwrap().to_owned().into_string().unwrap(), md5_file(&path).unwrap()))
		.collect()
}

fn main() {
	let malware_hashes = md5_apks(MALWARE);
	let check_files = WalkDir::new(std::env::args().nth(1).unwrap_or(MALWARE.to_owned()).as_str());
	for entry in check_files.into_iter().filter_map(|e| e.ok()).filter(|e| e.file_type().is_file()) {
		let name = entry.file_name();
		println!("Checking: {:?}", name);
		let hash = md5_file(entry.path()).unwrap();
		for &(ref malname, ref malhash) in &malware_hashes {
			if &hash == malhash {
				println!("{:?} is malware: {}", name, malname);
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use test::Bencher;

	#[bench]
	fn bench_apks1(b: &mut Bencher) {
		b.iter(|| md5_apks(MALWARE))
	}

	#[bench]
	fn bench_apks2(b: &mut Bencher) {
		b.iter(|| md5_apks2(MALWARE))
	}
}