shaha 0.2.1

Hash database builder and reverse lookup tool
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use std::path::PathBuf;

use anyhow::Result;
use clap::{Args, ValueEnum};

use crate::storage::{ParquetStorage, R2Storage, Storage};

#[derive(Clone, ValueEnum)]
pub enum OutputFormat {
    Plain,
    Json,
}

#[derive(Args)]
pub struct InfoArgs {
    #[arg(default_value = "hashes.parquet")]
    pub database: PathBuf,

    #[arg(short, long, default_value = "plain")]
    pub format: OutputFormat,
    #[command(flatten)]
    pub r2: super::R2Args,
}

pub fn run(args: InfoArgs) -> Result<()> {
    let (stats, location) = if args.r2.enabled {
        let r2_config = args.r2.build_config(&args.database)?;
        let url = r2_config.s3_url();
        let storage = R2Storage::new(r2_config)?;
        (storage.stats()?, url)
    } else {
        let storage = ParquetStorage::new(&args.database);
        (storage.stats()?, args.database.display().to_string())
    };

    match args.format {
        OutputFormat::Plain => print_plain(&location, &stats),
        OutputFormat::Json => print_json(&location, &stats)?,
    }

    Ok(())
}

fn print_plain(location: &str, stats: &crate::storage::Stats) {
    println!("Database:   {}", location);
    println!("Records:    {}", stats.total_records);
    if stats.file_size_bytes > 0 {
        println!("Size:       {}", format_bytes(stats.file_size_bytes));
    }
    println!(
        "Algorithms: {}",
        if stats.algorithms.is_empty() {
            "-".to_string()
        } else {
            stats.algorithms.join(", ")
        }
    );
    println!(
        "Sources:    {}",
        if stats.sources.is_empty() {
            "-".to_string()
        } else {
            stats.sources.join(", ")
        }
    );
}

fn print_json(location: &str, stats: &crate::storage::Stats) -> Result<()> {
    #[derive(serde::Serialize)]
    struct JsonInfo {
        database: String,
        total_records: usize,
        #[serde(skip_serializing_if = "Option::is_none")]
        file_size_bytes: Option<u64>,
        algorithms: Vec<String>,
        sources: Vec<String>,
    }

    let info = JsonInfo {
        database: location.to_string(),
        total_records: stats.total_records,
        file_size_bytes: if stats.file_size_bytes > 0 {
            Some(stats.file_size_bytes)
        } else {
            None
        },
        algorithms: stats.algorithms.clone(),
        sources: stats.sources.clone(),
    };

    println!("{}", serde_json::to_string_pretty(&info)?);
    Ok(())
}


fn format_bytes(bytes: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;

    if bytes >= GB {
        format!("{:.2} GB", bytes as f64 / GB as f64)
    } else if bytes >= MB {
        format!("{:.2} MB", bytes as f64 / MB as f64)
    } else if bytes >= KB {
        format!("{:.2} KB", bytes as f64 / KB as f64)
    } else {
        format!("{} B", bytes)
    }
}