malwaredb-client 0.3.4

// SPDX-License-Identifier: Apache-2.0

pub mod cart;
pub mod report;
pub mod retrieval;
pub mod search;
pub mod similar;
pub mod submit;

use crate::MdbClient;

use std::path::PathBuf;
use std::process::ExitCode;

use anyhow::Result;
use clap::{Command, CommandFactory, Parser, Subcommand};
use clap_complete::{generate, Shell};
use dialoguer::Password;

pub const VERSION: &str = concat!(env!("MDB_VERSION"), " ", env!("MDB_BUILD_DATE"));

/// Malware Database Client
///
/// Malware Database maintains the bookkeeping for unknown, malicious, and benign binaries
/// using a database, and optionally storing the files in a given location for later retrieval.
#[derive(Debug, Parser)]
#[command(author, about, version = VERSION)]
pub struct Options {
    #[clap(subcommand)]
    cmd: Subcommands,
}

impl Options {
    #[allow(clippy::too_many_lines)]
    pub async fn execute(&self) -> Result<ExitCode> {
        // Special case: there's no config file yet, which the others require.
        if let Subcommands::Login(cmd) = &self.cmd {
            return cmd.execute().await;
        }

        // Special case: utility to encode, decode, or view information for CaRT files.
        if let Subcommands::Cart(cmd) = &self.cmd {
            return cmd.execute();
        }

        // Special case: generating auto-complete script for a supported shell.
        if let Subcommands::Generate(cmd) = &self.cmd {
            return Ok(cmd.execute());
        }

        if self.cmd == Subcommands::Discover {
            let servers = malwaredb_client::discover_servers()?;
            if servers.is_empty() {
                println!("No MalwareDB servers found.");
            } else {
                for server in servers {
                    let extra = if let Ok(info) = server.server_info().await {
                        format!(
                            " — Version {} with {} samples, {}",
                            info.mdb_version, info.num_samples, info.db_size
                        )
                    } else {
                        String::new()
                    };
                    println!("{}\t{server}{extra}", server.name);
                }
            }
            return Ok(ExitCode::SUCCESS);
        }

        let mdb_client = MdbClient::load()?;

        match &self.cmd {
            Subcommands::Whoami => {
                let resp = mdb_client.whoami().await?;

                println!("UserID: {}", resp.id);
                if resp.groups.is_empty() {
                    println!("You aren't part of any groups.");
                } else {
                    println!("You're part of {} groups:", resp.groups.len());
                    for group in resp.groups {
                        println!("\t{group}");
                    }
                }
                // Access to sources is done by groups, so no groups already means no sources.
                if !resp.sources.is_empty() {
                    println!("You have access to {} sources:", resp.sources.len());
                    for source in resp.sources {
                        println!("\t{source}");
                    }
                }

                println!("Account creation: {}", resp.created);
                if resp.is_readonly {
                    println!("You are a read-only user.");
                }

                if resp.is_admin {
                    println!("You are an administrator of this MalwareDB instance.");
                }

                Ok(ExitCode::SUCCESS)
            }
            Subcommands::Login(_)
            | Subcommands::Cart(_)
            | Subcommands::Generate(_)
            | Subcommands::Discover => {
                unreachable!()
            }
            Subcommands::Logout => {
                mdb_client.delete()?;
                Ok(ExitCode::SUCCESS)
            }
            Subcommands::ResetKey => {
                mdb_client.reset_key().await?;
                mdb_client.delete()?;
                Ok(ExitCode::SUCCESS)
            }
            Subcommands::ServerInfo => {
                let resp = mdb_client.server_info().await?;
                println!("-- {} --", resp.instance_name);
                println!("MalwareDB version {} on {}", resp.mdb_version, resp.os_name);
                println!("Memory Used: {}", resp.memory_used);
                println!("Users: {}", resp.num_users);
                println!("Samples: {}", resp.num_samples);
                println!("Database size: {}", resp.db_size);
                println!("Uptime: {}", resp.uptime);

                if resp.mdb_version > *malwaredb_client::MDB_VERSION_SEMVER {
                    eprintln!(
                        "Server version {:?} is newer than client {:?}, consider updating.",
                        resp.mdb_version,
                        malwaredb_client::MDB_VERSION_SEMVER
                    );
                }

                Ok(ExitCode::SUCCESS)
            }
            Subcommands::ServerTypes => {
                let resp = mdb_client.supported_types().await?;
                for data_type in resp.types {
                    print!("{}", data_type.name);
                    if let Some(desc) = data_type.description {
                        print!(" {desc}");
                    }
                    if data_type.is_executable {
                        print!(" -- is executable");
                    }
                    println!();
                    for magic in data_type.magic {
                        println!("\t{magic}");
                    }
                }
                Ok(ExitCode::SUCCESS)
            }
            Subcommands::ListLabels => {
                let labels = mdb_client.labels().await?;
                println!("{labels}");
                Ok(ExitCode::SUCCESS)
            }
            Subcommands::ListSources => {
                let sources = mdb_client.sources().await?;
                if sources.sources.is_empty() {
                    println!("No sources available, are you part of any groups?");
                }
                for source in sources.sources {
                    print!("{}: {}", source.id, source.name);
                    if let Some(url) = source.url {
                        print!(" {url}");
                    }
                    if let Some(malicious) = source.malicious {
                        if malicious {
                            print!(" - malicious!");
                        }
                    }
                    println!();
                }
                Ok(ExitCode::SUCCESS)
            }
            Subcommands::SubmitSamples(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::SampleReport(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::RetrieveSample(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::FindSimilar(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::Search(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::YaraSearch(cmd) => cmd.exec(&mdb_client).await,
            Subcommands::YaraResult(cmd) => cmd.exec(&mdb_client).await,
        }
    }
}

#[derive(Subcommand, Clone, Debug, PartialEq)]
pub enum Subcommands {
    /// Discover local Malware DB servers via Multicast DNS (also known as Bonjour or Zeroconf)
    Discover,
    /// Show information about your account, including available groups and data sources
    Whoami,
    /// Login with your username and password to fetch & store your API key
    Login(Login),
    /// Remove your API key from this system
    Logout,
    /// Invalidate your API key, and logout from this system
    ResetKey,
    /// Show information about the server
    ServerInfo,
    /// Show the data types and magic numbers of supported file types
    ServerTypes,
    /// List labels, a hierarchical taxonomy for file samples
    ListLabels,
    /// List the sources available to the user
    ListSources,
    /// Submit one or more samples to the server by source ID
    SubmitSamples(submit::SubmitSamples),
    /// Retrieve a sample from the server by hash
    RetrieveSample(retrieval::RetrieveSample),
    /// Retrieve a report for a sample from the server by hash
    SampleReport(report::SampleReport),
    /// Find similar samples
    FindSimilar(similar::Similar),
    /// Pack or unpack `CaRT` files
    Cart(cart::CartIO),
    /// Shell autocomplete generation
    Generate(Generator),
    /// Search for samples based on partial hash and/or file name
    Search(search::SearchRequest),
    /// Search for samples using Yara rules
    YaraSearch(search::YaraSearch),
    /// Get the results of a Yara search
    YaraResult(search::YaraResult),
}

/// Login to a Malware DB instance, saving your API key in `.mdb_client.toml` in your home directory
#[derive(Clone, Debug, Parser, PartialEq)]
pub struct Login {
    /// URL of the server to connect to
    pub url: String,

    /// Username
    pub uname: String,

    /// Path to server cert or CA cert, if needed
    pub cert_path: Option<PathBuf>,
}

impl Login {
    async fn execute(&self) -> Result<ExitCode> {
        let password = Password::new()
            .with_prompt(format!("Password for {}", self.uname))
            .interact()?;

        if let Err(e) = MdbClient::login(
            self.url.clone(),
            self.uname.clone(),
            password,
            true,
            self.cert_path.clone(),
        )
        .await
        {
            eprintln!("{e}");
            Ok(ExitCode::FAILURE)
        } else {
            Ok(ExitCode::SUCCESS)
        }
    }
}

#[derive(Parser, Debug, Clone, PartialEq)]
pub struct Generator {
    #[arg(value_enum)]
    pub(crate) generator: Shell,
}

impl Generator {
    pub fn execute(&self) -> ExitCode {
        let mut cmd = Options::command();
        eprintln!("Generating completion file for {:?}...", self.generator);
        print_completions(self.generator, &mut cmd);
        ExitCode::SUCCESS
    }
}

fn print_completions<G: clap_complete::Generator>(gen: G, cmd: &mut Command) {
    generate(gen, cmd, cmd.get_name().to_string(), &mut std::io::stdout());
}

#[test]
fn verify_cli() {
    Options::command().debug_assert();
    Login::command().debug_assert();
}