iocutil/
virustotal.rs

1//! VirusTotal client and its utilities
2
3use chrono::Utc;
4use failure::Fail;
5use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::convert::TryInto;
9
10use crate::contenthash::ContentHash;
11use crate::util::unwrap_try_into;
12use crate::{GenericResult, SampleHash};
13
14/// client for VirusTotal API (default use `$VTAPIKEY` environment variable as apikey)
15pub struct VirusTotalClient {
16    apikey: String,
17}
18
19impl VirusTotalClient {
20    /// new client with apikey
21    pub fn new(apikey: impl AsRef<str>) -> Self {
22        VirusTotalClient {
23            apikey: apikey.as_ref().to_owned(),
24        }
25    }
26
27    fn file_report_url(&self, resource: impl AsRef<str>, allinfo: bool) -> String {
28        format!(
29            "https://www.virustotal.com/vtapi/v2/file/report?apikey={}&allinfo={}&resource={}",
30            self.apikey,
31            allinfo,
32            resource.as_ref()
33        )
34    }
35
36    fn download_url(&self, hash: impl AsRef<str>) -> String {
37        format!(
38            "https://www.virustotal.com/vtapi/v2/file/download?apikey={}&hash={}",
39            self.apikey,
40            hash.as_ref()
41        )
42    }
43
44    fn internal_query<T>(&self, resource: impl AsRef<str>, allinfo: bool) -> GenericResult<T>
45    where
46        T: serde::de::DeserializeOwned,
47    {
48        let mut res = reqwest::get(self.file_report_url(resource, allinfo).as_str())?;
49        if res.status().is_success() == false {
50            return Err(VTError::RequestFailed.into());
51        }
52        Ok(res.json()?)
53    }
54
55    /// get file report of VirusTotal (with allinfo option)
56    ///
57    /// # Example
58    ///
59    /// ```ignore
60    /// use iocutil::prelude::*;
61    /// use serde::Deserialize;
62    ///
63    /// #[derive(Deserialize)]
64    /// struct FieldsWhatYouNeed {
65    ///     response_code: i32,
66    ///     // fields you want to retrieve
67    /// }
68    ///
69    /// let client = VirusTotalClient::default();
70    /// let sample = SampleHash::new("d41d8cd98f00b204e9800998ecf8427e").expect("failed to parse hash");
71    /// let report: FieldsWhatYouNeed = client.query_filereport_allinfo(sample).expect("failed to retrieve hash");
72    /// assert_eq!(report.response_code, 1);
73    /// ```
74    ///
75    pub fn query_filereport_allinfo<T>(&self, resource: impl AsRef<str>) -> GenericResult<T>
76    where
77        T: serde::de::DeserializeOwned,
78    {
79        self.internal_query(resource, true)
80    }
81
82    /// get raw filereport as text
83    /// # Example
84    ///
85    /// ```ignore
86    /// use iocutil::prelude::*;
87    ///
88    /// let client = VirusTotalClient::default();
89    /// let json_text = client.get_raw_filereport_json(
90    ///         "d41d8cd98f00b204e9800998ecf8427e",
91    ///         false,
92    ///     ).expect("failed to get report");
93    /// ```
94    pub fn get_raw_filereport_json(
95        &self,
96        resource: impl AsRef<str>,
97        allinfo: bool,
98    ) -> GenericResult<String> {
99        let mut res = reqwest::get(self.file_report_url(resource, allinfo).as_str())?;
100        if res.status().is_success() == false {
101            return Err(VTError::RequestFailed.into());
102        }
103        Ok(res.text()?)
104    }
105
106    /// get raw filereport json at specified datetime
107    ///
108    /// # Example
109    ///
110    /// ```ignore
111    /// use iocutil::prelude::*;
112    ///
113    /// let client = VirusTotalClient::default();
114    /// let json_text = client.get_raw_filereport_json_at(
115    ///         "d41d8cd98f00b204e9800998ecf8427e",
116    ///         false,
117    ///         days_ago(7)
118    ///     ).expect("failed to get report");
119    /// ```
120    pub fn get_raw_filereport_json_at(
121        &self,
122        hash: impl TryInto<SampleHash>,
123        allinfo: bool,
124        datetime: chrono::DateTime<Utc>,
125    ) -> GenericResult<String> {
126        let hash = unwrap_try_into(hash)?;
127        let r = scan_id(hash, datetime);
128        self.get_raw_filereport_json(r, allinfo)
129    }
130
131    /// query_filereport_at
132    ///
133    /// # Example
134    ///
135    /// ```ignore
136    /// use iocutil::prelude::*;
137    /// let client = VirusTotalClient::default();
138    ///
139    /// let report = client.query_filereport_at(
140    ///         "d41d8cd98f00b204e9800998ecf8427e",
141    ///         days_ago(7)
142    ///     ).expect("failed to query");
143    /// ```
144    pub fn query_filereport_at(
145        &self,
146        hash: impl TryInto<SampleHash>,
147        datetime: chrono::DateTime<Utc>,
148    ) -> GenericResult<FileReport> {
149        let hash = unwrap_try_into(hash)?;
150        let r = scan_id(hash, datetime);
151        self.query_filereport(r)
152    }
153
154    /// query file report (without allinfo)
155    ///
156    /// # Example
157    ///
158    /// ```ignore
159    /// use iocutil::prelude::*;
160    ///
161    /// let client = VirusTotalClient::default();
162    ///
163    /// let report = client.query_filereport("d41d8cd98f00b204e9800998ecf8427e").unwrap();
164    /// ```
165    pub fn query_filereport(
166        &self,
167        resource: impl AsRef<str>,
168    ) -> Result<FileReport, failure::Error> {
169        let report: RawFileReport = self.internal_query(resource, false)?;
170        Ok(report.try_into()?)
171    }
172
173    /// batch query file report
174    ///
175    /// # Example
176    ///
177    /// ```ignore
178    /// use iocutil::prelude::*;
179    /// use serde::Deserialize;
180    ///
181    /// #[derive(Deserialize)]
182    /// struct FieldsWhatYouNeed {
183    ///     response_code: i32,
184    ///     // fields you want to retrieve
185    /// }
186    ///
187    /// let vtclient = VirusTotalClient::default();
188    /// let hashes = &["d41d8cd98f00b204e9800998ecf8427e"];
189    /// let items: Vec<Result<FieldsWhatYouNeed, failure::Error>> = vtclient.batch_query_allinfo(hashes);
190    /// for item in items {
191    ///     item.expect("failed to retrieve");
192    /// }
193    /// ```
194    pub fn batch_query_allinfo<T>(
195        &self,
196        resources: impl IntoIterator<Item = impl AsRef<str>>,
197    ) -> Vec<Result<T, failure::Error>>
198    where
199        T: serde::de::DeserializeOwned,
200    {
201        resources
202            .into_iter()
203            .enumerate()
204            .inspect(|(idx, _)| {
205                if *idx != 0 {
206                    std::thread::sleep(std::time::Duration::from_secs(1));
207                }
208            })
209            .map(|(_idx, item)| self.query_filereport_allinfo(item))
210            .collect()
211    }
212
213    /// batch query file report
214    ///
215    /// # Example
216    ///
217    /// ```ignore
218    /// use iocutil::prelude::*;
219    ///
220    /// let vtclient = VirusTotalClient::default();
221    /// let hashes = &["e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"];
222    /// let items = vtclient.batch_query(hashes, true);
223    /// for item in items {
224    ///     item.expect("failed to retrieve");
225    /// }
226    /// ```
227    ///
228    pub fn batch_query(
229        &self,
230        resources: impl IntoIterator<Item = impl AsRef<str>>,
231        public_api: bool,
232    ) -> Vec<Result<FileReport, failure::Error>> {
233        let sleeptime = if public_api {
234            std::time::Duration::from_secs(15)
235        } else {
236            std::time::Duration::from_secs(1)
237        };
238        resources
239            .into_iter()
240            .enumerate()
241            .inspect(|(idx, _item)| {
242                if *idx != 0 {
243                    std::thread::sleep(sleeptime);
244                }
245            })
246            .map(|(_idx, item)| self.query_filereport(item))
247            .collect()
248    }
249
250    /// download a file from hash
251    ///
252    /// # Example
253    ///
254    /// ```ignore
255    /// use iocutil::prelude::*;
256    ///
257    /// let client = VirusTotalClient::default();
258    /// client.download(
259    ///         "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
260    ///         "./e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
261    ///     ).expect("failed to download file");
262    ///
263    /// std::fs::remove_file("./e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
264    ///     .expect("failed to remove file");
265    /// ```
266    pub fn download(
267        &self,
268        hash: impl TryInto<SampleHash>,
269        into: impl AsRef<std::path::Path>,
270    ) -> Result<(), failure::Error> {
271        let h = unwrap_try_into(hash)?;
272        let h = h.as_ref();
273
274        let mut res = reqwest::get(self.download_url(h).as_str())?;
275        if !res.status().is_success() {
276            return Err(VTError::DownloadFailed(h.to_owned()).into());
277        }
278
279        let mut f = std::fs::File::create(into)?;
280        std::io::copy(&mut res, &mut f)?;
281
282        Ok(())
283    }
284
285    /// search by page (Private API required)
286    /// https://www.virustotal.com/intelligence/help/file-search/#search-modifiers
287    ///
288    /// # Example
289    ///
290    /// ```ignore
291    /// use iocutil::prelude::*;
292    ///
293    /// let client = VirusTotalClient::default();
294    /// let mut pages = client.search_by_pages("p:5+ AND submitter:CN", Some(600));
295    ///
296    /// let samples: Vec<_> = pages.do_search().expect("failed to search");
297    /// assert_eq!(samples.len(), 300)
298    /// ```
299    pub fn search_by_pages(&self, query: impl AsRef<str>, goal: Option<usize>) -> Search {
300        Search::new(&self.apikey, query, goal)
301    }
302
303    /// search samples (Private API required)
304    /// https://www.virustotal.com/intelligence/help/file-search/#search-modifiers
305    ///
306    /// # Example
307    ///
308    /// ```ignore
309    /// use iocutil::prelude::*;
310    ///
311    /// let client = VirusTotalClient::default();
312    ///
313    /// let samples: Vec<_> = client.search("p:5+ AND submitter:CN", Some(600));
314    /// assert_eq!(samples.len(), 600)
315    /// ```
316    pub fn search<T>(&self, query: impl AsRef<str>, goal: Option<usize>) -> T
317    where
318        T: std::iter::FromIterator<SampleHash>,
319    {
320        self.search_by_pages(query, goal)
321            .into_iter()
322            .flat_map(|x| x)
323            .collect()
324    }
325}
326
327/// context object for search api
328pub struct Search {
329    apikey: String,
330    query: String,
331    goal: Option<usize>,
332    offset: Option<String>,
333    current: usize,
334    has_done: bool,
335}
336
337impl Search {
338    /// create new object
339    pub fn new(apikey: impl AsRef<str>, query: impl AsRef<str>, goal: Option<usize>) -> Self {
340        Search {
341            apikey: apikey.as_ref().to_owned(),
342            query: Search::escape_search_query(query),
343            offset: None,
344            current: 0,
345            has_done: false,
346            goal,
347        }
348    }
349
350    fn escape_search_query(query: impl AsRef<str>) -> String {
351        utf8_percent_encode(query.as_ref(), NON_ALPHANUMERIC).to_string()
352    }
353
354    fn search_url(&self, offset: &Option<String>) -> String {
355        match offset {
356            Some(o) => format!(
357                "https://www.virustotal.com/vtapi/v2/file/search?apikey={}&query={}&offset={}",
358                self.apikey.as_str(),
359                self.query.as_str(),
360                o,
361            ),
362            None => format!(
363                "https://www.virustotal.com/vtapi/v2/file/search?apikey={}&query={}",
364                self.apikey.as_str(),
365                self.query.as_str(),
366            ),
367        }
368    }
369
370    /// request once (most 300 samples per a page)
371    pub fn do_search<T>(&mut self) -> GenericResult<T>
372    where
373        T: std::iter::FromIterator<SampleHash>,
374    {
375        if self.has_done {
376            return Err(VTError::AlreadyReachToGoal.into());
377        }
378
379        let url = self.search_url(&self.offset);
380
381        let mut res = reqwest::get(url.as_str())?;
382        if !res.status().is_success() {
383            return Err(VTError::RequestFailed.into());
384        }
385
386        let result: SearchResponse = res.json()?;
387        if result.response_code != 1 {
388            return Err(VTError::ResponseCodeError(result.response_code).into());
389        }
390
391        let hashes = result.hashes.ok_or(VTError::RequestFailed)?;
392
393        if let Some(x) = self.goal {
394            self.current += hashes.len();
395            if x <= self.current {
396                self.has_done = true;
397            }
398        }
399
400        if result.offset.is_none() {
401            self.has_done = true;
402        }
403
404        self.offset = result.offset;
405
406        SampleHash::try_map(hashes)
407    }
408}
409
410impl Iterator for Search {
411    type Item = Vec<SampleHash>;
412
413    fn next(&mut self) -> Option<Self::Item> {
414        self.do_search().ok()
415    }
416}
417
418/// scan_id for virustotal
419/// You can use this to get a report at the specific time.
420///
421/// # Example
422///
423/// ```
424/// use iocutil::prelude::*;
425///
426/// let client = VirusTotalClient::default();
427///
428/// let sample =  SampleHash::new("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855").unwrap();
429/// let sid = scan_id(sample, at!(1, weeks ago));
430///
431/// // let report_at_one_week_ago = client.query_filereport(sid).unwrap();
432/// ```
433pub fn scan_id(sample: crate::SampleHash, datetime: impl Into<chrono::DateTime<Utc>>) -> String {
434    format!("{}-{}", sample.as_ref(), datetime.into().timestamp())
435}
436
437impl Default for VirusTotalClient {
438    fn default() -> Self {
439        VirusTotalClient {
440            apikey: std::env::var("VTAPIKEY")
441                .expect("please set VirusTotal API key to environment var $VTAPIKEY"),
442        }
443    }
444}
445
446/// Errors in VirusTotal operation
447#[derive(Fail, Debug)]
448pub enum VTError {
449    #[fail(display = "VT not returned status code 1")]
450    ResponseCodeError(i32),
451
452    #[fail(display = "record missing field(s)")]
453    MissingFields(String),
454
455    #[fail(display = "download failed")]
456    DownloadFailed(String),
457
458    #[fail(
459        display = "request failed. Usually, it caused by wrong query. Or it's a Private API if you use public API key."
460    )]
461    RequestFailed,
462
463    #[fail(display = "already reach to goal")]
464    AlreadyReachToGoal,
465}
466
467/// ScanResult item of "scans"
468#[derive(Deserialize, Serialize, Debug)]
469pub struct ScanResult {
470    pub detected: bool,
471    pub version: Option<String>,
472    pub result: Option<String>,
473    pub update: Option<String>,
474}
475
476/// RawFileReport structure (without fields included only in allinfo option)
477#[derive(Deserialize, Serialize, Debug)]
478pub struct RawFileReport {
479    response_code: i32,
480    verbose_msg: String,
481    sha1: Option<String>,
482    sha256: Option<String>,
483    md5: Option<String>,
484    scan_date: Option<String>,
485    permalink: Option<String>,
486    positives: Option<u32>,
487    total: Option<u32>,
488    scans: Option<HashMap<String, ScanResult>>,
489}
490
491impl std::convert::TryInto<FileReport> for RawFileReport {
492    type Error = VTError;
493
494    fn try_into(self) -> Result<FileReport, Self::Error> {
495        if self.response_code != 1 {
496            // virustotal returns reposnse code 1 when succeeded to retrieve scan result.
497            return Err(VTError::ResponseCodeError(self.response_code));
498        }
499
500        Ok(FileReport {
501            sha1: self
502                .sha1
503                .ok_or(VTError::MissingFields("sha1".to_string()))?,
504            sha256: self
505                .sha256
506                .ok_or(VTError::MissingFields("sha256".to_string()))?,
507            md5: self.md5.ok_or(VTError::MissingFields("md5".to_string()))?,
508            scan_date: self
509                .scan_date
510                .ok_or(VTError::MissingFields("scan_date".to_string()))?,
511            permalink: self
512                .permalink
513                .ok_or(VTError::MissingFields("permalink".to_string()))?,
514            positives: self
515                .positives
516                .ok_or(VTError::MissingFields("positives".to_string()))?,
517            total: self
518                .total
519                .ok_or(VTError::MissingFields("total".to_string()))?,
520            scans: self
521                .scans
522                .ok_or(VTError::MissingFields("scans".to_string()))?,
523        })
524    }
525}
526
527/// FileReport (without fields included only in allinfo option)
528#[derive(Debug, Serialize, Deserialize)]
529pub struct FileReport {
530    pub sha1: String,
531    pub sha256: String,
532    pub md5: String,
533    pub scan_date: String,
534    pub permalink: String,
535    pub positives: u32,
536    pub total: u32,
537    pub scans: HashMap<String, ScanResult>,
538}
539
540impl Into<ContentHash> for FileReport {
541    fn into(self) -> ContentHash {
542        ContentHash {
543            sha256: SampleHash::new(self.sha256).unwrap(),
544            sha1: SampleHash::new(self.sha1).unwrap(),
545            md5: SampleHash::new(self.md5).unwrap(),
546        }
547    }
548}
549
550#[derive(Debug, Serialize, Deserialize)]
551pub struct SearchResponse {
552    response_code: i32,
553    offset: Option<String>,
554    hashes: Option<Vec<String>>,
555}
556
557/// first submission search modifier macro
558///
559/// # Example
560///
561/// ```
562/// use iocutil::prelude::*;
563///
564/// let f = day!(2019, 11, 01).unwrap();
565/// let t = day!(2019, 11, 02).unwrap();
566/// let fs1 = fs!(f => t);
567/// assert_eq!(fs1.as_str(), "(fs:2019-11-01T00:00:00+ AND fs:2019-11-02T00:00:00-)");
568///
569/// let fs2 = fs!(f =>);
570/// assert_eq!(fs2.as_str(), "fs:2019-11-01T00:00:00+");
571///
572/// let fs3 = fs!(=> t);
573/// assert_eq!(fs3.as_str(), "fs:2019-11-02T00:00:00-");
574///
575/// let for_a_week = fs!(at!(1, weeks ago) =>);
576/// ```
577#[macro_export]
578macro_rules! fs {
579    ($from:expr => $to:expr) => {
580        format!(
581            "(fs:{}+ AND fs:{}-)",
582            $crate::datetime::vtdatetime($from),
583            $crate::datetime::vtdatetime($to)
584        )
585    };
586    ($from:expr =>) => {
587        format!("fs:{}+", $crate::datetime::vtdatetime($from))
588    };
589    (=> $to:expr) => {
590        format!("fs:{}-", $crate::datetime::vtdatetime($to))
591    };
592}
593
594/// last submission search modifier macro
595///
596/// # Example
597///
598/// ```
599/// use iocutil::prelude::*;
600///
601/// let f = day!(2019, 11, 01).unwrap();
602/// let t = day!(2019, 11, 02).unwrap();
603/// let ls1 = ls!(f => t);
604/// assert_eq!(ls1.as_str(), "(ls:2019-11-01T00:00:00+ AND ls:2019-11-02T00:00:00-)");
605///
606/// let ls2 = ls!(f =>);
607/// assert_eq!(ls2.as_str(), "ls:2019-11-01T00:00:00+");
608///
609/// let ls3 = ls!(=> t);
610/// assert_eq!(ls3.as_str(), "ls:2019-11-02T00:00:00-");
611///
612/// let for_a_week = ls!(at!(1, weeks ago) =>);
613/// ```
614#[macro_export]
615macro_rules! ls {
616    ($from:expr => $to:expr) => {
617        format!(
618            "(ls:{}+ AND ls:{}-)",
619            $crate::datetime::vtdatetime($from),
620            $crate::datetime::vtdatetime($to)
621        )
622    };
623    ($from:expr =>) => {
624        format!("ls:{}+", $crate::datetime::vtdatetime($from))
625    };
626    (=> $to:expr) => {
627        format!("ls:{}-", $crate::datetime::vtdatetime($to))
628    };
629}
630
631/// last analysis search modifier macro
632///
633/// # Example
634///
635/// ```
636/// use iocutil::prelude::*;
637///
638/// let f = day!(2019, 11, 01).unwrap();
639/// let t = day!(2019, 11, 02).unwrap();
640///
641/// let la1 = la!(f => t);
642/// assert_eq!(la1.as_str(), "(la:2019-11-01T00:00:00+ AND la:2019-11-02T00:00:00-)");
643///
644/// let la2 = la!(f =>);
645/// assert_eq!(la2.as_str(), "la:2019-11-01T00:00:00+");
646///
647/// let la3 = la!(=> t);
648/// assert_eq!(la3.as_str(), "la:2019-11-02T00:00:00-");
649///
650/// let for_a_week = la!(at!(1, weeks ago) =>);
651/// ```
652#[macro_export]
653macro_rules! la {
654    ($from:expr => $to:expr) => {
655        format!(
656            "(la:{}+ AND la:{}-)",
657            $crate::datetime::vtdatetime($from),
658            $crate::datetime::vtdatetime($to)
659        )
660    };
661    ($from:expr =>) => {
662        format!("la:{}+", $crate::datetime::vtdatetime($from))
663    };
664    (=> $to:expr) => {
665        format!("la:{}-", $crate::datetime::vtdatetime($to))
666    };
667}
668
669/// positive numbers search modifier macro
670///
671/// # Example
672///
673/// ```
674/// use iocutil::prelude::*;
675///
676/// let p1 = p!(1 => 10);
677/// assert_eq!(p1.as_str(), "(p:1+ AND p:10-)");
678///
679/// let p2 = p!(1 =>);
680/// assert_eq!(p2.as_str(), "p:1+");
681///
682/// let p3 = p!(=> 10);
683/// assert_eq!(p3.as_str(), "p:10-");
684/// ```
685#[macro_export]
686macro_rules! p {
687    ($from:expr => $to:expr) => {
688        format!("(p:{}+ AND p:{}-)", $from, $to)
689    };
690    ($num:expr =>) => {
691        format!("p:{}+", $num)
692    };
693    (=> $num:expr) => {
694        format!("p:{}-", $num)
695    };
696}
iocutil/virustotal.rs

iocutil/
virustotal.rs
