malwaredb_client/

lib.rs

// SPDX-License-Identifier: Apache-2.0

#![doc = include_str!("../README.md")]
#![deny(missing_docs)]
#![deny(clippy::all)]
#![deny(clippy::pedantic)]
#![forbid(unsafe_code)]

/// Non-async version of the Malware DB client
#[cfg(feature = "blocking")]
pub mod blocking;
pub(crate) mod option_cert_path_serialization;

pub use malwaredb_api;
use malwaredb_lzjd::{LZDict, Murmur3HashState};
use malwaredb_types::exec::pe32::EXE;
use malwaredb_types::utils::entropy_calc;

use std::fmt::{Debug, Formatter};
use std::io::Cursor;
use std::path::{Path, PathBuf};

use anyhow::{bail, ensure, Context, Result};
use base64::engine::general_purpose;
use base64::Engine;
use cart_container::JsonMap;
use fuzzyhash::FuzzyHash;
use home::home_dir;
use malwaredb_api::{PartialHashSearchType, SearchRequest};
use reqwest::Certificate;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256, Sha384, Sha512};
use tlsh_fixed::TlshBuilder;
use tracing::{error, info};
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Local directory for the Malware DB client configs
const MDB_CLIENT_DIR: &str = "malwaredb_client";

/// Config file name expected by Malware DB client
const MDB_CLIENT_CONFIG_TOML: &str = "mdb_client.toml";

/// MDB version
pub const MDB_VERSION: &str = env!("CARGO_PKG_VERSION");

/// Asynchronous Malware DB Client Configuration and connection
#[derive(Deserialize, Serialize, Zeroize, ZeroizeOnDrop)]
pub struct MdbClient {
    /// URL of the Malware DB server, including http and port number, ending without a slash
    pub url: String,

    /// User's API key for Malware DB
    api_key: String,

    /// Certificate and Path, if needed
    /// The path is serialized; deserialization loads & parses the certificate file specified.
    #[zeroize(skip)]
    #[serde(default, with = "option_cert_path_serialization")]
    cert: Option<(Certificate, PathBuf)>,
}

impl MdbClient {
    /// MDB Client from components, doesn't test connectivity
    ///
    /// # Errors
    ///
    /// Returns an error if a list of certificates was passed and any were not in the expected
    /// DER or PEM format or could not be parsed.
    pub fn new(url: String, api_key: String, cert_path: Option<PathBuf>) -> Result<Self> {
        let mut url = url;
        let url = if url.ends_with('/') {
            url.pop();
            url
        } else {
            url
        };

        let cert = if let Some(path) = cert_path {
            Some((path_load_cert(&path)?, path))
        } else {
            None
        };

        Ok(Self { url, api_key, cert })
    }

