malware-modeler 0.0.5

// SPDX-License-Identifier: Apache-2.0

#![deny(clippy::all)]
//#![deny(clippy::cargo)]
#![deny(clippy::pedantic)]
#![forbid(unsafe_code)]

use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::process::ExitCode;

use anyhow::{Result, bail};
use clap::{Command, CommandFactory, Parser, Subcommand, ValueHint};
use clap_complete::{Shell, generate};
use walkdir::WalkDir;

/// Malware Modeler
///
/// Create malicious vs benign models
#[derive(Parser)]
#[command(author, about, version = malware_modeler::VERSION)]
pub struct Args {
    /// Subcommands (with their own options)
    #[clap(subcommand)]
    cmd: Actions,

    /// Number of threads to use for parallel operations.
    #[clap(long, default_value_t = num_cpus::get())]
    threads: usize,
}

#[derive(Subcommand)]
enum Actions {
    /// Find n-grams
    Ngram(Ngram),

    /// Create dataset from samples and existing n-grams file
    Dataset(Dataset),

    /// Convert between dataset file formats, using the file extension to detect file types
    DatasetConvert(DatasetConvert),

    /// Train a malware model
    Train(Train),

    /// Evaluate a dataset against a model
    Evaluate(Evaluate),

    /// Model operations
    Model(ModelInfo),

    /// Show type information or sort files by type with symbolic links and SHA-256 hashes.
    Types(Types),

    /// Check similarity between files using LZJD
    Similarity(Similarity),

    /// Unzip files from a Zip archive by the type, if the type is known
    SelectiveUnzip(SelectiveUnzip),

    /// What's inside that Zip file??
    ZipSummary(ZipSummary),

    /// Generate shell autocomplete script
    Generate(Generator),
}

/// Find n-grams
#[derive(Parser)]
struct Ngram {
    /// Directory containing benign and malicious samples
    #[arg(long, value_hint = ValueHint::DirPath)]
    path: PathBuf,

    /// Optional: specify the type of file to be processed.
    /// If not specified, the file type will be detected automatically.
    #[arg(short = 't', long = "type", ignore_case = true)]
    ftype: Option<malware_modeler::ftype::FileType>,

    /// Size of the byte ngram
    #[arg(default_value = "6")]
    n: u16,

    /// Max number of n-grams to keep
    #[arg(default_value = "10000")]
    k: usize,

    /// Output n-grams path
    #[arg(long, value_hint = ValueHint::FilePath, default_value = "ngrams.txt")]
    output: PathBuf,

    /// Store the number of n-gram occurrences
    #[clap(long, short, action)]
    counts: bool,

    /// Show some additional information
    #[clap(long, default_value = "false", action)]
    verbose: bool,
}

impl Ngram {
    fn execute(&self) -> Result<ExitCode> {
        if self.k < 10 {
            eprintln!("NGrams must be at least 10, should be ~100k");
            return Ok(ExitCode::FAILURE);
        }
        if self.n < 2 {
            eprintln!("NGrams must be at least 2, should be >4");
            return Ok(ExitCode::FAILURE);
        }

        let mut ngrammer = malware_modeler::ngram::Ngrammer::new(
            self.ftype,
            &self.path,
            self.n,
            self.k,
            self.verbose,
        )?;
        ngrammer.find();
        ngrammer.save(&self.output, self.counts)?;

        Ok(ExitCode::SUCCESS)
    }
}

/// Create dataset from samples and existing n-grams file
#[derive(Parser)]
struct Dataset {
    /// Directory containing malicious samples
    #[arg(short, long, value_hint = ValueHint::DirPath)]
    pub malicious: PathBuf,

    /// Directory containing benign samples
    #[arg(short, long, value_hint = ValueHint::DirPath)]
    pub benign: PathBuf,

    /// N-grams file for featurization
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub ngrams: PathBuf,

    /// Output dataset file
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub output: PathBuf,
}

impl Dataset {
    fn execute(&self) -> Result<ExitCode> {
        malware_modeler::dataset::Dataset::create_save_from_benign_malicious_files_and_ngrams(
            &self.malicious,
            &self.benign,
            &self.ngrams,
            &self.output,
        )?;

        Ok(ExitCode::SUCCESS)
    }
}

/// Convert between dataset file formats, using the file extension to detect file types
#[derive(Parser)]
struct DatasetConvert {
    /// Input file
    #[arg(short, long = "in", value_hint = ValueHint::FilePath)]
    pub input: PathBuf,

