malwaredb-client-py 0.3.4

// SPDX-License-Identifier: Apache-2.0

use std::collections::HashMap;
use std::path::PathBuf;

use pyo3::{pyclass, pymethods, PyResult};
use uuid::Uuid;

/// Information about a label
#[pyclass(frozen)]
#[derive(Debug)]
pub struct Label {
    /// Label ID
    #[pyo3(get)]
    pub id: u64,

    /// Label value
    #[pyo3(get)]
    pub name: String,

    /// Label parent
    #[pyo3(get)]
    pub parent: Option<String>,
}

#[pymethods]
impl Label {
    /// Label printable representation
    #[must_use]
    pub fn __repr__(&self) -> String {
        if let Some(parent) = &self.parent {
            format!("Label {}: {parent}/{}", self.id, self.name)
        } else {
            format!("Label {}: {}", self.id, self.name)
        }
    }

    /// Label information, just the name
    #[must_use]
    pub fn __str__(&self) -> String {
        self.name.clone()
    }
}

impl From<malwaredb_client::malwaredb_api::Label> for Label {
    fn from(label: malwaredb_client::malwaredb_api::Label) -> Self {
        Self {
            id: label.id,
            name: label.name,
            parent: label.parent,
        }
    }
}

/// Information about the Malware DB server
#[pyclass(frozen, from_py_object)]
#[derive(Debug, Clone)]
pub struct ServerInfo {
    /// Operating System used
    #[pyo3(get)]
    pub os_name: String,

    /// Memory footprint
    #[pyo3(get)]
    pub memory_used: String,

    /// MDB version
    #[pyo3(get)]
    pub mdb_version: String,

    /// Type and version of the database
    #[pyo3(get)]
    pub db_version: String,

    /// Size of the database on disk
    #[pyo3(get)]
    pub db_size: String,

    /// Total number of samples
    #[pyo3(get)]
    pub num_samples: u64,

    /// Total number of users
    #[pyo3(get)]
    pub num_users: u32,

    /// Uptime of Malware DB in a human-readable format
    #[pyo3(get)]
    pub uptime: String,

    /// The name of the Malware DB instance
    #[pyo3(get)]
    pub instance_name: String,
}

#[pymethods]
impl ServerInfo {
    /// Server Info printable representation
    #[must_use]
    pub fn __repr__(&self) -> String {
        format!(
            "MalwareDB {} running on {} for {}",
            self.mdb_version, self.os_name, self.uptime
        )
    }

    /// Server info, just the version
    #[must_use]
    pub fn __str__(&self) -> String {
        format!("MalwareDB {}", self.mdb_version)
    }
}

impl From<malwaredb_client::malwaredb_api::ServerInfo> for ServerInfo {
    fn from(value: malwaredb_client::malwaredb_api::ServerInfo) -> Self {
        Self {
            os_name: value.os_name,
            memory_used: value.memory_used,
            mdb_version: value.mdb_version.to_string(),
            db_version: value.db_version,
            db_size: value.db_size,
            num_samples: value.num_samples,
            num_users: value.num_users,
            uptime: value.uptime,
            instance_name: value.instance_name,
        }
    }
}

/// Information about a sample source
#[pyclass(frozen)]
#[derive(Debug)]
pub struct Source {
    /// Identifier of the source
    #[pyo3(get)]
    pub id: u32,

    /// Name of the source
    #[pyo3(get)]
    pub name: String,

    /// Description of the source
    #[pyo3(get)]
    pub description: Option<String>,

    /// URL of the source, or where the files were found
    #[pyo3(get)]
    pub url: Option<String>,

    /// Creation date or first acquisition date of or from the source
    #[pyo3(get)]
    pub first_acquisition: String,

    /// Whether the source holds malware
    #[pyo3(get)]
    pub malicious: Option<bool>,
}

