Add stats command

Cargo.lock

@@ -1593,6 +1593,12 @@ version = "1.3.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "615caabe2c3160b313d52ccc905335f4ed5f10881dd63dc5699d47e90be85691"
 
+[[package]]
+name = "humansize"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6cab2627acfc432780848602f3f558f7e9dd427352224b0d9324025796d2a5e"
+
 [[package]]
 name = "humantime"
 version = "1.3.0"
@@ -3928,6 +3934,7 @@ dependencies = [
  "glob",
  "hlua-badtouch",
  "hmac",
+ "humansize",
  "ipnetwork",
  "lazy_static",
  "libsqlite3-sys",

Cargo.toml

@@ -79,6 +79,7 @@ ipnetwork = "0.17"
 strum = "0.20"
 strum_macros = "0.20"
 embedded-triple = "0.1.0"
+humansize = "1.1.0"
 
 digest = "0.9"
 md-5 = "0.9"

src/args.rs

@@ -81,6 +81,9 @@ pub enum SubCommand {
     /// Export a workspace for external processing
     #[structopt(name="export")]
     Export(cmd::export_cmd::Args),
+    /// Show statistics about your current workspace
+    #[structopt(name="stats")]
+    Stats(cmd::stats_cmd::Args),
     /// Run a lua repl
     #[structopt(name="repl")]
     Repl,

src/blobs.rs

@@ -93,6 +93,14 @@ impl BlobStorage {
         }
         Ok(blobs)
     }
+
+    pub fn stat(&self, id: &str) -> Result<u64> {
+        let path = self.join(id)?;
+        debug!("Stat-ing blob: {:?}", path);
+        let md = fs::metadata(path)
+            .context("Failed to stat blob")?;
+        Ok(md.len())
+    }
 }
 
 

src/cmd/mod.rs

@@ -34,6 +34,7 @@ pub mod notify_cmd;
 pub mod pkg_cmd;
 pub mod set_cmd;
 pub mod scope_cmd;
+pub mod stats_cmd;
 pub mod target_cmd;
 pub mod quickstart_cmd;
 pub mod workspace_cmd;

src/cmd/stats_cmd.rs

@@ -0,0 +1,84 @@
+use crate::errors::*;
+
+use colored::Colorize;
+use crate::cmd::Cmd;
+use crate::db::Database;
+use crate::db::ttl;
+use crate::models::*;
+use crate::shell::Shell;
+use humansize::{FileSize, file_size_opts};
+use separator::Separatable;
+use structopt::StructOpt;
+use structopt::clap::AppSettings;
+
+#[derive(Debug, StructOpt)]
+#[structopt(global_settings = &[AppSettings::ColoredHelp])]
+pub struct Args {
+    /// Exclude blob storage
+    #[structopt(long)]
+    short: bool,
+}
+
+fn show_amount(label: &str, count: usize, amount: &str) -> Result<()> {
+    let label = format!("{:20}", label);
+    let amount = format!("{:>20}", amount);
+
+    if count > 0 {
+        println!("{} {}", label.green(), amount.yellow());
+    } else {
+        println!("{} {}", label, amount);
+    }
+
+    Ok(())
+}
+
+fn show_count(label: &str, count: usize) -> Result<()> {
+    show_amount(label, count, &count.separated_string())
+}
+
+fn count_models<T: Model>(label: &str, db: &Database) -> Result<()> {
+    let query = db.list::<T>()?;
+    let count = query.len();
+    show_count(label, count)
+}
+
+impl Cmd for Args {
+    #[inline]
+    fn run(self, rl: &mut Shell) -> Result<()> {
+        ttl::reap_expired(rl)?;
+
+        let db = rl.db();
+        println!("{:>41}", rl.workspace().bold());
+        count_models::<Domain>("domain", &db)?;
+        count_models::<Subdomain>("subdomains", &db)?;
+        count_models::<IpAddr>("ipaddrs", &db)?;
+        count_models::<Url>("urls", &db)?;
+        count_models::<Email>("emails", &db)?;
+        count_models::<PhoneNumber>("phonenumbers", &db)?;
+        count_models::<Device>("devices", &db)?;
+        count_models::<Network>("networks", &db)?;
+        count_models::<Account>("accounts", &db)?;
+        count_models::<Breach>("breaches", &db)?;
+        count_models::<Image>("images", &db)?;
+        count_models::<Port>("ports", &db)?;
+        count_models::<Netblock>("netblocks", &db)?;
+        count_models::<CryptoAddr>("cryptoaddrs", &db)?;
+        show_count("activity", Activity::count(&db)?)?;
+
+        if !self.short {
+            let storage = rl.blobs();
+            let blobs = storage.list()?;
+            show_count("blobs", blobs.len())?;
+
+            let mut total_size = 0;
+            for blob in &blobs {
+                total_size += storage.stat(blob)?;
+            }
+            let human_size = total_size.file_size(file_size_opts::CONVENTIONAL)
+                .map_err(|e| format_err!("Failed to format size: {}", e))?;
+            show_amount("blobs (size)", total_size as usize, &human_size)?;
+        }
+
+        Ok(())
+    }
+}