    /// Output file
    #[arg(short, long = "out", value_hint = ValueHint::FilePath)]
    pub output: PathBuf,

    /// Optionally provide a model, which is used to reduce the dataset by removing unneeded features.
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub model: Option<PathBuf>,
}

impl DatasetConvert {
    fn execute(&self) -> Result<ExitCode> {
        if self.input.extension() == self.output.extension() && self.model.is_none() {
            eprintln!("Refusing to attempt 'conversion' between files with the same extension.");
            return Ok(ExitCode::FAILURE);
        }

        let mut dataset = malware_modeler::dataset::Dataset::load(&self.input)?;

        if let Some(model) = &self.model {
            let original = dataset.data[0].len();
            let model_contents = std::fs::read_to_string(model)?;
            let model: malware_modeler::model::LogisticRegression =
                serde_json::from_str(model_contents.as_str())?;
            let removed = dataset.reduce(&model)?.len();
            println!("Removed {removed} features from the original {original}");
        }

        dataset.save(&self.output)?;
        Ok(ExitCode::SUCCESS)
    }
}

/// Train a malware model
#[derive(Parser)]
struct Train {
    /// Dataset for training
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub dataset: PathBuf,

    /// Dataset for testing
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub test: Option<PathBuf>,

    /// Output model path
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    pub model: PathBuf,

    /// Number of iterations for training
    #[arg(short, long, default_value_t = 100)]
    pub epochs: u32,

    /// Learning rate
    #[arg(long, default_value_t = 0.1)]
    pub learning_rate: f32,

    /// L1 penalty to reduce features
    #[arg(long, default_value_t = 0.1)]
    pub l1_pentalty: f32,

    /// L2 penalty
    #[arg(long, default_value_t = 0.1)]
    pub l2_penalty: f32,
}

impl Train {
    fn execute(&self) -> Result<ExitCode> {
        let mut dataset = malware_modeler::dataset::Dataset::load(&self.dataset)?;
        let mut model = malware_modeler::model::LogisticRegression::new(
            dataset.data[0].len(),
            self.learning_rate,
            self.l1_pentalty,
            self.l2_penalty,
        );
        match model.train(self.epochs, &mut dataset) {
            Ok(error) => {
                println!("Training complete, final error {error}");
            }
            Err(error) => {
                eprintln!("Training failed: {error}");
                return Ok(ExitCode::FAILURE);
            }
        }

        let result = model.evaluate_dataset(&dataset)?;
        println!("{result}");
        println!("Accuracy: {:.2}", result.accuracy());
        println!("Precision: {:.2}", result.precision());
        println!("Recall: {:.2}", result.recall());
        println!("F1: {:.2}", result.f1());
        println!("Auc: {:.2}", result.auc());

        if let Some(test_dataset) = &self.test {
            let dataset = malware_modeler::dataset::Dataset::load(test_dataset)?;
            let result = model.evaluate_dataset(&dataset)?;
            println!("{result}");
            println!("Accuracy: {:.2}", result.accuracy());
            println!("Precision: {:.2}", result.precision());
            println!("Recall: {:.2}", result.recall());
            println!("F1: {:.2}", result.f1());
            println!("Auc: {:.2}", result.auc());

            model.test_performance = Some(result.into());
        }

        let original = model.weights.len();
        model.set_features_and_reduce(&dataset)?;
        let reduced = model.weights.len();
        let diff = original - reduced;
        println!("Features reduced by {diff}: {original} - {reduced}");

        let model_json = serde_json::to_string_pretty(&model)?;
        std::fs::write(&self.model, model_json)?;
        Ok(ExitCode::SUCCESS)
    }
}

/// Evaluate a dataset against a model
#[derive(Parser)]
struct Evaluate {
    /// Path to the dataset
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    dataset: PathBuf,

    /// Path to the model
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    model: PathBuf,

    /// If the model file should be updated to reflect evaluation performance of this test dataset
    #[arg(long, default_value_t = false)]
    update_model: bool,
}

impl Evaluate {
    fn execute(&self) -> Result<ExitCode> {
        let dataset = malware_modeler::dataset::Dataset::load(&self.dataset)?;
        let model_contents = std::fs::read_to_string(&self.model)?;
        let mut model: malware_modeler::model::LogisticRegression =
            serde_json::from_str(model_contents.as_str())?;
        let result = model.evaluate_dataset(&dataset)?;
        println!("{result}");
        println!("Accuracy: {:.2}", result.accuracy());
        println!("Precision: {:.2}", result.precision());
        println!("Recall: {:.2}", result.recall());
        println!("F1: {:.2}", result.f1());
        println!("Auc: {:.2}", result.auc());

        if self.update_model {
            model.test_performance = Some(result.into());
            let model_json = serde_json::to_string_pretty(&model)?;
            std::fs::write(&self.model, model_json)?;
        }

        Ok(ExitCode::SUCCESS)
    }
}