#[pymethods]
impl Source {
    /// Source printable representation
    #[must_use]
    pub fn __repr__(&self) -> String {
        let url = if let Some(url) = &self.url {
            format!(" from {url}")
        } else {
            String::new()
        };

        let desc = if let Some(desc) = &self.description {
            format!(" -- {desc}")
        } else {
            String::new()
        };

        format!("{}({}){url}{desc}", self.name, self.id)
    }

    /// Simpler source printable representation
    #[must_use]
    pub fn __str__(&self) -> String {
        self.name.clone()
    }
}

impl From<malwaredb_client::malwaredb_api::SourceInfo> for Source {
    fn from(value: malwaredb_client::malwaredb_api::SourceInfo) -> Self {
        Self {
            id: value.id,
            name: value.name,
            description: value.description,
            url: value.url,
            first_acquisition: value.first_acquisition.to_rfc3339(),
            malicious: value.malicious,
        }
    }
}

/// Information about a file type supported by Malware DB
#[pyclass(frozen)]
#[derive(Debug)]
pub struct SupportedFileType {
    /// Common name of the file type
    #[pyo3(get)]
    pub name: String,

    /// Magic number bytes in hex of the file type
    #[pyo3(get)]
    pub magic: Vec<String>,

    /// Whether the file type is executable
    #[pyo3(get)]
    pub is_executable: bool,

    /// Description of the file type
    #[pyo3(get)]
    pub description: Option<String>,
}

#[pymethods]
impl SupportedFileType {
    /// Supported file type as a printable representation
    #[must_use]
    pub fn __repr__(&self) -> String {
        format!("{}, starting with {}", self.name, self.magic.join(" or "))
    }

    /// Supported file type as a printable simpler representation
    #[must_use]
    pub fn __str__(&self) -> String {
        self.name.clone()
    }
}

impl From<malwaredb_client::malwaredb_api::SupportedFileType> for SupportedFileType {
    fn from(value: malwaredb_client::malwaredb_api::SupportedFileType) -> Self {
        Self {
            name: value.name,
            magic: value.magic,
            is_executable: value.is_executable,
            description: value.description,
        }
    }
}

/// Information about the user's account
#[pyclass(frozen)]
#[derive(Debug)]
pub struct UserInfo {
    /// User's numeric ID
    #[pyo3(get)]
    pub id: u32,

    /// User's name
    #[pyo3(get)]
    pub username: String,

    /// User's group memberships, if any
    #[pyo3(get)]
    pub groups: Vec<String>,

    /// User's available sample sources, if any
    #[pyo3(get)]
    pub sources: Vec<String>,

    /// If the user is an admin
    #[pyo3(get)]
    pub is_admin: bool,

    /// When the account was created
    #[pyo3(get)]
    pub created: String,

    /// User has read-only access, perhaps a guest or demo account
    #[pyo3(get)]
    pub is_readonly: bool,
}

#[pymethods]
impl UserInfo {
    /// Simple user information
    #[must_use]
    pub fn __str__(&self) -> String {
        self.username.clone()
    }
}

impl From<malwaredb_client::malwaredb_api::GetUserInfoResponse> for UserInfo {
    fn from(value: malwaredb_client::malwaredb_api::GetUserInfoResponse) -> Self {
        Self {
            id: value.id,
            username: value.username,
            groups: value.groups,
            sources: value.sources,
            is_admin: value.is_admin,
            created: value.created.to_rfc3339(),
            is_readonly: value.is_readonly,
        }
    }
}

/// Search results
#[pyclass(frozen)]
#[derive(Debug)]
pub struct SearchResults {
    /// Hashes of results
    #[pyo3(get)]
    pub hashes: Vec<String>,

    /// Total samples matching search
    #[pyo3(get)]
    pub results: u64,

    /// Identifier for continuing to the next batch of search results
    #[pyo3(get)]
    pub pagination: Option<Uuid>,

    /// Message from the server, if available
    #[pyo3(get)]
    pub message: Option<String>,
}