src/main.rs

@@ -128,6 +128,7 @@ fn run() -> Result<()> {
         Some(SubCommand::Export(export)) => run_cmd(&args, export, &config),
         Some(SubCommand::Cal(cal)) => run_cmd(&args, cal, &config),
         Some(SubCommand::Notify(notify)) => run_cmd(&args, notify, &config),
+        Some(SubCommand::Stats(stats)) => run_cmd(&args, stats, &config),
         Some(SubCommand::Repl) => repl::run(&config),
         Some(SubCommand::Paths) => paths::run(&config),
         Some(SubCommand::Completions(completions)) => complete::run_generate(&completions),

src/models/activity.rs

@@ -83,6 +83,14 @@ impl Activity {
             .optional()
             .map_err(Error::from)
     }
+
+    pub fn count(db: &Database) -> Result<usize> {
+        use crate::schema::activity::dsl::*;
+        activity.count()
+            .get_result::<i64>(db.db())
+            .map(|x| x as usize)
+            .map_err(Error::from)
+    }
 }
 
 pub struct ActivityFilter {

src/shell/mod.rs

@@ -45,6 +45,7 @@ pub enum Command {
     Scope,
     Set,
     Select,
+    Stats,
     Target,
     Use,
     Quickstart,
@@ -76,6 +77,7 @@ impl Command {
             Command::Scope => "scope",
             Command::Set => "set",
             Command::Select => "select",
+            Command::Stats => "stats",
             Command::Target => "target",
             Command::Use => "use",
             Command::Quickstart => "quickstart",
@@ -104,6 +106,7 @@ impl Command {
                 Command::Scope.as_str(),
                 Command::Set.as_str(),
                 Command::Select.as_str(),
+                Command::Stats.as_str(),
                 Command::Target.as_str(),
                 Command::Use.as_str(),
                 Command::Quit.as_str(),
@@ -137,6 +140,7 @@ impl FromStr for Command {
             "scope"  => Ok(Command::Scope),
             "set"  => Ok(Command::Set),
             "select" => Ok(Command::Select),
+            "stats" => Ok(Command::Stats),
             "target"  => Ok(Command::Target),
             "use"  => Ok(Command::Use),
             "quickstart"  => Ok(Command::Quickstart),
@@ -504,6 +508,7 @@ pub fn run_once(rl: &mut Shell) -> Result<bool> {
         Some((Command::Scope, args)) => scope_cmd::run(rl, &args)?,
         Some((Command::Set, args)) => set_cmd::run(rl, &args)?,
         Some((Command::Select, args)) => cmd::<select_cmd::Args>(rl, &args)?,
+        Some((Command::Stats, args)) => cmd::<stats_cmd::Args>(rl, &args)?,
         Some((Command::Target, args)) => target_cmd::run(rl, &args)?,
         Some((Command::Use, args)) => use_cmd::run(rl, &args)?,
         Some((Command::Quickstart, args)) => quickstart_cmd::run(rl, &args)?,