#[derive(Parser)]
struct ModelInfo {
    /// Path to the model
    #[arg(short, long, value_hint = ValueHint::FilePath)]
    model: PathBuf,

    /// Extract n-grams, useful for when a new dataset might be used and the model
    /// features were reduced
    #[arg(short = 'n', long, value_hint = ValueHint::FilePath)]
    extract_ngrams: Option<PathBuf>,

    /// Evaluate a file or directory of files
    #[arg(short, long, value_hint = ValueHint::AnyPath)]
    eval: Option<PathBuf>,
}

impl ModelInfo {
    fn execute(&self) -> Result<ExitCode> {
        let model_contents = std::fs::read_to_string(&self.model)?;
        let model: malware_modeler::model::LogisticRegression =
            serde_json::from_str(model_contents.as_str())?;

        println!("Model has {} weights", model.weights.len());

        if let Some(extract_ngrams) = &self.extract_ngrams {
            let features = model
                .features
                .keys()
                .map(hex::encode)
                .collect::<Vec<String>>()
                .join("\n");
            let mut file = File::create(extract_ngrams)?;
            file.write_all(features.as_bytes())?;
            file.write_all(b"\n")?;
            file.sync_all()?;
        }

        if let Some(eval) = &self.eval {
            if eval.is_file() {
                let (label, prediction, features) = model.evaluate_file(eval)?;
                println!(
                    "{} is predicted to be {label} (raw: {prediction}) and had {features} of {} features",
                    eval.display(),
                    model.features.len()
                );
            } else if eval.is_dir() {
                println!("Path, PredictionLabel, PredictionRaw, FoundFeatures/AvailableFeatures");
                for entry in WalkDir::new(eval)
                    .max_depth(malware_modeler::MAX_RECURSION_DEPTH)
                    .follow_links(true)
                    .into_iter()
                    .flatten()
                {
                    if entry.file_type().is_file() {
                        let (label, prediction, features) = model.evaluate_file(entry.path())?;
                        println!(
                            "{}, {label}, {prediction}, {features}/{}",
                            entry.path().display(),
                            model.features.len()
                        );
                    }
                }
            } else {
                eprintln!("{} is not a file or directory", eval.display());
            }
        }

        Ok(ExitCode::SUCCESS)
    }
}

/// Show type information or sort files by type with symbolic links and SHA-256 hashes.
#[derive(Parser)]
struct Types {
    #[arg(value_hint = ValueHint::AnyPath)]
    path: PathBuf,

    /// Optionally create symbolic links based on file type to separate types for further actions.
    #[arg(value_hint = ValueHint::DirPath)]
    destination: Option<PathBuf>,

    /// Specify depth for optional type sorting.
    #[arg(long, default_value_t = 3)]
    depth: u8,
}

impl Types {
    pub fn execute(&self) -> Result<ExitCode> {
        if self.path.is_file() {
            let contents = std::fs::read(&self.path)?;
            let ftype = malware_modeler::sorting::FileTypeUnion::from_bytes(&contents);
            println!("{ftype}");
        } else if self.path.is_dir() && self.destination.is_none() {
            for entry in WalkDir::new(&self.path)
                .max_depth(malware_modeler::MAX_RECURSION_DEPTH)
                .follow_links(true)
                .into_iter()
                .flatten()
            {
                if entry.file_type().is_file() {
                    let contents = std::fs::read(entry.path())?;
                    let ftype = malware_modeler::sorting::FileTypeUnion::from_bytes(&contents);
                    println!("{}: {ftype}", entry.path().display());
                }
            }
        } else if self.path.is_dir()
            && let Some(dest) = &self.destination
        {
            let destination = dest.to_owned();

            if destination.is_file() {
                eprintln!("Destination must be a directory, not a file");
                return Ok(ExitCode::FAILURE);
            }

            let result =
                malware_modeler::sorting::file_sorting(&self.path, &destination, self.depth)?;
            println!(
                "Countered {} files, {} of which were duplicates.",
                result.total_files, result.files_removed
            );
            if result.errors > 0 {
                eprintln!("Encountered {} errors while sorting files.", result.errors);
            }
        }
        Ok(ExitCode::SUCCESS)
    }
}