    /// Generate a client which already knows to send the API key and asks for gzip or zstd responses.
    #[inline]
    fn client(&self) -> reqwest::Result<reqwest::Client> {
        let builder = reqwest::ClientBuilder::new()
            .gzip(true)
            .zstd(true)
            .use_rustls_tls()
            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));

        if let Some(cert) = &self.cert {
            builder.add_root_certificate(cert.0.clone()).build()
        } else {
            builder.build()
        }
    }

    /// Login to a server, optionally save the config file, and return a client object
    ///
    /// # Errors
    ///
    /// Returns an error if the server URL, username, or password were incorrect, or if a network
    /// issue occurred.
    pub async fn login(
        url: String,
        username: String,
        password: String,
        save: bool,
        cert_path: Option<PathBuf>,
    ) -> Result<Self> {
        let mut url = url;
        let url = if url.ends_with('/') {
            url.pop();
            url
        } else {
            url
        };

        let api_request = malwaredb_api::GetAPIKeyRequest {
            user: username,
            password,
        };

        let builder = reqwest::ClientBuilder::new()
            .gzip(true)
            .zstd(true)
            .use_rustls_tls()
            .user_agent(concat!("mdb_client/", env!("CARGO_PKG_VERSION")));

        let cert = if let Some(path) = cert_path {
            Some((path_load_cert(&path)?, path))
        } else {
            None
        };

        let client = if let Some(cert) = &cert {
            builder.add_root_certificate(cert.0.clone()).build()
        } else {
            builder.build()
        }?;

        let res = client
            .post(format!("{url}{}", malwaredb_api::USER_LOGIN_URL))
            .json(&api_request)
            .send()
            .await?
            .json::<malwaredb_api::GetAPIKeyResponse>()
            .await?;

        if let Some(key) = &res.key {
            let client = MdbClient {
                url,
                api_key: key.clone(),
                cert,
            };

            if save {
                if let Err(e) = client.save() {
                    error!("Login successful but failed to save config: {e}");
                    bail!("Login successful but failed to save config: {e}");
                }
            }
            Ok(client)
        } else {
            if let Some(msg) = &res.message {
                error!("Login failed, response: {msg}");
            }
            bail!("server error or bad credentials");
        }
    }

    /// Reset one's own API key to effectively logout & disable all clients who are using the key
    ///
    /// # Errors
    ///
    /// Returns an error if there was a network issue or the user wasn't properly logged in.
    pub async fn reset_key(&self) -> Result<()> {
        let response = self
            .client()?
            .get(format!("{}{}", self.url, malwaredb_api::USER_LOGOUT_URL))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await
            .context("server error, or invalid API key")?;
        if !response.status().is_success() {
            bail!("failed to reset API key, was it correct?");
        }
        Ok(())
    }

    /// MDB Client loaded from a specified path
    ///
    /// # Errors
    ///
    /// Returns an error if the configuration file cannot be read, possibly because it
    /// doesn't exist or due to a permission error or a parsing error.
    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
        let name = path.as_ref().display();
        let config =
            std::fs::read_to_string(&path).context(format!("failed to read config file {name}"))?;
        let cfg: MdbClient =
            toml::from_str(&config).context(format!("failed to parse config file {name}"))?;
        Ok(cfg)
    }

    /// MDB Client from user's home directory
    ///
    /// # Errors
    ///
    /// Returns an error if the configuration file cannot be read, possibly because it
    /// doesn't exist or due to a permission error or a parsing error.
    pub fn load() -> Result<Self> {
        let path = get_config_path(false)?;
        if path.exists() {
            return Self::from_file(path);
        }
        bail!("config file not found")
    }

    /// Save MDB Client to the user's home directory
    ///
    /// # Errors
    ///
    /// Returns an error if there was a problem saving the configuration file.
    pub fn save(&self) -> Result<()> {
        let toml = toml::to_string(self)?;
        let path = get_config_path(true)?;
        std::fs::write(&path, toml)
            .context(format!("failed to write mdb config to {}", path.display()))
    }

    /// Delete the `MalwareDB` client config file
    ///
    /// # Errors
    ///
    /// Returns an error if there isn't a configuration file to delete, or if it cannot be deleted,
    /// possibly due to a permissions error.
    pub fn delete(&self) -> Result<()> {
        let path = get_config_path(false)?;
        if path.exists() {
            std::fs::remove_file(&path).context(format!(
                "failed to delete client config file {}",
                path.display()
            ))?;
        }
        Ok(())
    }

    // Actions of the client

    /// Get information about the server, unauthenticated
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation.
    pub async fn server_info(&self) -> Result<malwaredb_api::ServerInfo> {
        self.client()?
            .get(format!("{}{}", self.url, malwaredb_api::SERVER_INFO))
            .send()
            .await?
            .json::<malwaredb_api::ServerInfo>()
            .await
            .context("failed to receive or decode server info")
    }

    /// Get file types supported by the server, unauthenticated
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation.
    pub async fn supported_types(&self) -> Result<malwaredb_api::SupportedFileTypes> {
        self.client()?
            .get(format!(
                "{}{}",
                self.url,
                malwaredb_api::SUPPORTED_FILE_TYPES
            ))
            .send()
            .await?
            .json::<malwaredb_api::SupportedFileTypes>()
            .await
            .context("failed to receive or decode server-supported file types")
    }

    /// Get information about the user
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn whoami(&self) -> Result<malwaredb_api::GetUserInfoResponse> {
        self.client()?
            .get(format!("{}{}", self.url, malwaredb_api::USER_INFO_URL))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::GetUserInfoResponse>()
            .await
            .context("failed to receive or decode user info, or invalid API key")
    }

    /// Get the sample labels known to the server
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn labels(&self) -> Result<malwaredb_api::Labels> {
        self.client()?
            .get(format!("{}{}", self.url, malwaredb_api::LIST_LABELS))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Labels>()
            .await
            .context("failed to receive or decode available labels, or invalid API key")
    }

    /// Get the sources available to the current user
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn sources(&self) -> Result<malwaredb_api::Sources> {
        self.client()?
            .get(format!("{}{}", self.url, malwaredb_api::LIST_SOURCES))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Sources>()
            .await
            .context("failed to receive or decode available labels, or invalid API key")
    }

    /// Submit one file to `MalwareDB`: provide the contents, file name, and source ID
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn submit(
        &self,
        contents: impl AsRef<[u8]>,
        file_name: impl AsRef<str>,
        source_id: u32,
    ) -> Result<bool> {
        let mut hasher = Sha256::new();
        hasher.update(&contents);
        let result = hasher.finalize();

        let encoded = general_purpose::STANDARD.encode(contents);

        let payload = malwaredb_api::NewSample {
            file_name: file_name.as_ref().to_string(),
            source_id,
            file_contents_b64: encoded,
            sha256: hex::encode(result),
        };

        match self
            .client()?
            .post(format!("{}{}", self.url, malwaredb_api::UPLOAD_SAMPLE))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .json(&payload)
            .send()
            .await
        {
            Ok(res) => {
                if !res.status().is_success() {
                    info!("Code {} sending {}", res.status(), payload.file_name);
                }
                Ok(res.status().is_success())
            }
            Err(e) => {
                let status: String = e
                    .status()
                    .map(|s| s.as_str().to_string())
                    .unwrap_or_default();
                error!("Error{status} sending {}: {e}", payload.file_name);
                bail!(e.to_string())
            }
        }
    }

    /// Search for a file based on partial hash and/or partial file name, returns a list of hashes
    ///
    /// # Errors
    ///
    /// * 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
    pub async fn partial_search(
        &self,
        partial_hash: Option<(PartialHashSearchType, String)>,
        name: Option<String>,
        response: PartialHashSearchType,
        limit: u32,
    ) -> Result<Vec<String>> {
        let query = SearchRequest {
            partial_hash,
            file_name: name,
            response,
            limit,
        };

        ensure!(
            query.is_valid(),
            "Query isn't valid: hash isn't hexidecimal or both the hashes and file name are empty"
        );

        self.client()?
            .post(format!("{}{}", self.url, malwaredb_api::SEARCH))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .json(&query)
            .send()
            .await?
            .json::<Vec<String>>()
            .await
            .context("failed to receive or decode hash list, or invalid API key")
    }

    /// Retrieve sample by hash, optionally in the `CaRT` format
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn retrieve(&self, hash: &str, cart: bool) -> Result<Vec<u8>> {
        let api_endpoint = if cart {
            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE_CART)
        } else {
            format!("{}{hash}", malwaredb_api::DOWNLOAD_SAMPLE)
        };

        let res = self
            .client()?
            .get(format!("{}{api_endpoint}", self.url))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?;

        if !res.status().is_success() {
            bail!("Received code {}", res.status());
        }

        let body = res.bytes().await?;
        Ok(body.to_vec())
    }

    /// Fetch a report for a sample
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn report(&self, hash: &str) -> Result<malwaredb_api::Report> {
        self.client()?
            .get(format!(
                "{}{}/{hash}",
                self.url,
                malwaredb_api::SAMPLE_REPORT
            ))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .send()
            .await?
            .json::<malwaredb_api::Report>()
            .await
            .context("failed to receive or decode sample report, or invalid API key")
    }

    /// Find similar samples in `MalwareDB` based on the contents of a given file.
    /// This does not submit the sample to `MalwareDB`.
    ///
    /// # Errors
    ///
    /// This may return an error if there's a network situation or if the user is not logged in
    /// or not properly authorized to connect.
    pub async fn similar(&self, contents: &[u8]) -> Result<malwaredb_api::SimilarSamplesResponse> {
        let mut hashes = vec![];
        let ssdeep_hash = FuzzyHash::new(contents);

        let build_hasher = Murmur3HashState::default();
        let lzjd_str =
            LZDict::from_bytes_stream(contents.iter().copied(), &build_hasher).to_string();
        hashes.push((malwaredb_api::SimilarityHashType::LZJD, lzjd_str));
        hashes.push((
            malwaredb_api::SimilarityHashType::SSDeep,
            ssdeep_hash.to_string(),
        ));

        let mut builder = TlshBuilder::new(
            tlsh_fixed::BucketKind::Bucket256,
            tlsh_fixed::ChecksumKind::ThreeByte,
            tlsh_fixed::Version::Version4,
        );

        builder.update(contents);
        if let Ok(hasher) = builder.build() {
            hashes.push((malwaredb_api::SimilarityHashType::TLSH, hasher.hash()));
        }

        if let Ok(exe) = EXE::from(contents) {
            if let Some(imports) = exe.imports {
                hashes.push((
                    malwaredb_api::SimilarityHashType::ImportHash,
                    hex::encode(imports.hash()),
                ));
                hashes.push((
                    malwaredb_api::SimilarityHashType::FuzzyImportHash,
                    imports.fuzzy_hash(),
                ));
            }
        }

        let request = malwaredb_api::SimilarSamplesRequest { hashes };

        self.client()?
            .post(format!("{}{}", self.url, malwaredb_api::SIMILAR_SAMPLES))
            .header(malwaredb_api::MDB_API_HEADER, &self.api_key)
            .json(&request)
            .send()
            .await?
            .json::<malwaredb_api::SimilarSamplesResponse>()
            .await
            .context("failed to receive or decode similarity response, or invalid API key")
    }
}

