use chrono::Utc;
use failure::Fail;
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::TryInto;
use crate::contenthash::ContentHash;
use crate::util::unwrap_try_into;
use crate::{GenericResult, SampleHash};
pub struct VirusTotalClient {
apikey: String,
}
impl VirusTotalClient {
pub fn new(apikey: impl AsRef<str>) -> Self {
VirusTotalClient {
apikey: apikey.as_ref().to_owned(),
}
}
fn file_report_url(&self, resource: impl AsRef<str>, allinfo: bool) -> String {
format!(
"https://www.virustotal.com/vtapi/v2/file/report?apikey={}&allinfo={}&resource={}",
self.apikey,
allinfo,
resource.as_ref()
)
}
fn download_url(&self, hash: impl AsRef<str>) -> String {
format!(
"https://www.virustotal.com/vtapi/v2/file/download?apikey={}&hash={}",
self.apikey,
hash.as_ref()
)
}
fn internal_query<T>(&self, resource: impl AsRef<str>, allinfo: bool) -> GenericResult<T>
where
T: serde::de::DeserializeOwned,
{
let mut res = reqwest::get(self.file_report_url(resource, allinfo).as_str())?;
if res.status().is_success() == false {
return Err(VTError::RequestFailed.into());
}
Ok(res.json()?)
}
pub fn query_filereport_allinfo<T>(&self, resource: impl AsRef<str>) -> GenericResult<T>
where
T: serde::de::DeserializeOwned,
{
self.internal_query(resource, true)
}
pub fn get_raw_filereport_json(
&self,
resource: impl AsRef<str>,
allinfo: bool,
) -> GenericResult<String> {
let mut res = reqwest::get(self.file_report_url(resource, allinfo).as_str())?;
if res.status().is_success() == false {
return Err(VTError::RequestFailed.into());
}
Ok(res.text()?)
}
pub fn get_raw_filereport_json_at(
&self,
hash: impl TryInto<SampleHash>,
allinfo: bool,
datetime: chrono::DateTime<Utc>,
) -> GenericResult<String> {
let hash = unwrap_try_into(hash)?;
let r = scan_id(hash, datetime);
self.get_raw_filereport_json(r, allinfo)
}
pub fn query_filereport_at(
&self,
hash: impl TryInto<SampleHash>,
datetime: chrono::DateTime<Utc>,
) -> GenericResult<FileReport> {
let hash = unwrap_try_into(hash)?;
let r = scan_id(hash, datetime);
self.query_filereport(r)
}
pub fn query_filereport(
&self,
resource: impl AsRef<str>,
) -> Result<FileReport, failure::Error> {
let report: RawFileReport = self.internal_query(resource, false)?;
Ok(report.try_into()?)
}
pub fn batch_query_allinfo<T>(
&self,
resources: impl IntoIterator<Item = impl AsRef<str>>,
) -> Vec<Result<T, failure::Error>>
where
T: serde::de::DeserializeOwned,
{
resources
.into_iter()
.enumerate()
.inspect(|(idx, _)| {
if *idx != 0 {
std::thread::sleep(std::time::Duration::from_secs(1));
}
})
.map(|(_idx, item)| self.query_filereport_allinfo(item))
.collect()
}
pub fn batch_query(
&self,
resources: impl IntoIterator<Item = impl AsRef<str>>,
public_api: bool,
) -> Vec<Result<FileReport, failure::Error>> {
let sleeptime = if public_api {
std::time::Duration::from_secs(15)
} else {
std::time::Duration::from_secs(1)
};
resources
.into_iter()
.enumerate()
.inspect(|(idx, _item)| {
if *idx != 0 {
std::thread::sleep(sleeptime);
}
})
.map(|(_idx, item)| self.query_filereport(item))
.collect()
}
pub fn download(
&self,
hash: impl TryInto<SampleHash>,
into: impl AsRef<std::path::Path>,
) -> Result<(), failure::Error> {
let h = unwrap_try_into(hash)?;
let h = h.as_ref();
let mut res = reqwest::get(self.download_url(h).as_str())?;
if !res.status().is_success() {
return Err(VTError::DownloadFailed(h.to_owned()).into());
}
let mut f = std::fs::File::create(into)?;
std::io::copy(&mut res, &mut f)?;
Ok(())
}
pub fn search_by_pages(&self, query: impl AsRef<str>, goal: Option<usize>) -> Search {
Search::new(&self.apikey, query, goal)
}
pub fn search<T>(&self, query: impl AsRef<str>, goal: Option<usize>) -> T
where
T: std::iter::FromIterator<SampleHash>,
{
self.search_by_pages(query, goal)
.into_iter()
.flat_map(|x| x)
.collect()
}
}
pub struct Search {
apikey: String,
query: String,
goal: Option<usize>,
offset: Option<String>,
current: usize,
has_done: bool,
}
impl Search {
pub fn new(apikey: impl AsRef<str>, query: impl AsRef<str>, goal: Option<usize>) -> Self {
Search {
apikey: apikey.as_ref().to_owned(),
query: Search::escape_search_query(query),
offset: None,
current: 0,
has_done: false,
goal,
}
}
fn escape_search_query(query: impl AsRef<str>) -> String {
utf8_percent_encode(query.as_ref(), NON_ALPHANUMERIC).to_string()
}
fn search_url(&self, offset: &Option<String>) -> String {
match offset {
Some(o) => format!(
"https://www.virustotal.com/vtapi/v2/file/search?apikey={}&query={}&offset={}",
self.apikey.as_str(),
self.query.as_str(),
o,
),
None => format!(
"https://www.virustotal.com/vtapi/v2/file/search?apikey={}&query={}",
self.apikey.as_str(),
self.query.as_str(),
),
}
}
pub fn do_search<T>(&mut self) -> GenericResult<T>
where
T: std::iter::FromIterator<SampleHash>,
{
if self.has_done {
return Err(VTError::AlreadyReachToGoal.into());
}
let url = self.search_url(&self.offset);
let mut res = reqwest::get(url.as_str())?;
if !res.status().is_success() {
return Err(VTError::RequestFailed.into());
}
let result: SearchResponse = res.json()?;
if result.response_code != 1 {
return Err(VTError::ResponseCodeError(result.response_code).into());
}
let hashes = result.hashes.ok_or(VTError::RequestFailed)?;
if let Some(x) = self.goal {
self.current += hashes.len();
if x <= self.current {
self.has_done = true;
}
}
if result.offset.is_none() {
self.has_done = true;
}
self.offset = result.offset;
SampleHash::try_map(hashes)
}
}
impl Iterator for Search {
type Item = Vec<SampleHash>;
fn next(&mut self) -> Option<Self::Item> {
self.do_search().ok()
}
}
pub fn scan_id(sample: crate::SampleHash, datetime: impl Into<chrono::DateTime<Utc>>) -> String {
format!("{}-{}", sample.as_ref(), datetime.into().timestamp())
}
impl Default for VirusTotalClient {
fn default() -> Self {
VirusTotalClient {
apikey: std::env::var("VTAPIKEY")
.expect("please set VirusTotal API key to environment var $VTAPIKEY"),
}
}
}
#[derive(Fail, Debug)]
pub enum VTError {
#[fail(display = "VT not returned status code 1")]
ResponseCodeError(i32),
#[fail(display = "record missing field(s)")]
MissingFields(String),
#[fail(display = "download failed")]
DownloadFailed(String),
#[fail(
display = "request failed. Usually, it caused by wrong query. Or it's a Private API if you use public API key."
)]
RequestFailed,
#[fail(display = "already reach to goal")]
AlreadyReachToGoal,
}
#[derive(Deserialize, Serialize, Debug)]
pub struct ScanResult {
pub detected: bool,
pub version: Option<String>,
pub result: Option<String>,
pub update: Option<String>,
}
#[derive(Deserialize, Serialize, Debug)]
pub struct RawFileReport {
response_code: i32,
verbose_msg: String,
sha1: Option<String>,
sha256: Option<String>,
md5: Option<String>,
scan_date: Option<String>,
permalink: Option<String>,
positives: Option<u32>,
total: Option<u32>,
scans: Option<HashMap<String, ScanResult>>,
}
impl std::convert::TryInto<FileReport> for RawFileReport {
type Error = VTError;
fn try_into(self) -> Result<FileReport, Self::Error> {
if self.response_code != 1 {
return Err(VTError::ResponseCodeError(self.response_code));
}
Ok(FileReport {
sha1: self
.sha1
.ok_or(VTError::MissingFields("sha1".to_string()))?,
sha256: self
.sha256
.ok_or(VTError::MissingFields("sha256".to_string()))?,
md5: self.md5.ok_or(VTError::MissingFields("md5".to_string()))?,
scan_date: self
.scan_date
.ok_or(VTError::MissingFields("scan_date".to_string()))?,
permalink: self
.permalink
.ok_or(VTError::MissingFields("permalink".to_string()))?,
positives: self
.positives
.ok_or(VTError::MissingFields("positives".to_string()))?,
total: self
.total
.ok_or(VTError::MissingFields("total".to_string()))?,
scans: self
.scans
.ok_or(VTError::MissingFields("scans".to_string()))?,
})
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FileReport {
pub sha1: String,
pub sha256: String,
pub md5: String,
pub scan_date: String,
pub permalink: String,
pub positives: u32,
pub total: u32,
pub scans: HashMap<String, ScanResult>,
}
impl Into<ContentHash> for FileReport {
fn into(self) -> ContentHash {
ContentHash {
sha256: SampleHash::new(self.sha256).unwrap(),
sha1: SampleHash::new(self.sha1).unwrap(),
md5: SampleHash::new(self.md5).unwrap(),
}
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SearchResponse {
response_code: i32,
offset: Option<String>,
hashes: Option<Vec<String>>,
}
#[macro_export]
macro_rules! fs {
($from:expr => $to:expr) => {
format!(
"(fs:{}+ AND fs:{}-)",
$crate::datetime::vtdatetime($from),
$crate::datetime::vtdatetime($to)
)
};
($from:expr =>) => {
format!("fs:{}+", $crate::datetime::vtdatetime($from))
};
(=> $to:expr) => {
format!("fs:{}-", $crate::datetime::vtdatetime($to))
};
}
#[macro_export]
macro_rules! ls {
($from:expr => $to:expr) => {
format!(
"(ls:{}+ AND ls:{}-)",
$crate::datetime::vtdatetime($from),
$crate::datetime::vtdatetime($to)
)
};
($from:expr =>) => {
format!("ls:{}+", $crate::datetime::vtdatetime($from))
};
(=> $to:expr) => {
format!("ls:{}-", $crate::datetime::vtdatetime($to))
};
}
#[macro_export]
macro_rules! la {
($from:expr => $to:expr) => {
format!(
"(la:{}+ AND la:{}-)",
$crate::datetime::vtdatetime($from),
$crate::datetime::vtdatetime($to)
)
};
($from:expr =>) => {
format!("la:{}+", $crate::datetime::vtdatetime($from))
};
(=> $to:expr) => {
format!("la:{}-", $crate::datetime::vtdatetime($to))
};
}
#[macro_export]
macro_rules! p {
($from:expr => $to:expr) => {
format!("(p:{}+ AND p:{}-)", $from, $to)
};
($num:expr =>) => {
format!("p:{}+", $num)
};
(=> $num:expr) => {
format!("p:{}-", $num)
};
}