/// Check similarity between files using LZJD
#[derive(Parser, Debug, Clone, PartialEq)]
struct Similarity {
    /// If a file, show the LZJD hash
    one: PathBuf,

    /// Optionally a file if comparing two files
    two: Option<PathBuf>,

    /// Threshold when reporting files too similar
    #[arg(long, default_value_t = 0.8)]
    threshold: f32,

    /// Delete a file which is too similar
    #[arg(long, default_value_t = false)]
    delete_file_if_too_similar_yes_i_really_mean_it: bool,
}

impl Similarity {
    pub fn execute(&self) -> Result<ExitCode> {
        if self.one.is_file() {
            if let Some(two) = &self.two {
                if two.is_file() {
                    let sim = malware_modeler::similarity::lzjd_compare_paths(&self.one, two)?;
                    println!("{sim:.4}");
                }
            } else {
                let lzjd = malware_modeler::similarity::lzjd_from_path(&self.one)?;
                println!("{lzjd}");
            }
        } else {
            if self.two.is_some() {
                bail!("Similarity does not a directory argument with another path");
            }

            let similarity_check = malware_modeler::similarity::Similarity {
                path: &self.one,
                threshold: self.threshold,
            };

            let similar_count = similarity_check.find(|a, b, sim| {
                eprintln!("{} & {} too similar: {sim}", a.display(), b.display());
                if self.delete_file_if_too_similar_yes_i_really_mean_it
                    && std::fs::exists(a).unwrap_or_default()
                    && let Err(e) = std::fs::remove_file(a)
                {
                    eprintln!("Failed to remove {}: {e}", a.display());
                }
            })?;
            println!("Total similar file pairs: {similar_count}");
        }

        Ok(ExitCode::SUCCESS)
    }
}

/// Unzip files from a Zip archive by the type, if the type is known
/// Made with [VirusShare](https://virusshare.com) in mind.
#[derive(Parser, Debug, Clone, PartialEq)]
struct SelectiveUnzip {
    /// Path to the Zip archive
    #[arg(value_hint = ValueHint::FilePath)]
    source: PathBuf,

    /// Destination for the extracted files
    #[arg(value_hint = ValueHint::DirPath)]
    destination: PathBuf,

    /// Password if the Zip archive is encrypted.
    #[arg(short, long)]
    password: Option<String>,

    /// Specify depth for optional type sorting.
    #[arg(long, default_value_t = 3)]
    depth: u8,

    /// The specific file type which should be kept
    #[arg(short = 't', long = "type", value_parser = malware_modeler::sorting::parse_file_type_union)]
    file_type: Option<malware_modeler::sorting::FileTypeUnion>,

    /// Specify the string which should be checked against libmagic, also known as the `file` command.
    /// Ideally should not be used with `-t`, `--type`.
    #[cfg(feature = "libmagic")]
    #[arg(short = 'm', long = "magic")]
    magic: Option<String>,

    /// If unknown files should be kept when keeping all known types
    #[arg(long, default_value_t = false)]
    keep_unknowns: bool,
}

impl SelectiveUnzip {
    pub fn execute(self) -> Result<ExitCode> {
        let extracted_count = malware_modeler::unzip::unzip_files_by_type(
            self.source,
            self.destination,
            self.password.as_ref(),
            self.depth,
            self.file_type,
            #[cfg(feature = "libmagic")]
            self.magic.as_deref(),
            self.keep_unknowns,
        )?;

        println!("Extracted {extracted_count} files");
        Ok(ExitCode::SUCCESS)
    }
}

/// What's inside that Zip file??
#[derive(Parser, Debug, Clone, PartialEq)]
struct ZipSummary {
    /// Path to the Zip archive
    #[arg(value_hint = ValueHint::FilePath)]
    source: PathBuf,

    /// Password if the Zip archive is encrypted.
    #[arg(short, long)]
    password: Option<String>,

    /// If unknowns are encountered, how many leading bytes should be captured to display?
    #[arg(short, long, default_value_t = 0)]
    unknown_magic: usize,

    /// If unknowns are encountered, each has to have this minimum occurrence to be displayed.
    #[arg(long, default_value_t = 2)]
    unknown_minimum: usize,
}

