malwaredb_client/
blocking.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::fmt::{Debug, Formatter};
4use std::fs::OpenOptions;
5use std::io::Write;
6use std::path::{Path, PathBuf};
7
8use crate::{MDB_CLIENT_ERROR_CONTEXT, get_config_path};
9use malwaredb_api::{
10    GetAPIKeyResponse, GetUserInfoResponse, Labels, PartialHashSearchType, Report, SearchRequest,
11    SearchRequestParameters, SearchResponse, SearchType, ServerInfo, ServerResponse,
12    SimilarSamplesResponse, Sources, SupportedFileTypes, YaraSearchRequest,
13    YaraSearchRequestResponse, YaraSearchResponse, digest::HashType,
14};
15use malwaredb_types::exec::pe32::EXE;
16
17use anyhow::{Context, Result, bail, ensure};
18use base64::Engine;
19use base64::engine::general_purpose;
20use fuzzyhash::FuzzyHash;
21use malwaredb_lzjd2::lzjd::LzDigest;
22use serde::{Deserialize, Serialize};
23use sha2::{Digest, Sha256};
24use tlsh_fixed::TlshBuilder;
25use tracing::{error, info, trace, warn};
26use uuid::Uuid;
27use zeroize::{Zeroize, ZeroizeOnDrop};
28
29/// Blocking Malware DB Client Configuration and connection which requires the `blocking` feature
30#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
31#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
32pub struct MdbClient {
33    /// URL of the Malware DB server, including http and port number, ending without a slash
34    pub url: String,
35
36    /// User's API key for Malware DB
37    api_key: String,
38
39    /// Blocking http client which stores the optional server certificate
40    #[zeroize(skip)]
41    #[serde(skip)]
42    client: reqwest::blocking::Client,
43
44    /// Server's certificate
45    #[cfg(target_os = "macos")]
46    #[zeroize(skip)]
47    #[serde(skip)]
48    cert: Option<crate::macos::CertificateData>,
49}
50
51impl MdbClient {
52    /// MDB Client from components, doesn't test connectivity
53    ///
54    /// # Errors
55    ///
56    /// Returns an error if a list of certificates was passed and any were not in the expected
57    /// DER or PEM format or could not be parsed.
58    ///
59    /// # Panics
60    ///
61    /// This method panics if called from within an async runtime.
62    pub fn new(url: String, api_key: String, cert_path: Option<PathBuf>) -> Result<Self> {
63        let mut url = url;
64        let url = if url.ends_with('/') {
65            url.pop();
66            url
67        } else {
68            url
69        };
70
71        let cert = if let Some(path) = cert_path {
72            Some((crate::path_load_cert(&path)?, path))
73        } else {
74            None
75        };
76
77        let builder = reqwest::blocking::ClientBuilder::new()
78            .gzip(true)
79            .zstd(true)
80            .use_rustls_tls()
81            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
82
83        let client = if let Some(((_cert_type, cert), _path)) = &cert {
84            builder.add_root_certificate(cert.clone()).build()
85        } else {
86            builder.build()
87        }?;
88
89        #[cfg(target_os = "macos")]
90        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
91            Some(crate::macos::CertificateData {
92                cert_type: *cert_type,
93                cert_bytes: std::fs::read(cert_path)?,
94            })
95        } else {
96            None
97        };
98
99        Ok(Self {
100            url,
101            api_key,
102            client,
103
104            #[cfg(target_os = "macos")]
105            cert,
106        })
107    }
108
109    /// Connect to a server anonymously.
110    ///
111    /// # Errors
112    ///
113    /// Networking errors may result, and there will be an error if the server doesn't support anonymous connections.
114    pub fn anonymous(url: String, save: bool, cert_path: Option<PathBuf>) -> Result<Self> {
115        let mut url = url;
116        let url = if url.ends_with('/') {
117            url.pop();
118            url
119        } else {
120            url
121        };
122
123        let builder = reqwest::blocking::ClientBuilder::new()
124            .gzip(true)
125            .zstd(true)
126            .use_rustls_tls()
127            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
128
129        let cert = if let Some(path) = cert_path {
130            Some((crate::path_load_cert(&path)?, path))
131        } else {
132            None
133        };
134
135        let client = if let Some(((_cert_type, cert), _cert_path)) = &cert {
136            builder.add_root_certificate(cert.clone()).build()
137        } else {
138            builder.build()
139        }?;
140
141        let info = client
142            .get(format!("{url}{}", malwaredb_api::USER_INFO_URL))
143            .send()?
144            .json::<ServerResponse<GetUserInfoResponse>>()
145            .context(MDB_CLIENT_ERROR_CONTEXT)?;
146
147        if let ServerResponse::Success(info) = info {
148            ensure!(info.id > 0);
149        } else {
150            bail!("Anonymous access is not allowed");
151        }
152
153        #[cfg(target_os = "macos")]
154        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
155            Some(crate::macos::CertificateData {
156                cert_type: *cert_type,
157                cert_bytes: std::fs::read(cert_path)?,
158            })
159        } else {
160            None
161        };
162
163        let client = Self {
164            url,
165            api_key: "anonymous".to_string(),
166            client,
167
168            #[cfg(target_os = "macos")]
169            cert,
170        };
171
172        let server_info = client.server_info()?;
173        if server_info.mdb_version > *crate::MDB_VERSION_SEMVER {
174            warn!(
175                "Server version {:?} is newer than client {:?}, consider updating.",
176                server_info.mdb_version,
177                crate::MDB_VERSION_SEMVER
178            );
179        }
180
181        if save && let Err(e) = client.save() {
182            error!("Anonymous connection successful but failed to save config: {e}");
183            bail!("Anonymous connection successful but failed to save config: {e}");
184        }
185
186        Ok(client)
187    }
188
189    /// Login to a server, optionally save the configuration file, and return a client object
190    ///
191    /// # Errors
192    ///
193    /// Returns an error if the server URL, username, or password were incorrect, or if a network
194    /// issue occurred.
195    ///
196    /// # Panics
197    ///
198    /// This method panics if called from within an async runtime.
199    pub fn login(
200        url: String,
201        username: String,
202        password: String,
203        save: bool,
204        cert_path: Option<PathBuf>,
205    ) -> Result<Self> {
206        let mut url = url;
207        let url = if url.ends_with('/') {
208            url.pop();
209            url
210        } else {
211            url
212        };
213
214        let api_request = malwaredb_api::GetAPIKeyRequest {
215            user: username,
216            password,
217        };
218
219        let builder = reqwest::blocking::ClientBuilder::new()
220            .gzip(true)
221            .zstd(true)
222            .use_rustls_tls()
223            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
224
225        let cert = if let Some(path) = cert_path {
226            Some((crate::path_load_cert(&path)?, path))
227        } else {
228            None
229        };
230
231        let client = if let Some(((_cert_type, cert), _path)) = &cert {
232            builder.add_root_certificate(cert.clone()).build()
233        } else {
234            builder.build()
235        }?;
236
237        let res = client
238            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
239            .json(&api_request)
240            .send()?
241            .json::<ServerResponse<GetAPIKeyResponse>>()
242            .context(MDB_CLIENT_ERROR_CONTEXT)?;
243
244        let res = match res {
245            ServerResponse::Success(res) => res,
246            ServerResponse::Error(err) => return Err(err.into()),
247        };
248
249        #[cfg(target_os = "macos")]
250        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
251            Some(crate::macos::CertificateData {
252                cert_type: *cert_type,
253                cert_bytes: std::fs::read(cert_path)?,
254            })
255        } else {
256            None
257        };
258
259        let client = MdbClient {
260            url,
261            api_key: res.key.clone(),
262            client,
263
264            #[cfg(target_os = "macos")]
265            cert,
266        };
267
268        let server_info = client.server_info()?;
269        if server_info.mdb_version > *crate::MDB_VERSION_SEMVER {
270            warn!(
271                "Server version {:?} is newer than client {:?}, consider updating.",
272                server_info.mdb_version,
273                crate::MDB_VERSION_SEMVER
274            );
275        }
276
277        if save && let Err(e) = client.save() {
278            error!("Login successful but failed to save config: {e}");
279            bail!("Login successful but failed to save config: {e}");
280        }
281        Ok(client)
282    }
283
284    /// Reset one's own API key to effectively logout & disable all clients who are using the key
285    ///
286    /// # Errors
287    ///
288    /// Returns an error if there was a network issue or the user wasn't properly logged in.
289    pub fn reset_key(&self) -> Result<()> {
290        let response = self
291            .client
292            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
293            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
294            .send()
295            .context(MDB_CLIENT_ERROR_CONTEXT)?;
296        if !response.status().is_success() {
297            bail!("failed to reset API key, was it correct?");
298        }
299        Ok(())
300    }
301
302    /// Malware DB Client configuration loaded from a specified path
303    ///
304    /// # Errors
305    ///
306    /// Returns an error if the configuration file cannot be read, possibly because it
307    /// doesn't exist or due to a permission error or a parsing error.
308    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
309        let name = path.as_ref().display();
310        let config =
311            std::fs::read_to_string(&path).context(format!("failed to read config file {name}"))?;
312        let cfg: MdbClient =
313            toml::from_str(&config).context(format!("failed to parse config file {name}"))?;
314        Ok(cfg)
315    }
316
317    /// Malware DB Client configuration from user's home directory
318    ///
319    /// On macOS, it will attempt to load this information in the Keychain, which isn't required.
320    ///
321    /// # Errors
322    ///
323    /// Returns an error if the configuration file cannot be read, possibly because it
324    /// doesn't exist or due to a permission error or a parsing error.
325    ///
326    /// # Panics
327    ///
328    /// This method panics if called from within an async runtime.
329    pub fn load() -> Result<Self> {
330        #[cfg(target_os = "macos")]
331        {
332            if let Ok((api_key, url, cert)) = crate::macos::retrieve_credentials() {
333                let builder = reqwest::blocking::ClientBuilder::new()
334                    .gzip(true)
335                    .zstd(true)
336                    .use_rustls_tls()
337                    .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
338
339                let client = if let Some(cert) = &cert {
340                    builder.add_root_certificate(cert.as_cert()?).build()
341                } else {
342                    builder.build()
343                }?;
344
345                return Ok(Self {
346                    url,
347                    api_key,
348                    client,
349                    cert,
350                });
351            }
352        }
353
354        let path = get_config_path(false)?;
355        if path.exists() {
356            return Self::from_file(path);
357        }
358        bail!("config file not found")
359    }
360
361    /// Save Malware DB Client configuration to the user's home directory
362    ///
363    /// On macOS, it will attempt to save this information in the Keychain, which isn't required.
364    ///
365    /// # Errors
366    ///
367    /// Returns an error if there was a problem saving the configuration file.
368    pub fn save(&self) -> Result<()> {
369        #[cfg(target_os = "macos")]
370        {
371            if crate::macos::save_credentials(&self.url, &self.api_key, self.cert.clone()).is_ok() {
372                return Ok(());
373            }
374        }
375
376        let toml = toml::to_string(self)?;
377        let path = get_config_path(true)?;
378
379        let mut options = OpenOptions::new();
380        options
381            .write(true)
382            .create(true)
383            .append(false)
384            .truncate(false);
385
386        #[cfg(target_family = "unix")]
387        {
388            use std::os::unix::fs::OpenOptionsExt;
389
390            options.mode(0o600);
391        }
392
393        let mut file = options.open(&path)?;
394        write!(file, "{toml}").context(format!("failed to write mdb config to {}", path.display()))
395    }
396
397    /// Delete the Malware DB client configuration file
398    ///
399    /// # Errors
400    ///
401    /// Returns an error if there isn't a configuration file to delete, or if it cannot be deleted,
402    /// possibly due to a permissions error.
403    pub fn delete(&self) -> Result<()> {
404        #[cfg(target_os = "macos")]
405        crate::macos::clear_credentials();
406
407        let path = get_config_path(false)?;
408        if path.exists() {
409            std::fs::remove_file(&path)
410                .context(format!("failed to delete client config file {}", path.display()))?;
411        }
412        Ok(())
413    }
414
415    // Actions of the client
416
417    /// Get information about the server, unauthenticated
418    ///
419    /// # Errors
420    ///
421    /// This may return an error if there's a network situation.
422    pub fn server_info(&self) -> Result<ServerInfo> {
423        let response = self
424            .client
425            .get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO_URL))
426            .send()?
427            .json::<ServerResponse<ServerInfo>>()
428            .context(MDB_CLIENT_ERROR_CONTEXT)?;
429
430        match response {
431            ServerResponse::Success(info) => Ok(info),
432            ServerResponse::Error(e) => Err(e.into()),
433        }
434    }
435
436    /// Get file types supported by the server, unauthenticated
437    ///
438    /// # Errors
439    ///
440    /// This may return an error if there's a network situation.
441    pub fn supported_types(&self) -> Result<SupportedFileTypes> {
442        let response = self
443            .client
444            .get(format!("{}{}", self.url, malwaredb_api::SUPPORTED_FILE_TYPES_URL))
445            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
446            .send()?
447            .json::<ServerResponse<SupportedFileTypes>>()
448            .context(MDB_CLIENT_ERROR_CONTEXT)?;
449
450        match response {
451            ServerResponse::Success(types) => Ok(types),
452            ServerResponse::Error(e) => Err(e.into()),
453        }
454    }
455
456    /// Get information about the user
457    ///
458    /// # Errors
459    ///
460    /// This may return an error if there's a network situation or if the user is not logged in
461    /// or not properly authorized to connect.
462    pub fn whoami(&self) -> Result<GetUserInfoResponse> {
463        let response = self
464            .client
465            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
466            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
467            .send()?
468            .json::<ServerResponse<GetUserInfoResponse>>()
469            .context(MDB_CLIENT_ERROR_CONTEXT)?;
470
471        match response {
472            ServerResponse::Success(info) => Ok(info),
473            ServerResponse::Error(e) => Err(e.into()),
474        }
475    }
476
477    /// Get the sample labels known to the server
478    ///
479    /// # Errors
480    ///
481    /// This may return an error if there's a network situation or if the user is not logged in
482    /// or not properly authorized to connect.
483    pub fn labels(&self) -> Result<Labels> {
484        let response = self
485            .client
486            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS_URL))
487            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
488            .send()?
489            .json::<ServerResponse<Labels>>()
490            .context(MDB_CLIENT_ERROR_CONTEXT)?;
491
492        match response {
493            ServerResponse::Success(labels) => Ok(labels),
494            ServerResponse::Error(e) => Err(e.into()),
495        }
496    }
497
498    /// Get the sources available to the current user
499    ///
500    /// # Errors
501    ///
502    /// This may return an error if there's a network situation or if the user is not logged in
503    /// or not properly authorized to connect.
504    pub fn sources(&self) -> Result<Sources> {
505        let response = self
506            .client
507            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES_URL))
508            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
509            .send()?
510            .json::<ServerResponse<Sources>>()
511            .context(MDB_CLIENT_ERROR_CONTEXT)?;
512
513        match response {
514            ServerResponse::Success(sources) => Ok(sources),
515            ServerResponse::Error(e) => Err(e.into()),
516        }
517    }
518
519    /// Submit one file to Malware DB: provide the contents, file name, and source ID
520    ///
521    /// # Errors
522    ///
523    /// This may return an error if there's a network situation or if the user is not logged in
524    /// or not properly authorized to connect.
525    pub fn submit(
526        &self,
527        contents: impl AsRef<[u8]>,
528        file_name: String,
529        source_id: u32,
530    ) -> Result<bool> {
531        let mut hasher = Sha256::new();
532        hasher.update(&contents);
533        let result = hasher.finalize();
534
535        let encoded = general_purpose::STANDARD.encode(contents);
536
537        let payload = malwaredb_api::NewSampleB64 {
538            file_name,
539            source_id,
540            file_contents_b64: encoded,
541            sha256: hex::encode(result),
542        };
543
544        match self
545            .client
546            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE_JSON_URL))
547            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
548            .json(&payload)
549            .send()
550        {
551            Ok(res) => {
552                if !res.status().is_success() {
553                    info!("Code {} sending {}", res.status(), payload.file_name);
554                }
555                Ok(res.status().is_success())
556            }
557            Err(e) => {
558                let status: String = e
559                    .status()
560                    .map(|s| s.as_str().to_string())
561                    .unwrap_or_default();
562                error!("Error{status} sending {}: {e}", payload.file_name);
563                bail!(e.to_string())
564            }
565        }
566    }
567
568    /// Submit one file to Malware DB: provide the contents, file name, and source ID
569    /// Experimental! May be removed at any point.
570    ///
571    /// # Errors
572    ///
573    /// This may return an error if there's a network situation or if the user is not logged in
574    /// or not properly authorized to connect.
575    pub fn submit_as_cbor(
576        &self,
577        contents: impl AsRef<[u8]>,
578        file_name: String,
579        source_id: u32,
580    ) -> Result<bool> {
581        let mut hasher = Sha256::new();
582        hasher.update(&contents);
583        let result = hasher.finalize();
584
585        let payload = malwaredb_api::NewSampleBytes {
586            file_name,
587            source_id,
588            file_contents: contents.as_ref().to_vec(),
589            sha256: hex::encode(result),
590        };
591
592        let mut bytes = Vec::with_capacity(payload.file_contents.len());
593        ciborium::ser::into_writer(&payload, &mut bytes)?;
594
595        match self
596            .client
597            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE_CBOR_URL))
598            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
599            .header("content-type", "application/cbor")
600            .body(bytes)
601            .send()
602        {
603            Ok(res) => {
604                if !res.status().is_success() {
605                    info!("Code {} sending {}", res.status(), payload.file_name);
606                }
607                Ok(res.status().is_success())
608            }
609            Err(e) => {
610                let status: String = e
611                    .status()
612                    .map(|s| s.as_str().to_string())
613                    .unwrap_or_default();
614                error!("Error{status} sending {}: {e}", payload.file_name);
615                bail!(e.to_string())
616            }
617        }
618    }
619
620    /// Search for a file based on partial hash and/or partial file name, returns a list of hashes
621    ///
622    /// # Errors
623    ///
624    /// * This may return an error if there's a network situation or if the user is not logged in or the request isn't valid
625    pub fn partial_search(
626        &self,
627        partial_hash: Option<(PartialHashSearchType, String)>,
628        name: Option<String>,
629        response: PartialHashSearchType,
630        limit: u32,
631    ) -> Result<SearchResponse> {
632        let query = SearchRequest {
633            search: SearchType::Search(SearchRequestParameters {
634                partial_hash,
635                file_name: name,
636                response,
637                limit,
638                labels: None,
639                file_type: None,
640                magic: None,
641            }),
642        };
643
644        self.do_search_request(&query)
645    }
646
647    /// Search for a file based on partial hash and/or partial file name, labels, file type; returns a list of hashes
648    ///
649    /// # Errors
650    ///
651    /// * This may return an error if there's a network situation or if the user is not logged in or the request isn't valid
652    #[allow(clippy::too_many_arguments)]
653    pub fn partial_search_labels_type(
654        &self,
655        partial_hash: Option<(PartialHashSearchType, String)>,
656        name: Option<String>,
657        response: PartialHashSearchType,
658        labels: Option<Vec<String>>,
659        file_type: Option<String>,
660        magic: Option<String>,
661        limit: u32,
662    ) -> Result<SearchResponse> {
663        let query = SearchRequest {
664            search: SearchType::Search(SearchRequestParameters {
665                partial_hash,
666                file_name: name,
667                response,
668                limit,
669                file_type,
670                magic,
671                labels,
672            }),
673        };
674
675        self.do_search_request(&query)
676    }
677
678    /// Return the next page from the search result
679    ///
680    /// # Errors
681    ///
682    /// Returns an error if there is a network problem, or pagination not available
683    pub fn next_page_search(&self, response: &SearchResponse) -> Result<SearchResponse> {
684        if let Some(uuid) = response.pagination {
685            let request = SearchRequest {
686                search: SearchType::Continuation(uuid),
687            };
688            return self.do_search_request(&request);
689        }
690
691        bail!("Pagination not available")
692    }
693
694    fn do_search_request(&self, query: &SearchRequest) -> Result<SearchResponse> {
695        ensure!(
696            query.is_valid(),
697            "Query isn't valid: hash isn't hexidecimal or both the hashes and file name are empty"
698        );
699
700        let response = self
701            .client
702            .post(format!("{}{}", self.url, malwaredb_api::SEARCH_URL))
703            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
704            .json(query)
705            .send()?
706            .json::<ServerResponse<SearchResponse>>()
707            .context(MDB_CLIENT_ERROR_CONTEXT)?;
708
709        match response {
710            ServerResponse::Success(search) => Ok(search),
711            ServerResponse::Error(e) => Err(e.into()),
712        }
713    }
714
715    /// Retrieve sample by hash, optionally in the `CaRT` format
716    ///
717    /// # Errors
718    ///
719    /// This may return an error if there's a network situation or if the user is not logged in
720    /// or not properly authorized to connect.
721    pub fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
722        let api_endpoint = if cart {
723            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART_URL)
724        } else {
725            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_URL)
726        };
727
728        let res = self
729            .client
730            .get(format!("{}{api_endpoint}", self.url))
731            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
732            .send()?;
733
734        if !res.status().is_success() {
735            bail!("Received code {}", res.status());
736        }
737
738        let content_digest = res.headers().get("content-digest").map(ToOwned::to_owned);
739        let body = res.bytes()?;
740        let bytes = body.to_vec();
741
742        // TODO: Make this required in v0.3
743        if let Some(digest) = content_digest {
744            let hash = HashType::from_content_digest_header(digest.to_str()?)?;
745            if hash.verify(&bytes) {
746                trace!("Hash verified for sample {hash}");
747            } else {
748                error!("Hash mismatch for sample {hash}");
749            }
750        } else {
751            warn!("No content digest header received for sample {hash}");
752        }
753
754        Ok(bytes)
755    }
756
757    /// Fetch a report for a sample
758    ///
759    /// # Errors
760    ///
761    /// This may return an error if there's a network situation or if the user is not logged in
762    /// or not properly authorized to connect.
763    pub fn report(&self, hash: &str) -> Result<Report> {
764        let response = self
765            .client
766            .get(format!("{}{}/{hash}", self.url, malwaredb_api::SAMPLE_REPORT_URL))
767            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
768            .send()?
769            .json::<ServerResponse<Report>>()
770            .context(MDB_CLIENT_ERROR_CONTEXT)?;
771
772        match response {
773            ServerResponse::Success(report) => Ok(report),
774            ServerResponse::Error(e) => Err(e.into()),
775        }
776    }
777
778    /// Find similar samples in `MalwareDB` based on the contents of a given file.
779    /// This does not submit the sample to `MalwareDB`.
780    ///
781    /// # Errors
782    ///
783    /// This may return an error if there's a network situation or if the user is not logged in
784    /// or not properly authorized to connect.
785    pub fn similar(&self, contents: &[u8]) -> Result<SimilarSamplesResponse> {
786        let mut hashes = vec![];
787        let ssdeep_hash = FuzzyHash::new(contents);
788
789        let lzjd_str = LzDigest::from(contents).to_string();
790        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));
791        hashes.push((malwaredb_api::SimilarityHashType::SSDeep, ssdeep_hash.to_string()));
792
793        let mut builder = TlshBuilder::new(
794            tlsh_fixed::BucketKind::Bucket256,
795            tlsh_fixed::ChecksumKind::ThreeByte,
796            tlsh_fixed::Version::Version4,
797        );
798
799        builder.update(contents);
800        if let Ok(hasher) = builder.build() {
801            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
802        }
803
804        if let Ok(exe) = EXE::from(contents)
805            && let Some(imports) = exe.imports
806        {
807            hashes
808                .push((malwaredb_api::SimilarityHashType::ImportHash, hex::encode(imports.hash())));
809            hashes.push((malwaredb_api::SimilarityHashType::FuzzyImportHash, imports.fuzzy_hash()));
810        }
811
812        let request = malwaredb_api::SimilarSamplesRequest { hashes };
813
814        let response = self
815            .client
816            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES_URL))
817            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
818            .json(&request)
819            .send()?
820            .json::<ServerResponse<SimilarSamplesResponse>>()
821            .context(MDB_CLIENT_ERROR_CONTEXT)?;
822
823        match response {
824            ServerResponse::Success(similar) => Ok(similar),
825            ServerResponse::Error(e) => Err(e.into()),
826        }
827    }
828
829    /// Submit a Yara rule and return the UUID of the search for later retrieval.
830    ///
831    /// # Errors
832    ///
833    /// Network or authentication errors
834    pub fn yara_search(&self, yara: &str) -> Result<YaraSearchRequestResponse> {
835        let yara = YaraSearchRequest {
836            rules: vec![yara.to_string()],
837            response: PartialHashSearchType::SHA256,
838        };
839
840        let response = self
841            .client
842            .post(format!("{}{}", self.url, malwaredb_api::YARA_SEARCH_URL))
843            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
844            .json(&yara)
845            .send()?
846            .json::<ServerResponse<YaraSearchRequestResponse>>()?;
847
848        match response {
849            ServerResponse::Success(similar) => Ok(similar),
850            ServerResponse::Error(e) => Err(e.into()),
851        }
852    }
853
854    /// Get the result from a Yara search
855    ///
856    /// # Errors
857    ///
858    /// Network or authentication errors
859    pub fn yara_result(&self, uuid: Uuid) -> Result<YaraSearchResponse> {
860        let response = self
861            .client
862            .get(format!("{}{}/{uuid}", self.url, malwaredb_api::YARA_SEARCH_URL))
863            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
864            .send()?
865            .json::<ServerResponse<YaraSearchResponse>>()?;
866
867        match response {
868            ServerResponse::Success(sources) => Ok(sources),
869            ServerResponse::Error(e) => Err(e.into()),
870        }
871    }
872}
873
874impl Debug for MdbClient {
875    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
876        use crate::MDB_VERSION;
877
878        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
879    }
880}
881
882/// Wrapper around search results for iterating over resulting hashes with the blocking client
883///
884/// ```rust,no_run
885/// use malwaredb_client::blocking::{MdbClient, IterableHashSearchResult};
886/// use malwaredb_client::malwaredb_api::{PartialHashSearchType, SearchType};
887///
888/// let client = MdbClient::load().expect("Failed to load client or parse config file");
889///
890/// // Get the first 100 files where the file name contains "foo", returning hashes as SHA-256
891/// let search_result = client.partial_search(None, Some("foo".into()), PartialHashSearchType::SHA256, 100).unwrap();
892/// for hash in IterableHashSearchResult::from(search_result, &client) {
893///     println!("{hash}");
894/// }
895/// ```
896#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
897pub struct IterableHashSearchResult<'a> {
898    /// Server search result
899    pub response: SearchResponse,
900
901    /// Blocking client
902    client: &'a MdbClient,
903}
904
905impl<'a> IterableHashSearchResult<'a> {
906    /// Iterate over the hashes from a search result and blocking client
907    #[must_use]
908    pub fn from(response: SearchResponse, client: &'a MdbClient) -> Self {
909        Self { response, client }
910    }
911}
912
913impl Iterator for IterableHashSearchResult<'_> {
914    type Item = String;
915
916    fn next(&mut self) -> Option<Self::Item> {
917        if let Some(hash) = self.response.hashes.pop() {
918            Some(hash)
919        } else if let Some(uuid) = self.response.pagination {
920            let request = SearchRequest {
921                search: SearchType::Continuation(uuid),
922            };
923
924            self.response = match self.client.do_search_request(&request) {
925                Ok(response) => response,
926                Err(e) => {
927                    warn!("Failed to continue search: {e}");
928                    return None;
929                }
930            };
931
932            self.response.hashes.pop()
933        } else {
934            None
935        }
936    }
937}
938
939/// Wrapper around search results for iterating over resulting binaries with the blocking client
940#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
941pub struct IterableSampleSearchResult<'a> {
942    /// Server search result
943    pub response: SearchResponse,
944
945    /// Blocking client
946    client: &'a MdbClient,
947}
948
949impl<'a> IterableSampleSearchResult<'a> {
950    /// Iterate over the hashes from a search result and blocking client, returning the binary
951    #[must_use]
952    pub fn from(response: SearchResponse, client: &'a MdbClient) -> Self {
953        Self { response, client }
954    }
955}
956
957impl Iterator for IterableSampleSearchResult<'_> {
958    type Item = Vec<u8>;
959
960    fn next(&mut self) -> Option<Self::Item> {
961        if let Some(hash) = self.response.hashes.pop() {
962            let binary = match self.client.retrieve(&hash, false) {
963                Ok(binary) => binary,
964                Err(e) => {
965                    error!("Failed to download {hash}: {e}");
966                    return None;
967                }
968            };
969            Some(binary)
970        } else if let Some(uuid) = self.response.pagination {
971            let request = SearchRequest {
972                search: SearchType::Continuation(uuid),
973            };
974
975            self.response = match self.client.do_search_request(&request) {
976                Ok(response) => response,
977                Err(e) => {
978                    warn!("Failed to continue search: {e}");
979                    return None;
980                }
981            };
982
983            if let Some(hash) = self.response.hashes.pop() {
984                let binary = match self.client.retrieve(&hash, false) {
985                    Ok(binary) => binary,
986                    Err(e) => {
987                        error!("Failed to download {hash}: {e}");
988                        return None;
989                    }
990                };
991                Some(binary)
992            } else {
993                None
994            }
995        } else {
996            None
997        }
998    }
999}
malwaredb_client/blocking.rs

malwaredb_client/
blocking.rs
