malwaredb 0.3.2

Service for storing malicious, benign, or unknown files and related metadata and relationships.
// SPDX-License-Identifier: Apache-2.0

use malwaredb_api::{Label, Labels};
use malwaredb_server::State;

use std::process::ExitCode;

use anyhow::{bail, Result};
use clap::Parser;

/// Create a new label
#[derive(Clone, Debug, Parser, PartialEq)]
pub struct Create {
    /// Label ID
    #[arg(short, long)]
    pub name: String,

    /// Identify the parent label by ID
    #[arg(long = "pid")]
    pub parent_id: Option<u64>,

    /// Identify the parent label by name
    #[arg(long = "pname")]
    pub parent_name: Option<String>,
}

impl Create {
    pub async fn execute(&self, state: State) -> Result<ExitCode> {
        if self.parent_id.is_some() && self.parent_name.is_some() {
            bail!("Only one parent type may be used, not both");
        }

        let parent_id = if let Some(name) = &self.parent_name {
            state
                .db_type
                .get_labels()
                .await
                .expect("unable to list labels")
                .0
                .into_iter()
                .filter(|l| l.name == *name)
                .collect::<Vec<Label>>()
                .first()
                .map(|l| l.id)
        } else {
            self.parent_id
        };

        let lid = state.db_type.create_label(&self.name, parent_id).await?;
        println!("Created label `{}` with id {lid}.", self.name);

        Ok(ExitCode::SUCCESS)
    }
}

/// List labels
#[derive(Clone, Debug, Parser, PartialEq)]
pub struct List {
    /// Optionally filter labels (simple contains, nothing fancy)
    pub filter: Option<String>,
}

impl List {
    pub async fn execute(&self, state: State) -> Result<ExitCode> {
        let labels = state
            .db_type
            .get_labels()
            .await
            .expect("unable to list labels");
        let labels = if let Some(filter) = &self.filter {
            Labels(
                labels
                    .0
                    .into_iter()
                    .filter(|l| l.name.contains(filter))
                    .collect::<Vec<Label>>(),
            )
        } else {
            labels
        };

        if labels.is_empty() {
            if self.filter.is_some() {
                println!("No labels matched the filter");
            } else {
                println!("No labels created yet.");
            }
            return Ok(ExitCode::SUCCESS);
        }

        println!("{} labels:\n{labels}", labels.len());

        Ok(ExitCode::SUCCESS)
    }
}