impl ZipSummary {
    pub fn execute(self) -> Result<ExitCode> {
        const WIDTH: usize = 14;

        let zip_summary = malware_modeler::unzip::zip_file_type_counts(
            self.source,
            self.password.as_ref(),
            self.unknown_magic,
        )?;

        println!("{0:<WIDTH$} {1:<WIDTH$}", "Type", "Count");
        let mut ordered_summary: Vec<_> = zip_summary.file_type_counts.into_iter().collect();
        ordered_summary.sort_by(|a, b| b.1.cmp(&a.1));
        for (ftype, count) in ordered_summary {
            println!("{:<WIDTH$} {count:<WIDTH$}", format!("{ftype}"));
        }

        println!("{:^WIDTH$} {:<WIDTH$}", "Total", zip_summary.total_files);

        if !zip_summary.unknown_magic_counts.is_empty() {
            #[cfg(feature = "libmagic")]
            println!("\n{0:<WIDTH$} {1:<WIDTH$} Magic", "Unknown Header", "Count");
            #[cfg(not(feature = "libmagic"))]
            println!("\n{0:<WIDTH$} {1:<WIDTH$}", "Unknown Header", "Count");

            let mut ordered_summary: Vec<_> =
                zip_summary.unknown_magic_counts.into_iter().collect();
            ordered_summary.sort_by(|a, b| b.1.cmp(&a.1));

            #[cfg(feature = "libmagic")]
            for (unknown_bytes, (count, magic)) in ordered_summary {
                if count >= self.unknown_minimum {
                    println!(
                        "{:<WIDTH$} {count:<WIDTH$} {magic}",
                        hex::encode(unknown_bytes)
                    );
                }
            }

            #[cfg(not(feature = "libmagic"))]
            for (unknown_bytes, count) in ordered_summary {
                if count >= self.unknown_minimum {
                    println!("{:<WIDTH$} {count:<WIDTH$}", hex::encode(unknown_bytes));
                }
            }
        }

        Ok(ExitCode::SUCCESS)
    }
}

/// How to generate shell completions.
/// Instructions based on [Clap](https://github.com/clap-rs/clap/blob/master/clap_complete/examples/completion-derive.rs).
///
/// Usage for zsh:
/// ```console
/// $ cargo run -- generate zsh > /usr/local/share/zsh/site-functions/_completion_derive
/// $ compinit
/// $ ./target/debug/malware-modeler <TAB>
/// ```
///
/// Usage for Fish:
/// ```console
/// $ cargo run -- generate fish > completion_derive.fish
/// $ . ./completion_derive.fish
/// $ ./target/debug/malware-modeler <TAB>
/// ```
///
/// Usage for Bash:
/// ```console
/// $ cargo run -- generate bash > completion_derive.bash
/// $ source completion_derive.bash
/// $ ./target/debug/malware-modeler <TAB>
/// ```
#[derive(Parser, Debug, Clone, PartialEq)]
pub struct Generator {
    /// Shell for which to create the completion script
    #[arg(value_enum)]
    shell: Shell,
}

impl Generator {
    #[must_use]
    pub fn execute(&self) -> ExitCode {
        let mut cmd = Args::command();
        self.print_completions(&mut cmd);
        ExitCode::SUCCESS
    }

    fn print_completions(&self, cmd: &mut Command) {
        generate(
            self.shell,
            cmd,
            cmd.get_name().to_string(),
            &mut std::io::stdout(),
        );
    }
}

fn main() -> Result<ExitCode> {
    let args = Args::parse();
    rayon::ThreadPoolBuilder::new()
        .num_threads(args.threads)
        .build_global()?;

    let ret = match args.cmd {
        Actions::Ngram(ngram) => ngram.execute(),
        Actions::Dataset(dataset) => dataset.execute(),
        Actions::DatasetConvert(dataset) => dataset.execute(),
        Actions::Train(training) => training.execute(),
        Actions::Evaluate(evaluate) => evaluate.execute(),
        Actions::Model(model) => model.execute(),
        Actions::Types(types) => types.execute(),
        Actions::Similarity(sim) => sim.execute(),
        Actions::SelectiveUnzip(unzip) => unzip.execute(),
        Actions::ZipSummary(summary) => summary.execute(),
        Actions::Generate(generator) => Ok(generator.execute()),
    };

    if let Some(mem) = app_memory_usage_fetcher::get_memory_usage_string() {
        println!("Memory used: {mem}");
    }

    ret
}

#[test]
fn verify_cli() {
    use clap::CommandFactory;

    Args::command().debug_assert();
}