Skip to main content

malwaredb_client/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2
3#![doc = include_str!("../README.md")]
4#![cfg_attr(docsrs, feature(doc_cfg))]
5#![deny(missing_docs)]
6#![deny(clippy::all)]
7#![deny(clippy::pedantic)]
8#![forbid(unsafe_code)]
9
10/// Non-async version of the Malware DB client
11#[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
12#[cfg(feature = "blocking")]
13pub mod blocking;
14
15pub use malwaredb_api;
16use malwaredb_api::{
17    GetAPIKeyResponse, GetUserInfoResponse, Labels, PartialHashSearchType, Report, SearchRequest,
18    SearchRequestParameters, SearchResponse, SearchType, ServerInfo, ServerResponse,
19    SimilarSamplesResponse, Sources, SupportedFileTypes, YaraSearchRequest,
20    YaraSearchRequestResponse, YaraSearchResponse, digest::HashType,
21};
22use malwaredb_types::exec::pe32::EXE;
23use malwaredb_types::utils::entropy_calc;
24
25use std::collections::HashSet;
26use std::fmt::{Debug, Display, Formatter};
27use std::fs::OpenOptions;
28use std::io::{Cursor, Write};
29use std::path::{Path, PathBuf};
30use std::sync::LazyLock;
31
32use anyhow::{Context, Result, bail, ensure};
33use base64::Engine;
34use base64::engine::general_purpose;
35use cart_container::JsonMap;
36use fuzzyhash::FuzzyHash;
37use home::home_dir;
38use malwaredb_lzjd2::lzjd::LzDigest;
39use mdns_sd::{ServiceDaemon, ServiceEvent};
40use reqwest::Certificate;
41use serde::{Deserialize, Serialize};
42use sha2::{Digest, Sha256, Sha384, Sha512};
43use tlsh_fixed::TlshBuilder;
44use tracing::{debug, error, info, trace, warn};
45use uuid::Uuid;
46use zeroize::{Zeroize, ZeroizeOnDrop};
47
48/// Local directory for the Malware DB client configs
49const MDB_CLIENT_DIR: &str = "malwaredb_client";
50
51/// Error for Anyhow's `context()` function with regard to what is expected just a network error
52pub(crate) const MDB_CLIENT_ERROR_CONTEXT: &str =
53    "Network error connecting to MalwareDB, or failure to decode server response.";
54
55/// Config file name expected by Malware DB client
56const MDB_CLIENT_CONFIG_TOML: &str = "mdb_client.toml";
57
58/// MDB version
59pub const MDB_VERSION: &str = env!("CARGO_PKG_VERSION");
60
61/// MDB version as a semantic version object
62pub static MDB_VERSION_SEMVER: LazyLock<semver::Version> =
63    LazyLock::new(|| semver::Version::parse(MDB_VERSION).unwrap());
64
65/// macOS Keychain functionality
66#[cfg(target_os = "macos")]
67pub(crate) mod macos {
68    use crate::CertificateType;
69
70    use anyhow::Result;
71    use reqwest::Certificate;
72    use security_framework::os::macos::keychain::SecKeychain;
73    use tracing::error;
74
75    /// Application identifier for macOS Keychain
76    const KEYCHAIN_ID: &str = "malwaredb-client";
77
78    /// Entry ID for the Malware DB server URL
79    const KEYCHAIN_URL: &str = "URL";
80
81    /// Entry ID for the user's Malware DB API Key
82    const KEYCHAIN_API_KEY: &str = "API_KEY";
83
84    /// Entry ID for the custom PEM-encoded certificate for the Malware DB server
85    const KEYCHAIN_CERTIFICATE_PEM: &str = "CERT_PEM";
86
87    /// Entry ID for the custom DER-encoded certificate for the Malware DB server
88    const KEYCHAIN_CERTIFICATE_DER: &str = "CERT_DER";
89
90    #[derive(Clone)]
91    pub(crate) struct CertificateData {
92        pub cert_type: CertificateType,
93        pub cert_bytes: Vec<u8>,
94    }
95
96    impl CertificateData {
97        pub(crate) fn as_cert(&self) -> Result<Certificate> {
98            Ok(match self.cert_type {
99                CertificateType::PEM => Certificate::from_pem(&self.cert_bytes)?,
100                CertificateType::DER => Certificate::from_der(&self.cert_bytes)?,
101            })
102        }
103    }
104
105    /// Save the elements to they Keychain
106    pub fn save_credentials(url: &str, key: &str, cert: Option<CertificateData>) -> Result<()> {
107        let keychain = SecKeychain::default()?;
108
109        keychain.add_generic_password(KEYCHAIN_ID, KEYCHAIN_URL, url.as_bytes())?;
110        keychain.add_generic_password(KEYCHAIN_ID, KEYCHAIN_API_KEY, key.as_bytes())?;
111
112        if let Some(cert) = cert {
113            match cert.cert_type {
114                CertificateType::PEM => keychain.add_generic_password(
115                    KEYCHAIN_ID,
116                    KEYCHAIN_CERTIFICATE_PEM,
117                    &cert.cert_bytes,
118                )?,
119                CertificateType::DER => keychain.add_generic_password(
120                    KEYCHAIN_ID,
121                    KEYCHAIN_CERTIFICATE_DER,
122                    &cert.cert_bytes,
123                )?,
124            }
125        }
126
127        Ok(())
128    }
129
130    /// Return key, url, and optionally the certificate in that order. Errors silently discarded.
131    pub fn retrieve_credentials() -> Result<(String, String, Option<CertificateData>)> {
132        let keychain = SecKeychain::default()?;
133        let (api_key, _item) = keychain.find_generic_password(KEYCHAIN_ID, KEYCHAIN_API_KEY)?;
134        let api_key = String::from_utf8(api_key.as_ref().to_vec())?;
135        let (url, _item) = keychain.find_generic_password(KEYCHAIN_ID, KEYCHAIN_URL)?;
136        let url = String::from_utf8(url.as_ref().to_vec())?;
137
138        if let Ok((cert, _item)) =
139            keychain.find_generic_password(KEYCHAIN_ID, KEYCHAIN_CERTIFICATE_PEM)
140        {
141            let cert = CertificateData {
142                cert_type: CertificateType::PEM,
143                cert_bytes: cert.to_vec(),
144            };
145            return Ok((api_key, url, Some(cert)));
146        }
147
148        if let Ok((cert, _item)) =
149            keychain.find_generic_password(KEYCHAIN_ID, KEYCHAIN_CERTIFICATE_DER)
150        {
151            let cert = CertificateData {
152                cert_type: CertificateType::DER,
153                cert_bytes: cert.to_vec(),
154            };
155            return Ok((api_key, url, Some(cert)));
156        }
157
158        Ok((api_key, url, None))
159    }
160
161    /// Delete Malware DB client information from the Keychain
162    pub fn clear_credentials() {
163        if let Ok(keychain) = SecKeychain::default() {
164            for element in [
165                KEYCHAIN_API_KEY,
166                KEYCHAIN_URL,
167                KEYCHAIN_CERTIFICATE_PEM,
168                KEYCHAIN_CERTIFICATE_DER,
169            ] {
170                if let Ok((_, item)) = keychain.find_generic_password(KEYCHAIN_ID, element) {
171                    item.delete();
172                }
173            }
174        } else {
175            error!("Failed to get access to the Keychain to clear credentials");
176        }
177    }
178}
179
180#[allow(clippy::upper_case_acronyms)]
181#[derive(Copy, Clone, PartialEq, Eq)]
182enum CertificateType {
183    DER,
184    PEM,
185}
186
187/// Asynchronous Malware DB Client Configuration and connection
188#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
189pub struct MdbClient {
190    /// URL of the Malware DB server, including http and port number, ending without a slash
191    pub url: String,
192
193    /// User's API key for Malware DB
194    /// "anonymous" for anonymous access
195    api_key: String,
196
197    /// Async http client which stores the optional server certificate
198    #[zeroize(skip)]
199    #[serde(skip)]
200    client: reqwest::Client,
201
202    /// Server's certificate
203    #[cfg(target_os = "macos")]
204    #[zeroize(skip)]
205    #[serde(skip)]
206    cert: Option<macos::CertificateData>,
207}
208
209impl MdbClient {
210    /// MDB Client from components, doesn't test connectivity
211    ///
212    /// # Errors
213    ///
214    /// Returns an error if a list of certificates was passed and any were not in the expected
215    /// DER or PEM format or could not be parsed.
216    pub fn new(url: String, api_key: String, cert_path: Option<PathBuf>) -> Result<Self> {
217        let mut url = url;
218        let url = if url.ends_with('/') {
219            url.pop();
220            url
221        } else {
222            url
223        };
224
225        let cert = if let Some(path) = cert_path {
226            Some((path_load_cert(&path)?, path))
227        } else {
228            None
229        };
230
231        let builder = reqwest::ClientBuilder::new()
232            .gzip(true)
233            .zstd(true)
234            .use_rustls_tls()
235            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
236
237        let client = if let Some(((_cert_type, cert), _cert_path)) = &cert {
238            builder.add_root_certificate(cert.clone()).build()
239        } else {
240            builder.build()
241        }?;
242
243        #[cfg(target_os = "macos")]
244        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
245            Some(macos::CertificateData {
246                cert_type: *cert_type,
247                cert_bytes: std::fs::read(cert_path)?,
248            })
249        } else {
250            None
251        };
252
253        Ok(Self {
254            url,
255            api_key,
256            client,
257
258            #[cfg(target_os = "macos")]
259            cert,
260        })
261    }
262
263    /// Connect to a server anonymously.
264    ///
265    /// # Errors
266    ///
267    /// Networking errors may result, and there will be an error if the server doesn't support anonymous connections.
268    pub async fn anonymous(url: String, save: bool, cert_path: Option<PathBuf>) -> Result<Self> {
269        let mut url = url;
270        let url = if url.ends_with('/') {
271            url.pop();
272            url
273        } else {
274            url
275        };
276
277        let builder = reqwest::ClientBuilder::new()
278            .gzip(true)
279            .zstd(true)
280            .use_rustls_tls()
281            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
282
283        let cert = if let Some(path) = cert_path {
284            Some((path_load_cert(&path)?, path))
285        } else {
286            None
287        };
288
289        let client = if let Some(((_cert_type, cert), _cert_path)) = &cert {
290            builder.add_root_certificate(cert.clone()).build()
291        } else {
292            builder.build()
293        }?;
294
295        let info = client
296            .get(format!("{url}{}", malwaredb_api::USER_INFO_URL))
297            .send()
298            .await?
299            .json::<ServerResponse<GetUserInfoResponse>>()
300            .await
301            .context(MDB_CLIENT_ERROR_CONTEXT)?;
302
303        if let ServerResponse::Success(info) = info {
304            ensure!(info.id > 0);
305        } else {
306            bail!("Anonymous access is not allowed");
307        }
308
309        #[cfg(target_os = "macos")]
310        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
311            Some(macos::CertificateData {
312                cert_type: *cert_type,
313                cert_bytes: std::fs::read(cert_path)?,
314            })
315        } else {
316            None
317        };
318
319        let client = Self {
320            url,
321            api_key: "anonymous".to_string(),
322            client,
323
324            #[cfg(target_os = "macos")]
325            cert,
326        };
327
328        let server_info = client.server_info().await?;
329        if server_info.mdb_version > *MDB_VERSION_SEMVER {
330            warn!(
331                "Server version {:?} is newer than client {:?}, consider updating.",
332                server_info.mdb_version, MDB_VERSION_SEMVER
333            );
334        }
335
336        if save && let Err(e) = client.save() {
337            error!("Anonymous connection successful but failed to save config: {e}");
338            bail!("Anonymous connection successful but failed to save config: {e}");
339        }
340
341        Ok(client)
342    }
343
344    /// Login to a server, optionally save the configuration file, and return a client object
345    ///
346    /// # Errors
347    ///
348    /// Returns an error if the server URL, username, or password were incorrect, or if a network
349    /// issue occurred.
350    pub async fn login(
351        url: String,
352        username: String,
353        password: String,
354        save: bool,
355        cert_path: Option<PathBuf>,
356    ) -> Result<Self> {
357        let mut url = url;
358        let url = if url.ends_with('/') {
359            url.pop();
360            url
361        } else {
362            url
363        };
364
365        let api_request = malwaredb_api::GetAPIKeyRequest {
366            user: username,
367            password,
368        };
369
370        let builder = reqwest::ClientBuilder::new()
371            .gzip(true)
372            .zstd(true)
373            .use_rustls_tls()
374            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
375
376        let cert = if let Some(path) = cert_path {
377            Some((path_load_cert(&path)?, path))
378        } else {
379            None
380        };
381
382        let client = if let Some(((_cert_type, cert), _cert_path)) = &cert {
383            builder.add_root_certificate(cert.clone()).build()
384        } else {
385            builder.build()
386        }?;
387
388        let res = client
389            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
390            .json(&api_request)
391            .send()
392            .await?
393            .json::<ServerResponse<GetAPIKeyResponse>>()
394            .await
395            .context(MDB_CLIENT_ERROR_CONTEXT)?;
396
397        let res = match res {
398            ServerResponse::Success(res) => res,
399            ServerResponse::Error(err) => return Err(err.into()),
400        };
401
402        #[cfg(target_os = "macos")]
403        let cert = if let Some(((cert_type, _cert), cert_path)) = &cert {
404            Some(macos::CertificateData {
405                cert_type: *cert_type,
406                cert_bytes: std::fs::read(cert_path)?,
407            })
408        } else {
409            None
410        };
411
412        let client = MdbClient {
413            url,
414            api_key: res.key.clone(),
415            client,
416
417            #[cfg(target_os = "macos")]
418            cert,
419        };
420
421        let server_info = client.server_info().await?;
422        if server_info.mdb_version > *MDB_VERSION_SEMVER {
423            warn!(
424                "Server version {:?} is newer than client {:?}, consider updating.",
425                server_info.mdb_version, MDB_VERSION_SEMVER
426            );
427        }
428
429        if save && let Err(e) = client.save() {
430            error!("Login successful but failed to save config: {e}");
431            bail!("Login successful but failed to save config: {e}");
432        }
433        Ok(client)
434    }
435
436    /// Reset one's own API key to effectively logout & disable all clients who are using the key
437    ///
438    /// # Errors
439    ///
440    /// Returns an error if there was a network issue or the user wasn't properly logged in.
441    pub async fn reset_key(&self) -> Result<()> {
442        let response = self
443            .client
444            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
445            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
446            .send()
447            .await
448            .context(MDB_CLIENT_ERROR_CONTEXT)?;
449        if !response.status().is_success() {
450            bail!("failed to reset API key, was it correct?");
451        }
452        Ok(())
453    }
454
455    /// Malware DB Client configuration loaded from a specified path
456    ///
457    /// # Errors
458    ///
459    /// Returns an error if the configuration file cannot be read, possibly because it
460    /// doesn't exist or due to a permission error or a parsing error.
461    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
462        let name = path.as_ref().display();
463        let config =
464            std::fs::read_to_string(&path).context(format!("failed to read config file {name}"))?;
465        let cfg: MdbClient =
466            toml::from_str(&config).context(format!("failed to parse config file {name}"))?;
467        Ok(cfg)
468    }
469
470    /// Malware DB Client configuration from user's home directory
471    ///
472    /// On macOS, it will attempt to load this information in the Keychain, which isn't required.
473    ///
474    /// # Errors
475    ///
476    /// Returns an error if the configuration file cannot be read, possibly because it
477    /// doesn't exist or due to a permission error or a parsing error.
478    pub fn load() -> Result<Self> {
479        #[cfg(target_os = "macos")]
480        {
481            if let Ok((api_key, url, cert)) = macos::retrieve_credentials() {
482                let builder = reqwest::ClientBuilder::new()
483                    .gzip(true)
484                    .zstd(true)
485                    .use_rustls_tls()
486                    .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
487
488                let client = if let Some(cert) = &cert {
489                    builder.add_root_certificate(cert.as_cert()?).build()
490                } else {
491                    builder.build()
492                }?;
493
494                return Ok(Self {
495                    url,
496                    api_key,
497                    client,
498                    cert,
499                });
500            }
501        }
502
503        let path = get_config_path(false)?;
504        if path.exists() {
505            return Self::from_file(path);
506        }
507        bail!("config file not found")
508    }
509
510    /// Save Malware DB Client configuration to the user's home directory.
511    ///
512    /// On macOS, it will attempt to save this information in the Keychain, which isn't required.
513    ///
514    /// # Errors
515    ///
516    /// Returns an error if there was a problem saving the configuration file.
517    pub fn save(&self) -> Result<()> {
518        #[cfg(target_os = "macos")]
519        {
520            if macos::save_credentials(&self.url, &self.api_key, self.cert.clone()).is_ok() {
521                return Ok(());
522            }
523        }
524
525        let toml = toml::to_string(self)?;
526        let path = get_config_path(true)?;
527
528        let mut options = OpenOptions::new();
529        options
530            .write(true)
531            .create(true)
532            .append(false)
533            .truncate(false);
534
535        #[cfg(target_family = "unix")]
536        {
537            use std::os::unix::fs::OpenOptionsExt;
538
539            options.mode(0o600);
540        }
541
542        let mut file = options.open(&path)?;
543        write!(file, "{toml}").context(format!("failed to write mdb config to {}", path.display()))
544    }
545
546    /// Delete the Malware DB client configuration file
547    ///
548    /// # Errors
549    ///
550    /// Returns an error if there isn't a configuration file to delete, or if it cannot be deleted,
551    /// possibly due to a permissions error.
552    pub fn delete(&self) -> Result<()> {
553        #[cfg(target_os = "macos")]
554        macos::clear_credentials();
555
556        let path = get_config_path(false)?;
557        if path.exists() {
558            std::fs::remove_file(&path)
559                .context(format!("failed to delete client config file {}", path.display()))?;
560        }
561        Ok(())
562    }
563
564    // Actions of the client
565
566    /// Get information about the server, unauthenticated
567    ///
568    /// # Errors
569    ///
570    /// This may return an error if there's a network situation.
571    pub async fn server_info(&self) -> Result<ServerInfo> {
572        let response = self
573            .client
574            .get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO_URL))
575            .send()
576            .await?
577            .json::<ServerResponse<ServerInfo>>()
578            .await
579            .context(MDB_CLIENT_ERROR_CONTEXT)?;
580
581        match response {
582            ServerResponse::Success(info) => Ok(info),
583            ServerResponse::Error(e) => Err(e.into()),
584        }
585    }
586
587    /// Get file types supported by the server, unauthenticated
588    ///
589    /// # Errors
590    ///
591    /// This may return an error if there's a network situation.
592    pub async fn supported_types(&self) -> Result<SupportedFileTypes> {
593        let response = self
594            .client
595            .get(format!("{}{}", self.url, malwaredb_api::SUPPORTED_FILE_TYPES_URL))
596            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
597            .send()
598            .await?
599            .json::<ServerResponse<SupportedFileTypes>>()
600            .await
601            .context(MDB_CLIENT_ERROR_CONTEXT)?;
602
603        match response {
604            ServerResponse::Success(types) => Ok(types),
605            ServerResponse::Error(e) => Err(e.into()),
606        }
607    }
608
609    /// Get information about the user
610    ///
611    /// # Errors
612    ///
613    /// This may return an error if there's a network situation or if the user is not logged in
614    /// or not properly authorized to connect.
615    pub async fn whoami(&self) -> Result<GetUserInfoResponse> {
616        let response = self
617            .client
618            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
619            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
620            .send()
621            .await?
622            .json::<ServerResponse<GetUserInfoResponse>>()
623            .await
624            .context(MDB_CLIENT_ERROR_CONTEXT)?;
625
626        match response {
627            ServerResponse::Success(info) => Ok(info),
628            ServerResponse::Error(e) => Err(e.into()),
629        }
630    }
631
632    /// Get the sample labels known to the server
633    ///
634    /// # Errors
635    ///
636    /// This may return an error if there's a network situation or if the user is not logged in
637    /// or not properly authorized to connect.
638    pub async fn labels(&self) -> Result<Labels> {
639        let response = self
640            .client
641            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS_URL))
642            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
643            .send()
644            .await?
645            .json::<ServerResponse<Labels>>()
646            .await
647            .context(MDB_CLIENT_ERROR_CONTEXT)?;
648
649        match response {
650            ServerResponse::Success(labels) => Ok(labels),
651            ServerResponse::Error(e) => Err(e.into()),
652        }
653    }
654
655    /// Get the sources available to the current user
656    ///
657    /// # Errors
658    ///
659    /// This may return an error if there's a network situation or if the user is not logged in
660    /// or not properly authorized to connect.
661    pub async fn sources(&self) -> Result<Sources> {
662        let response = self
663            .client
664            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES_URL))
665            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
666            .send()
667            .await?
668            .json::<ServerResponse<Sources>>()
669            .await
670            .context(MDB_CLIENT_ERROR_CONTEXT)?;
671
672        match response {
673            ServerResponse::Success(sources) => Ok(sources),
674            ServerResponse::Error(e) => Err(e.into()),
675        }
676    }
677
678    /// Submit one file to `MalwareDB`: provide the contents, file name, and source ID
679    ///
680    /// # Errors
681    ///
682    /// This may return an error if there's a network situation or if the user is not logged in
683    /// or not properly authorized to connect.
684    pub async fn submit(
685        &self,
686        contents: impl AsRef<[u8]>,
687        file_name: impl AsRef<str>,
688        source_id: u32,
689    ) -> Result<bool> {
690        let mut hasher = Sha256::new();
691        hasher.update(&contents);
692        let result = hasher.finalize();
693
694        let encoded = general_purpose::STANDARD.encode(contents);
695
696        let payload = malwaredb_api::NewSampleB64 {
697            file_name: file_name.as_ref().to_string(),
698            source_id,
699            file_contents_b64: encoded,
700            sha256: hex::encode(result),
701        };
702
703        match self
704            .client
705            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE_JSON_URL))
706            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
707            .json(&payload)
708            .send()
709            .await
710        {
711            Ok(res) => {
712                if !res.status().is_success() {
713                    info!("Code {} sending {}", res.status(), payload.file_name);
714                }
715                Ok(res.status().is_success())
716            }
717            Err(e) => {
718                let status: String = e
719                    .status()
720                    .map(|s| s.as_str().to_string())
721                    .unwrap_or_default();
722                error!("Error{status} sending {}: {e}", payload.file_name);
723                bail!(e.to_string())
724            }
725        }
726    }
727
728    /// Submit one file to `MalwareDB` as a Cbor object: provide the contents, file name, and source ID
729    /// Experimental! May be removed at any point.
730    ///
731    /// # Errors
732    ///
733    /// This may return an error if there's a network situation or if the user is not logged in
734    /// or not properly authorized to connect.
735    pub async fn submit_as_cbor(
736        &self,
737        contents: impl AsRef<[u8]>,
738        file_name: impl AsRef<str>,
739        source_id: u32,
740    ) -> Result<bool> {
741        let mut hasher = Sha256::new();
742        hasher.update(&contents);
743        let result = hasher.finalize();
744
745        let payload = malwaredb_api::NewSampleBytes {
746            file_name: file_name.as_ref().to_string(),
747            source_id,
748            file_contents: contents.as_ref().to_vec(),
749            sha256: hex::encode(result),
750        };
751
752        let mut bytes = Vec::with_capacity(payload.file_contents.len());
753        ciborium::ser::into_writer(&payload, &mut bytes)?;
754
755        match self
756            .client
757            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE_CBOR_URL))
758            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
759            .header("content-type", "application/cbor")
760            .body(bytes)
761            .send()
762            .await
763        {
764            Ok(res) => {
765                if !res.status().is_success() {
766                    info!("Code {} sending {}", res.status(), payload.file_name);
767                }
768                Ok(res.status().is_success())
769            }
770            Err(e) => {
771                let status: String = e
772                    .status()
773                    .map(|s| s.as_str().to_string())
774                    .unwrap_or_default();
775                error!("Error{status} sending {}: {e}", payload.file_name);
776                bail!(e.to_string())
777            }
778        }
779    }
780
781    /// Search for a file based on partial hash and/or partial file name, returns a list of hashes
782    ///
783    /// # Errors
784    ///
785    /// * 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
786    pub async fn partial_search(
787        &self,
788        partial_hash: Option<(PartialHashSearchType, String)>,
789        name: Option<String>,
790        response: PartialHashSearchType,
791        limit: u32,
792    ) -> Result<SearchResponse> {
793        let query = SearchRequest {
794            search: SearchType::Search(SearchRequestParameters {
795                partial_hash,
796                file_name: name,
797                response,
798                limit,
799                labels: None,
800                file_type: None,
801                magic: None,
802            }),
803        };
804
805        self.do_search_request(&query).await
806    }
807
808    /// Search for a file based on partial hash and/or partial file name, labels, file type; returns a list of hashes
809    ///
810    /// # Errors
811    ///
812    /// * 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
813    #[allow(clippy::too_many_arguments)]
814    pub async fn partial_search_labels_type(
815        &self,
816        partial_hash: Option<(PartialHashSearchType, String)>,
817        name: Option<String>,
818        response: PartialHashSearchType,
819        labels: Option<Vec<String>>,
820        file_type: Option<String>,
821        magic: Option<String>,
822        limit: u32,
823    ) -> Result<SearchResponse> {
824        let query = SearchRequest {
825            search: SearchType::Search(SearchRequestParameters {
826                partial_hash,
827                file_name: name,
828                response,
829                limit,
830                file_type,
831                magic,
832                labels,
833            }),
834        };
835
836        self.do_search_request(&query).await
837    }
838
839    /// Return the next page from the search result
840    ///
841    /// # Errors
842    ///
843    /// Returns an error if there is a network problem, or pagination not available
844    pub async fn next_page_search(&self, response: &SearchResponse) -> Result<SearchResponse> {
845        if let Some(uuid) = response.pagination {
846            let request = SearchRequest {
847                search: SearchType::Continuation(uuid),
848            };
849            return self.do_search_request(&request).await;
850        }
851
852        bail!("Pagination not available")
853    }
854
855    async fn do_search_request(&self, query: &SearchRequest) -> Result<SearchResponse> {
856        ensure!(
857            query.is_valid(),
858            "Query isn't valid: hash isn't hexidecimal or both the hashes and file name are empty"
859        );
860
861        let response = self
862            .client
863            .post(format!("{}{}", self.url, malwaredb_api::SEARCH_URL))
864            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
865            .json(query)
866            .send()
867            .await?
868            .json::<ServerResponse<SearchResponse>>()
869            .await
870            .context(MDB_CLIENT_ERROR_CONTEXT)?;
871
872        match response {
873            ServerResponse::Success(search) => Ok(search),
874            ServerResponse::Error(e) => Err(e.into()),
875        }
876    }
877
878    /// Retrieve sample by hash, optionally in the `CaRT` format
879    ///
880    /// # Errors
881    ///
882    /// This may return an error if there's a network situation or if the user is not logged in
883    /// or not properly authorized to connect.
884    pub async fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
885        let api_endpoint = if cart {
886            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART_URL)
887        } else {
888            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_URL)
889        };
890
891        let res = self
892            .client
893            .get(format!("{}{api_endpoint}", self.url))
894            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
895            .send()
896            .await?;
897
898        if !res.status().is_success() {
899            bail!("Received code {}", res.status());
900        }
901
902        let content_digest = res.headers().get("content-digest").map(ToOwned::to_owned);
903        let body = res.bytes().await?;
904        let bytes = body.to_vec();
905
906        // TODO: Make this required in v0.3
907        if let Some(digest) = content_digest {
908            let hash = HashType::from_content_digest_header(digest.to_str()?)?;
909            if hash.verify(&bytes) {
910                trace!("Hash verified for sample {hash}");
911            } else {
912                error!("Hash mismatch for sample {hash}");
913            }
914        } else {
915            warn!("No content digest header received for sample {hash}");
916        }
917
918        Ok(bytes)
919    }
920
921    /// Fetch a report for a sample
922    ///
923    /// # Errors
924    ///
925    /// This may return an error if there's a network situation or if the user is not logged in
926    /// or not properly authorized to connect.
927    pub async fn report(&self, hash: &str) -> Result<Report> {
928        let response = self
929            .client
930            .get(format!("{}{}/{hash}", self.url, malwaredb_api::SAMPLE_REPORT_URL))
931            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
932            .send()
933            .await?
934            .json::<ServerResponse<Report>>()
935            .await
936            .context(MDB_CLIENT_ERROR_CONTEXT)?;
937
938        match response {
939            ServerResponse::Success(report) => Ok(report),
940            ServerResponse::Error(e) => Err(e.into()),
941        }
942    }
943
944    /// Find similar samples in `MalwareDB` based on the contents of a given file.
945    /// This does not submit the sample to `MalwareDB`.
946    ///
947    /// # Errors
948    ///
949    /// This may return an error if there's a network situation or if the user is not logged in
950    /// or not properly authorized to connect.
951    pub async fn similar(&self, contents: &[u8]) -> Result<SimilarSamplesResponse> {
952        let mut hashes = vec![];
953        let ssdeep_hash = FuzzyHash::new(contents);
954
955        let lzjd_str = LzDigest::from(contents).to_string();
956        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));
957        hashes.push((malwaredb_api::SimilarityHashType::SSDeep, ssdeep_hash.to_string()));
958
959        let mut builder = TlshBuilder::new(
960            tlsh_fixed::BucketKind::Bucket256,
961            tlsh_fixed::ChecksumKind::ThreeByte,
962            tlsh_fixed::Version::Version4,
963        );
964
965        builder.update(contents);
966        if let Ok(hasher) = builder.build() {
967            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
968        }
969
970        if let Ok(exe) = EXE::from(contents)
971            && let Some(imports) = exe.imports
972        {
973            hashes
974                .push((malwaredb_api::SimilarityHashType::ImportHash, hex::encode(imports.hash())));
975            hashes.push((malwaredb_api::SimilarityHashType::FuzzyImportHash, imports.fuzzy_hash()));
976        }
977
978        let request = malwaredb_api::SimilarSamplesRequest { hashes };
979
980        let response = self
981            .client
982            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES_URL))
983            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
984            .json(&request)
985            .send()
986            .await?
987            .json::<ServerResponse<SimilarSamplesResponse>>()
988            .await
989            .context(MDB_CLIENT_ERROR_CONTEXT)?;
990
991        match response {
992            ServerResponse::Success(similar) => Ok(similar),
993            ServerResponse::Error(e) => Err(e.into()),
994        }
995    }
996
997    /// Submit a Yara rule and return the UUID of the search for later retrieval.
998    ///
999    /// # Errors
1000    ///
1001    /// Network or authentication errors
1002    pub async fn yara_search(&self, yara: &str) -> Result<YaraSearchRequestResponse> {
1003        let yara = YaraSearchRequest {
1004            rules: vec![yara.to_string()],
1005            response: PartialHashSearchType::SHA256,
1006        };
1007
1008        let response = self
1009            .client
1010            .post(format!("{}{}", self.url, malwaredb_api::YARA_SEARCH_URL))
1011            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
1012            .json(&yara)
1013            .send()
1014            .await?
1015            .json::<ServerResponse<YaraSearchRequestResponse>>()
1016            .await?;
1017
1018        match response {
1019            ServerResponse::Success(sources) => Ok(sources),
1020            ServerResponse::Error(e) => Err(e.into()),
1021        }
1022    }
1023
1024    /// Get the result from a Yara search
1025    ///
1026    /// # Errors
1027    ///
1028    /// Network or authentication errors
1029    pub async fn yara_result(&self, uuid: Uuid) -> Result<YaraSearchResponse> {
1030        let response = self
1031            .client
1032            .get(format!("{}{}/{uuid}", self.url, malwaredb_api::YARA_SEARCH_URL))
1033            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
1034            .send()
1035            .await?
1036            .json::<ServerResponse<YaraSearchResponse>>()
1037            .await?;
1038
1039        match response {
1040            ServerResponse::Success(sources) => Ok(sources),
1041            ServerResponse::Error(e) => Err(e.into()),
1042        }
1043    }
1044}
1045
1046impl Debug for MdbClient {
1047    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1048        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
1049    }
1050}
1051
1052/// Convenience function for encoding bytes into a `CaRT` file using the default key. This also
1053/// adds SHA-384 and SHA-512 hashes plus the entropy of the original file.
1054/// See <https://github.com/CybercentreCanada/cart> for more information.
1055///
1056/// # Errors
1057///
1058/// There should not be any errors, but the underlying library can't guarantee that. However, any data
1059/// passed to it is correct.
1060pub fn encode_to_cart(data: &[u8]) -> Result<Vec<u8>> {
1061    let mut input_buffer = Cursor::new(data);
1062    let mut output_buffer = Cursor::new(vec![]);
1063    let mut output_metadata = JsonMap::new();
1064
1065    let mut sha384 = Sha384::new();
1066    sha384.update(data);
1067    let sha384 = hex::encode(sha384.finalize());
1068
1069    let mut sha512 = Sha512::new();
1070    sha512.update(data);
1071    let sha512 = hex::encode(sha512.finalize());
1072
1073    output_metadata.insert("sha384".into(), sha384.into());
1074    output_metadata.insert("sha512".into(), sha512.into());
1075    output_metadata.insert("entropy".into(), entropy_calc(data).into());
1076    cart_container::pack_stream(
1077        &mut input_buffer,
1078        &mut output_buffer,
1079        Some(output_metadata),
1080        None,
1081        cart_container::digesters::default_digesters(),
1082        None,
1083    )?;
1084
1085    Ok(output_buffer.into_inner())
1086}
1087
1088/// Convenience function for decoding a `CaRT` file using the default key, returning the bytes plus the
1089/// optional header and footer metadata, if present.
1090/// See <https://github.com/CybercentreCanada/cart> for more information.
1091///
1092/// # Errors
1093///
1094/// Returns an error if the file cannot be parsed or if this `CaRT` file didn't use the default key.
1095/// <https://github.com/CybercentreCanada/cart-rs/blob/7ad548143bb85b64f364804e90cfada6c31cf902/cart_container/src/cipher.rs#L14-L17>
1096pub fn decode_from_cart(data: &[u8]) -> Result<(Vec<u8>, Option<JsonMap>, Option<JsonMap>)> {
1097    let mut input_buffer = Cursor::new(data);
1098    let mut output_buffer = Cursor::new(vec![]);
1099    let (header, footer) =
1100        cart_container::unpack_stream(&mut input_buffer, &mut output_buffer, None)?;
1101    Ok((output_buffer.into_inner(), header, footer))
1102}
1103
1104/// Load a certificate from a path
1105///
1106/// # Errors
1107///
1108/// Returns errors if the file cannot be read or if the file isn't an ASN.1 DER file or
1109/// base64-encoded ASN.1 PEM file.
1110fn path_load_cert(path: &Path) -> Result<(CertificateType, Certificate)> {
1111    if !path.exists() {
1112        bail!("Certificate {} does not exist.", path.display());
1113    }
1114    let cert = match path
1115        .extension()
1116        .context("can't determine file extension")?
1117        .to_str()
1118        .context("unable to parse file extension")?
1119    {
1120        "pem" => {
1121            let contents = std::fs::read(path)?;
1122            (CertificateType::PEM, Certificate::from_pem(&contents)?)
1123        }
1124        "der" => {
1125            let contents = std::fs::read(path)?;
1126            (CertificateType::DER, Certificate::from_der(&contents)?)
1127        }
1128        ext => {
1129            bail!("Unknown extension {ext:?}")
1130        }
1131    };
1132    Ok(cert)
1133}
1134
1135/// Gets the configuration file in the following order:
1136///
1137/// 1. Current working directory: `./mdb_client.toml`
1138/// 2. Haiku-specific directory if on Haiku
1139/// 3. XDG Free Desktop directory if on Unix
1140/// 4. The user's home directory in `~/.config/malwaredb_client/mdb_client.toml`
1141/// 5. `mdb_client.toml` in the current directory (same as the first, but after checking others)
1142#[inline]
1143pub(crate) fn get_config_path(create: bool) -> Result<PathBuf> {
1144    // If there is a config file in the current working directory, use it
1145    let config = PathBuf::from(MDB_CLIENT_CONFIG_TOML);
1146    if config.exists() {
1147        return Ok(config);
1148    }
1149
1150    #[cfg(target_os = "haiku")]
1151    {
1152        let mut settings = PathBuf::from("/boot/home/config/settings/malwaredb");
1153        if create && !settings.exists() {
1154            std::fs::create_dir_all(&settings)?;
1155        }
1156        settings.push(MDB_CLIENT_CONFIG_TOML);
1157        return Ok(settings);
1158    }
1159
1160    #[cfg(unix)]
1161    {
1162        // Obey the Free Desktop standard, check the variable
1163        if let Some(xdg_home) = std::env::var_os("XDG_CONFIG_HOME") {
1164            let mut xdg_config_home = PathBuf::from(xdg_home);
1165            xdg_config_home.push(MDB_CLIENT_DIR);
1166            if create && !xdg_config_home.exists() {
1167                std::fs::create_dir_all(&xdg_config_home)?;
1168            }
1169            xdg_config_home.push(MDB_CLIENT_CONFIG_TOML);
1170            return Ok(xdg_config_home);
1171        }
1172    }
1173
1174    if let Some(mut home_config) = home_dir() {
1175        home_config.push(".config");
1176        home_config.push(MDB_CLIENT_DIR);
1177        if create && !home_config.exists() {
1178            std::fs::create_dir_all(&home_config)?;
1179        }
1180        home_config.push(MDB_CLIENT_CONFIG_TOML);
1181        return Ok(home_config);
1182    }
1183
1184    Ok(PathBuf::from(MDB_CLIENT_CONFIG_TOML))
1185}
1186
1187/// Malware DB entries found by Multicast DNS (also known as Bonjour or Zeroconf)
1188#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1189pub struct MalwareDBServer {
1190    /// Server IP or domain
1191    pub host: String,
1192
1193    /// Server port
1194    pub port: u16,
1195
1196    /// If the server expects an encrypted connection
1197    pub ssl: bool,
1198
1199    /// Malware DB server name
1200    pub name: String,
1201}
1202
1203impl Display for MalwareDBServer {
1204    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1205        if self.ssl {
1206            write!(f, "https://{}:{}", self.host, self.port)
1207        } else {
1208            write!(f, "http://{}:{}", self.host, self.port)
1209        }
1210    }
1211}
1212
1213impl MalwareDBServer {
1214    /// Retrieve details about the server
1215    ///
1216    /// # Errors
1217    ///
1218    /// An error will result if the server becomes unreachable or if a specific CA certificate is required
1219    pub async fn server_info(&self) -> Result<ServerInfo> {
1220        let client = reqwest::ClientBuilder::new()
1221            .gzip(true)
1222            .zstd(true)
1223            .use_rustls_tls()
1224            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")))
1225            .build()?;
1226
1227        let response = client
1228            .get(format!("{self}{}", malwaredb_api::SERVER_INFO_URL))
1229            .send()
1230            .await?
1231            .json::<ServerResponse<ServerInfo>>()
1232            .await
1233            .context(MDB_CLIENT_ERROR_CONTEXT)?;
1234
1235        match response {
1236            ServerResponse::Success(info) => Ok(info),
1237            ServerResponse::Error(e) => Err(e.into()),
1238        }
1239    }
1240
1241    /// Retrieve details about the server
1242    ///
1243    /// # Errors
1244    ///
1245    /// An error will result if the server becomes unreachable or if a specific CA certificate is required
1246    ///
1247    /// # Panics
1248    ///
1249    /// This method panics if called from within an async runtime.
1250    #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
1251    #[cfg(feature = "blocking")]
1252    pub fn server_info_blocking(&self) -> Result<ServerInfo> {
1253        let client = reqwest::blocking::ClientBuilder::new()
1254            .gzip(true)
1255            .zstd(true)
1256            .use_rustls_tls()
1257            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")))
1258            .build()?;
1259
1260        let response = client
1261            .get(format!("{self}{}", malwaredb_api::SERVER_INFO_URL))
1262            .send()?
1263            .json::<ServerResponse<ServerInfo>>()
1264            .context(MDB_CLIENT_ERROR_CONTEXT)?;
1265
1266        match response {
1267            ServerResponse::Success(similar) => Ok(similar),
1268            ServerResponse::Error(e) => Err(e.into()),
1269        }
1270    }
1271}
1272
1273/// Find servers using Multicast DNS (also known as Bonjour)
1274///
1275/// # Errors
1276///
1277/// This may fail if there's a networking issue.
1278pub fn discover_servers() -> Result<Vec<MalwareDBServer>> {
1279    const MAX_ITERS: usize = 5;
1280    let mdns = ServiceDaemon::new()?;
1281    let mut servers = HashSet::new();
1282    let receiver = mdns.browse(malwaredb_api::MDNS_NAME)?;
1283
1284    let mut counter = 0;
1285    while let Ok(event) = receiver.recv() {
1286        if let ServiceEvent::ServiceResolved(resolved) = event {
1287            let host = resolved.host.replace(".local.", "");
1288            let ssl = if let Some(ssl) = resolved.txt_properties.get("ssl") {
1289                ssl.val_str() == "true"
1290            } else {
1291                debug!(
1292                    "MalwareDB entry for {host}:{} doesn't specify ssl, assuming not",
1293                    resolved.port
1294                );
1295                false
1296            };
1297
1298            let server = MalwareDBServer {
1299                host,
1300                port: resolved.port,
1301                ssl,
1302                name: resolved.fullname.replace(malwaredb_api::MDNS_NAME, ""),
1303            };
1304
1305            servers.insert(server);
1306        }
1307        counter += 1;
1308        if counter > MAX_ITERS {
1309            break;
1310        }
1311    }
1312
1313    if mdns.shutdown().is_err() {
1314        // Pass
1315    }
1316    Ok(servers.into_iter().collect())
1317}
1318
1319#[cfg(test)]
1320mod tests {
1321    use super::*;
1322
1323    #[test]
1324    fn cart() {
1325        const BYTES: &[u8] = include_bytes!("../../crates/types/testdata/elf/elf_haiku_x86.cart");
1326        const ORIGINAL_SHA256: &str =
1327            "de10ba5e5402b46ea975b5cb8a45eb7df9e81dc81012fd4efd145ed2dce3a740";
1328
1329        let (decoded, header, footer) = decode_from_cart(BYTES).unwrap();
1330
1331        let mut sha256 = Sha256::new();
1332        sha256.update(&decoded);
1333        let sha256 = hex::encode(sha256.finalize());
1334        assert_eq!(sha256, ORIGINAL_SHA256);
1335
1336        let header = header.unwrap();
1337        let entropy = header.get("entropy").unwrap().as_f64().unwrap();
1338        assert!(entropy > 4.0 && entropy < 4.1);
1339
1340        let footer = footer.unwrap();
1341        assert_eq!(footer.get("length").unwrap(), "5093");
1342        assert_eq!(footer.get("sha256").unwrap(), ORIGINAL_SHA256);
1343    }
1344}