impl Debug for MdbClient {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "MDB Client v{MDB_VERSION}: {}", self.url)
    }
}

/// Convenience function for encoding bytes into a `CaRT` file using the default key. This also
/// adds SHA-384 and SHA-512 hashes plus the entropy of the original file.
/// See <https://github.com/CybercentreCanada/cart> for more information.
///
/// # Errors
///
/// There should not be any errors, but the underlying library can't guarantee that. However, any data
/// passed to it is correct.
pub fn encode_to_cart(data: &[u8]) -> Result<Vec<u8>> {
    let mut input_buffer = Cursor::new(data);
    let mut output_buffer = Cursor::new(vec![]);
    let mut output_metadata = JsonMap::new();

    let mut sha384 = Sha384::new();
    sha384.update(data);
    let sha384 = hex::encode(sha384.finalize());

    let mut sha512 = Sha512::new();
    sha512.update(data);
    let sha512 = hex::encode(sha512.finalize());

    output_metadata.insert("sha384".into(), sha384.into());
    output_metadata.insert("sha512".into(), sha512.into());
    output_metadata.insert("entropy".into(), entropy_calc(data).into());
    cart_container::pack_stream(
        &mut input_buffer,
        &mut output_buffer,
        Some(output_metadata),
        None,
        cart_container::digesters::default_digesters(),
        None,
    )?;

    Ok(output_buffer.into_inner())
}