#[pymethods]
impl SearchResults {
    /// Search results as a printable representation
    #[must_use]
    pub fn __repr__(&self) -> String {
        if self.pagination.is_some() && self.results > self.hashes.len() as u64 {
            format!(
                "Search results showing {} of {} available hashes",
                self.hashes.len(),
                self.results
            )
        } else if let Some(message) = &self.message {
            format!(
                "{message}: Search results showing {} hashes",
                self.hashes.len()
            )
        } else {
            format!("Search results showing {} hashes", self.hashes.len())
        }
    }

    /// Search results as a printable simpler representation
    #[must_use]
    pub fn __str__(&self) -> String {
        if let Some(message) = &self.message {
            format!(
                "{message}: Search results with {} hashes",
                self.hashes.len()
            )
        } else {
            format!("Search results with {} hashes", self.hashes.len())
        }
    }
}

impl From<malwaredb_client::malwaredb_api::SearchResponse> for SearchResults {
    fn from(value: malwaredb_client::malwaredb_api::SearchResponse) -> Self {
        Self {
            hashes: value.hashes,
            results: value.total_results,
            pagination: value.pagination,
            message: value.message,
        }
    }
}

/// Yara search results
#[pyclass(frozen)]
#[derive(Debug)]
pub struct YaraResult {
    /// Search results: file hashes by Yara rule name
    #[pyo3(get)]
    pub results: HashMap<String, Vec<String>>,
}

impl YaraResult {
    /// Search results length
    #[must_use]
    pub fn __len__(&self) -> usize {
        self.results.len()
    }
}

impl From<malwaredb_client::malwaredb_api::YaraSearchResponse> for YaraResult {
    fn from(value: malwaredb_client::malwaredb_api::YaraSearchResponse) -> Self {
        Self {
            results: value
                .results
                .into_iter()
                .map(|(k, v)| (k, v.into_iter().map(|h| h.to_string()).collect()))
                .collect(),
        }
    }
}

/// Discovered local Malware DB server
#[pyclass(frozen, from_py_object)]
#[derive(Debug, Clone)]
pub struct DiscoveredServer {
    /// Server IP or domain
    #[pyo3(get)]
    pub host: String,

    /// Server port
    #[pyo3(get)]
    pub port: u16,

    /// If the server expects an encrypted connection
    #[pyo3(get)]
    pub ssl: bool,

    /// Malware DB server name
    #[pyo3(get)]
    pub name: String,

    /// Additional server information, if retrievable
    #[pyo3(get)]
    pub info: Option<ServerInfo>,
}

#[pymethods]
impl DiscoveredServer {
    /// Connect to a discovered server using login credentials
    ///
    /// # Errors
    ///
    /// Returns an error if the server becomes unreachable or if credentials are incorrect
    #[pyo3(signature = (username, password, save = false, cert_path = None))]
    pub fn connect(
        &self,
        username: String,
        password: String,
        save: bool,
        cert_path: Option<PathBuf>,
    ) -> PyResult<crate::MalwareDBClient> {
        crate::MalwareDBClient::login(self.__str__(), username, password, save, cert_path)
    }

    /// Discovered server as a string with server name
    #[must_use]
    pub fn __repr__(&self) -> String {
        if self.ssl {
            format!("{} at https://{}:{}", self.name, self.host, self.port)
        } else {
            format!("{} at http://{}:{}", self.name, self.host, self.port)
        }
    }

    /// Discovered server as a URL string
    #[must_use]
    pub fn __str__(&self) -> String {
        if self.ssl {
            format!("https://{}:{}", self.host, self.port)
        } else {
            format!("http://{}:{}", self.host, self.port)
        }
    }
}

impl From<malwaredb_client::MalwareDBServer> for DiscoveredServer {
    fn from(value: malwaredb_client::MalwareDBServer) -> Self {
        let info = value.server_info_blocking().ok();
        Self {
            host: value.host,
            port: value.port,
            ssl: value.ssl,
            name: value.name,
            info: info.map(Into::into),
        }
    }
}