malwaredb_client/
blocking.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::fmt::{Debug, Formatter};
4use std::path::{Path, PathBuf};
5
6use crate::get_config_path;
7use malwaredb_types::exec::pe32::EXE;
8
9use anyhow::{bail, ensure, Context, Result};
10use base64::engine::general_purpose;
11use base64::Engine;
12use fuzzyhash::FuzzyHash;
13use malwaredb_api::{PartialHashSearchType, SearchRequest};
14use malwaredb_lzjd::{LZDict, Murmur3HashState};
15use serde::{Deserialize, Serialize};
16use sha2::{Digest, Sha256};
17use tlsh_fixed::TlshBuilder;
18use tracing::{error, info, warn};
19use zeroize::{Zeroize, ZeroizeOnDrop};
20
21/// Blocking Malware DB Client Configuration and connection
22#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
23pub struct MdbClient {
24    /// URL of the Malware DB server, including http and port number, ending without a slash
25    pub url: String,
26
27    /// User's API key for Malware DB
28    api_key: String,
29
30    /// Blocking http client which stores the optional server certificate
31    #[zeroize(skip)]
32    #[serde(skip)]
33    client: reqwest::blocking::Client,
34}
35
36impl MdbClient {
37    /// MDB Client from components, doesn't test connectivity
38    ///
39    /// # Errors
40    ///
41    /// Returns an error if a list of certificates was passed and any were not in the expected
42    /// DER or PEM format or could not be parsed.
43    pub fn new(url: String, api_key: String, cert_path: Option<PathBuf>) -> Result<Self> {
44        let mut url = url;
45        let url = if url.ends_with('/') {
46            url.pop();
47            url
48        } else {
49            url
50        };
51
52        let cert = if let Some(path) = cert_path {
53            Some((crate::path_load_cert(&path)?, path))
54        } else {
55            None
56        };
57
58        let builder = reqwest::blocking::ClientBuilder::new()
59            .gzip(true)
60            .zstd(true)
61            .use_rustls_tls()
62            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
63
64        let client = if let Some(cert) = cert {
65            builder.add_root_certificate(cert.0.clone()).build()
66        } else {
67            builder.build()
68        }?;
69
70        Ok(Self {
71            url,
72            api_key,
73            client,
74        })
75    }
76
77    /// Login to a server, optionally save the config file, and return a client object
78    ///
79    /// # Errors
80    ///
81    /// Returns an error if the server URL, username, or password were incorrect, or if a network
82    /// issue occurred.
83    pub fn login(
84        url: String,
85        username: String,
86        password: String,
87        save: bool,
88        cert_path: Option<PathBuf>,
89    ) -> Result<Self> {
90        let mut url = url;
91        let url = if url.ends_with('/') {
92            url.pop();
93            url
94        } else {
95            url
96        };
97
98        let api_request = malwaredb_api::GetAPIKeyRequest {
99            user: username,
100            password,
101        };
102
103        let builder = reqwest::blocking::ClientBuilder::new()
104            .gzip(true)
105            .zstd(true)
106            .use_rustls_tls()
107            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));
108
109        let cert = if let Some(path) = cert_path {
110            Some((crate::path_load_cert(&path)?, path))
111        } else {
112            None
113        };
114
115        let client = if let Some(cert) = &cert {
116            builder.add_root_certificate(cert.0.clone()).build()
117        } else {
118            builder.build()
119        }?;
120
121        let res = client
122            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
123            .json(&api_request)
124            .send()?
125            .json::<malwaredb_api::GetAPIKeyResponse>()?;
126
127        if let Some(key) = &res.key {
128            let client = MdbClient {
129                url,
130                api_key: key.clone(),
131                client,
132            };
133
134            let server_info = client.server_info()?;
135            if server_info.mdb_version > *crate::MDB_VERSION_SEMVER {
136                warn!(
137                    "Server version {:?} is newer than client {:?}, consider updating.",
138                    server_info.mdb_version,
139                    crate::MDB_VERSION_SEMVER
140                );
141            }
142
143            if save {
144                if let Err(e) = client.save() {
145                    error!("Login successful but failed to save config: {e}");
146                    bail!("Login successful but failed to save config: {e}");
147                }
148            }
149            Ok(client)
150        } else {
151            if let Some(msg) = &res.message {
152                error!("Login failed, response: {msg}");
153            }
154            bail!("server error or bad credentials");
155        }
156    }
157
158    /// Reset one's own API key to effectively logout & disable all clients who are using the key
159    ///
160    /// # Errors
161    ///
162    /// Returns an error if there was a network issue or the user wasn't properly logged in.
163    pub fn reset_key(&self) -> Result<()> {
164        let response = self
165            .client
166            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
167            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
168            .send()
169            .context("server error, or invalid API key")?;
170        if !response.status().is_success() {
171            bail!("failed to reset API key, was it correct?");
172        }
173        Ok(())
174    }
175
176    /// MDB Client loaded from a specified path
177    ///
178    /// # Errors
179    ///
180    /// Returns an error if the configuration file cannot be read, possibly because it
181    /// doesn't exist or due to a permission error or a parsing error.
182    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
183        let name = path.as_ref().display();
184        let config =
185            std::fs::read_to_string(&path).context(format!("failed to read config file {name}"))?;
186        let cfg: MdbClient =
187            toml::from_str(&config).context(format!("failed to parse config file {name}"))?;
188        Ok(cfg)
189    }
190
191    /// MDB Client from user's home directory
192    ///
193    /// # Errors
194    ///
195    /// Returns an error if the configuration file cannot be read, possibly because it
196    /// doesn't exist or due to a permission error or a parsing error.
197    pub fn load() -> Result<Self> {
198        let path = get_config_path(false)?;
199        if path.exists() {
200            return Self::from_file(path);
201        }
202        bail!("config file not found")
203    }
204
205    /// Save MDB Client to the user's home directory
206    ///
207    /// # Errors
208    ///
209    /// Returns an error if there was a problem saving the configuration file.
210    pub fn save(&self) -> Result<()> {
211        let toml = toml::to_string(self)?;
212        let path = get_config_path(true)?;
213        std::fs::write(&path, toml)
214            .context(format!("failed to write mdb config to {}", path.display()))
215    }
216
217    /// Delete the `MalwareDB` client config file
218    ///
219    /// # Errors
220    ///
221    /// Returns an error if there isn't a configuration file to delete, or if it cannot be deleted,
222    /// possibly due to a permissions error.
223    pub fn delete(&self) -> Result<()> {
224        let path = get_config_path(false)?;
225        if path.exists() {
226            std::fs::remove_file(&path).context(format!(
227                "failed to delete client config file {}",
228                path.display()
229            ))?;
230        }
231        Ok(())
232    }
233
234    // Actions of the client
235
236    /// Get information about the server, unauthenticated
237    ///
238    /// # Errors
239    ///
240    /// This may return an error if there's a network situation.
241    pub fn server_info(&self) -> Result<malwaredb_api::ServerInfo> {
242        self.client
243            .get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO))
244            .send()?
245            .json::<malwaredb_api::ServerInfo>()
246            .context("failed to receive or decode server info")
247    }
248
249    /// Get file types supported by the server, unauthenticated
250    ///
251    /// # Errors
252    ///
253    /// This may return an error if there's a network situation.
254    pub fn supported_types(&self) -> Result<malwaredb_api::SupportedFileTypes> {
255        self.client
256            .get(format!(
257                "{}{}",
258                self.url,
259                malwaredb_api::SUPPORTED_FILE_TYPES
260            ))
261            .send()?
262            .json::<malwaredb_api::SupportedFileTypes>()
263            .context("failed to receive or decode server-supported file types")
264    }
265
266    /// Get information about the user
267    ///
268    /// # Errors
269    ///
270    /// This may return an error if there's a network situation or if the user is not logged in
271    /// or not properly authorized to connect.
272    pub fn whoami(&self) -> Result<malwaredb_api::GetUserInfoResponse> {
273        self.client
274            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
275            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
276            .send()?
277            .json::<malwaredb_api::GetUserInfoResponse>()
278            .context("failed to receive or decode user info, or invalid API key")
279    }
280
281    /// Get the sample labels known to the server
282    ///
283    /// # Errors
284    ///
285    /// This may return an error if there's a network situation or if the user is not logged in
286    /// or not properly authorized to connect.
287    pub fn labels(&self) -> Result<malwaredb_api::Labels> {
288        self.client
289            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS))
290            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
291            .send()?
292            .json::<malwaredb_api::Labels>()
293            .context("failed to receive or decode available labels, or invalid API key")
294    }
295
296    /// Get the sources available to the current user
297    ///
298    /// # Errors
299    ///
300    /// This may return an error if there's a network situation or if the user is not logged in
301    /// or not properly authorized to connect.
302    pub fn sources(&self) -> Result<malwaredb_api::Sources> {
303        self.client
304            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES))
305            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
306            .send()?
307            .json::<malwaredb_api::Sources>()
308            .context("failed to receive or decode available labels, or invalid API key")
309    }
310
311    /// Submit one file to `MalwareDB`: provide the contents, file name, and source ID
312    ///
313    /// # Errors
314    ///
315    /// This may return an error if there's a network situation or if the user is not logged in
316    /// or not properly authorized to connect.
317    pub fn submit(
318        &self,
319        contents: impl AsRef<[u8]>,
320        file_name: String,
321        source_id: u32,
322    ) -> Result<bool> {
323        let mut hasher = Sha256::new();
324        hasher.update(&contents);
325        let result = hasher.finalize();
326
327        let encoded = general_purpose::STANDARD.encode(contents);
328
329        let payload = malwaredb_api::NewSample {
330            file_name,
331            source_id,
332            file_contents_b64: encoded,
333            sha256: hex::encode(result),
334        };
335
336        match self
337            .client
338            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE))
339            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
340            .json(&payload)
341            .send()
342        {
343            Ok(res) => {
344                if !res.status().is_success() {
345                    info!("Code {} sending {}", res.status(), payload.file_name);
346                }
347                Ok(res.status().is_success())
348            }
349            Err(e) => {
350                let status: String = e
351                    .status()
352                    .map(|s| s.as_str().to_string())
353                    .unwrap_or_default();
354                error!("Error{status} sending {}: {e}", payload.file_name);
355                bail!(e.to_string())
356            }
357        }
358    }
359
360    /// Search for a file based on partial hash and/or partial file name, returns a list of hashes
361    ///
362    /// # Errors
363    ///
364    /// * 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
365    pub fn partial_search(
366        &self,
367        partial_hash: Option<(PartialHashSearchType, String)>,
368        name: Option<String>,
369        response: PartialHashSearchType,
370        limit: u32,
371    ) -> Result<Vec<String>> {
372        let query = SearchRequest {
373            partial_hash,
374            file_name: name,
375            response,
376            limit,
377        };
378
379        ensure!(
380            query.is_valid(),
381            "Query isn't valid: hash isn't hexidecimal or both the hashes and file name are empty"
382        );
383
384        self.client
385            .post(format!("{}{}", self.url, malwaredb_api::SEARCH))
386            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
387            .json(&query)
388            .send()?
389            .json::<Vec<String>>()
390            .context("failed to receive or decode hash list, or invalid API key")
391    }
392
393    /// Retrieve sample by hash, optionally in the `CaRT` format
394    ///
395    /// # Errors
396    ///
397    /// This may return an error if there's a network situation or if the user is not logged in
398    /// or not properly authorized to connect.
399    pub fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
400        let api_endpoint = if cart {
401            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART)
402        } else {
403            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE)
404        };
405
406        let res = self
407            .client
408            .get(format!("{}{api_endpoint}", self.url))
409            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
410            .send()?;
411
412        if !res.status().is_success() {
413            bail!("Received code {}", res.status());
414        }
415
416        let body = res.bytes()?;
417        Ok(body.to_vec())
418    }
419
420    /// Fetch a report for a sample
421    ///
422    /// # Errors
423    ///
424    /// This may return an error if there's a network situation or if the user is not logged in
425    /// or not properly authorized to connect.
426    pub fn report(&self, hash: &str) -> Result<malwaredb_api::Report> {
427        self.client
428            .get(format!(
429                "{}{}/{hash}",
430                self.url,
431                malwaredb_api::SAMPLE_REPORT
432            ))
433            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
434            .send()?
435            .json::<malwaredb_api::Report>()
436            .context("failed to receive or decode sample report, or invalid API key")
437    }
438
439    /// Find similar samples in `MalwareDB` based on the contents of a given file.
440    /// This does not submit the sample to `MalwareDB`.
441    ///
442    /// # Errors
443    ///
444    /// This may return an error if there's a network situation or if the user is not logged in
445    /// or not properly authorized to connect.
446    pub fn similar(&self, contents: &[u8]) -> Result<malwaredb_api::SimilarSamplesResponse> {
447        let mut hashes = vec![];
448        let ssdeep_hash = FuzzyHash::new(contents);
449
450        let build_hasher = Murmur3HashState::default();
451        let lzjd_str =
452            LZDict::from_bytes_stream(contents.iter().copied(), &build_hasher).to_string();
453        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));
454        hashes.push((
455            malwaredb_api::SimilarityHashType::SSDeep,
456            ssdeep_hash.to_string(),
457        ));
458
459        let mut builder = TlshBuilder::new(
460            tlsh_fixed::BucketKind::Bucket256,
461            tlsh_fixed::ChecksumKind::ThreeByte,
462            tlsh_fixed::Version::Version4,
463        );
464
465        builder.update(contents);
466        if let Ok(hasher) = builder.build() {
467            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
468        }
469
470        if let Ok(exe) = EXE::from(contents) {
471            if let Some(imports) = exe.imports {
472                hashes.push((
473                    malwaredb_api::SimilarityHashType::ImportHash,
474                    hex::encode(imports.hash()),
475                ));
476                hashes.push((
477                    malwaredb_api::SimilarityHashType::FuzzyImportHash,
478                    imports.fuzzy_hash(),
479                ));
480            }
481        }
482
483        let request = malwaredb_api::SimilarSamplesRequest { hashes };
484
485        self.client
486            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES))
487            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
488            .json(&request)
489            .send()?
490            .json::<malwaredb_api::SimilarSamplesResponse>()
491            .context("failed to receive or decode similarity response, or invalid API key")
492    }
493}
494
495impl Debug for MdbClient {
496    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
497        use crate::MDB_VERSION;
498
499        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
500    }
501}
malwaredb_client/blocking.rs

malwaredb_client/
blocking.rs