/// Convenience function for decoding a `CaRT` file using the default key, returning the bytes plus the
/// optional header and footer metadata, if present.
/// See <https://github.com/CybercentreCanada/cart> for more information.
///
/// # Errors
///
/// Returns an error if the file cannot be parsed or if this `CaRT` file didn't use the default key.
/// <https://github.com/CybercentreCanada/cart-rs/blob/7ad548143bb85b64f364804e90cfada6c31cf902/cart_container/src/cipher.rs#L14-L17>
pub fn decode_from_cart(data: &[u8]) -> Result<(Vec<u8>, Option<JsonMap>, Option<JsonMap>)> {
    let mut input_buffer = Cursor::new(data);
    let mut output_buffer = Cursor::new(vec![]);
    let (header, footer) =
        cart_container::unpack_stream(&mut input_buffer, &mut output_buffer, None)?;
    Ok((output_buffer.into_inner(), header, footer))
}

/// Load a certificate from a path
///
/// # Errors
///
/// Returns errors if the file cannot be read or if the file isn't an ASN.1 DER file or
/// base64-encoded ASN.1 PEM file.
pub fn path_load_cert(path: &Path) -> Result<Certificate> {
    if !path.exists() {
        bail!("Certificate {path:?} does not exist.");
    }
    let cert = match path
        .extension()
        .context("can't determine file extension")?
        .to_str()
        .context("unable to parse file extension")?
    {
        "pem" => {
            let contents = std::fs::read(path)?;
            Certificate::from_pem(&contents)?
        }
        "der" => {
            let contents = std::fs::read(path)?;
            Certificate::from_der(&contents)?
        }
        ext => {
            bail!("Unknown extension {ext:?}")
        }
    };
    Ok(cert)
}

/// Gets the configuration file in the following order:
///
/// 1. Current working directory: `./mdb_client.toml`
/// 2. Haiku-specific directory if on Haiku
/// 3. XDG Free Desktop directory if on Unix
/// 4. The user's home directory in `~/.config/malwaredb_client/mdb_client.toml`
/// 5. `mdb_client.toml` in the current directory (same as the first, but after checking others)
#[inline]
pub(crate) fn get_config_path(create: bool) -> Result<PathBuf> {
    // If there is a config file in the current working directory, use it
    let config = PathBuf::from(MDB_CLIENT_CONFIG_TOML);
    if config.exists() {
        return Ok(config);
    }

    #[cfg(target_os = "haiku")]
    {
        let mut settings = PathBuf::from("/boot/home/config/settings/malwaredb");
        if create && !settings.exists() {
            std::fs::create_dir_all(&settings)?;
        }
        settings.push(MDB_CLIENT_CONFIG_TOML);
        return Ok(settings);
    }

    #[cfg(unix)]
    {
        // Obey the Free Desktop standard, check the variable
        if let Some(xdg_home) = std::env::var_os("XDG_CONFIG_HOME") {
            let mut xdg_config_home = PathBuf::from(xdg_home);
            xdg_config_home.push(MDB_CLIENT_DIR);
            if create && !xdg_config_home.exists() {
                std::fs::create_dir_all(&xdg_config_home)?;
            }
            xdg_config_home.push(MDB_CLIENT_CONFIG_TOML);
            return Ok(xdg_config_home);
        }
    }

    if let Some(mut home_config) = home_dir() {
        home_config.push(".config");
        home_config.push(MDB_CLIENT_DIR);
        if create && !home_config.exists() {
            std::fs::create_dir_all(&home_config)?;
        }
        home_config.push(MDB_CLIENT_CONFIG_TOML);
        return Ok(home_config);
    }

    Ok(PathBuf::from(MDB_CLIENT_CONFIG_TOML))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cart() {
        const BYTES: &[u8] = include_bytes!("../../crates/types/testdata/elf/elf_haiku_x86.cart");
        const ORIGINAL_SHA256: &str =
            "de10ba5e5402b46ea975b5cb8a45eb7df9e81dc81012fd4efd145ed2dce3a740";

        let (decoded, header, footer) = decode_from_cart(BYTES).unwrap();

        let mut sha256 = Sha256::new();
        sha256.update(&decoded);
        let sha256 = hex::encode(sha256.finalize());
        assert_eq!(sha256, ORIGINAL_SHA256);

        let header = header.unwrap();
        let entropy = header.get("entropy").unwrap().as_f64().unwrap();
        assert!(entropy > 4.0 && entropy < 4.1);

        let footer = footer.unwrap();
        assert_eq!(footer.get("length").unwrap(), "5093");
        assert_eq!(footer.get("sha256").unwrap(), ORIGINAL_SHA256);
    }
}