Merge pull request #199 from kpcyrd/stats

Improve stats command

src/cmd/select_cmd.rs

@@ -16,20 +16,24 @@ pub struct Args {
     #[structopt(subcommand)]
     subcommand: Target,
     /// Print json output
-    #[structopt(long="json", group="output")]
+    #[structopt(long, group="output")]
     json: bool,
+    /// Only print the value instead of the whole object
+    #[structopt(long, group="output")]
+    values: bool,
     /// Print paths to blobs
-    #[structopt(long="paths", group="output")]
+    #[structopt(long, group="output")]
     paths: bool,
     /// Count rows returned
-    #[structopt(short="c", long="count", group="output")]
+    #[structopt(short="c", group="output")]
     count: bool,
 }
 
 #[derive(PartialEq)]
 enum Output {
     Normal,
     Json,
+    Values,
     Paths,
     Count,
 }
@@ -43,6 +47,8 @@ impl<'a, 'b> Printer<'a, 'b> {
     pub fn new(rl: &'a mut Shell<'b>, args: &Args) -> Printer<'a, 'b> {
         let output = if args.json {
             Output::Json
+        } else if args.values {
+            Output::Values
         } else if args.paths {
             Output::Paths
         } else if args.count {
@@ -70,6 +76,9 @@ impl<'a, 'b> Printer<'a, 'b> {
                         let v = serde_json::to_string(&obj)?;
                         println!("{}", v);
                     },
+                    Output::Values => {
+                        println!("{}", obj.to_string());
+                    },
                     Output::Paths => {
                         let blob = obj.blob()
                             .ok_or_else(|| format_err!("This model isn't linked to blob storage"))?;

src/cmd/stats_cmd.rs

@@ -1,42 +1,54 @@
-use crate::errors::*;
-
+use crate::args;
 use colored::Colorize;
 use crate::blobs::BlobStorage;
 use crate::cmd::Cmd;
-use crate::db::Database;
-use crate::db::ttl;
+use crate::db::{ttl, Database};
+use crate::errors::*;
 use crate::models::*;
-use crate::shell::Shell;
+use crate::shell::{self, Shell};
+use crate::workspaces;
 use humansize::{FileSize, file_size_opts};
 use separator::Separatable;
 use serde::{Serialize, Deserialize};
 use structopt::StructOpt;
 use structopt::clap::AppSettings;
 
-#[derive(Debug, StructOpt)]
+#[derive(Debug, Clone, StructOpt)]
 #[structopt(global_settings = &[AppSettings::ColoredHelp])]
 pub struct Args {
     /// Exclude blob storage
-    #[structopt(long)]
+    #[structopt(short, long)]
     short: bool,
+    /// Exclude categories that don't contain any structs
+    #[structopt(short, long)]
+    quiet: bool,
     /// Show workspace statistics in json
-    #[structopt(long)]
+    #[structopt(short, long)]
     json: bool,
+    /// Go through all workspaces
+    #[structopt(short, long)]
+    all: bool,
 }
 
-fn show_amount(label: &str, count: usize, amount: &str) {
-    let label = format!("{:20}", label);
-    let amount = format!("{:>20}", amount);
+impl Args {
+    fn show_amount(&self, label: &str, count: usize, amount: &str) {
+        if self.quiet && count == 0 {
+            return;
+        }
+
+        let label = format!("{:20}", label);
+        let amount = format!("{:>20}", amount);
 
-    if count > 0 {
-        println!("{} {}", label.green(), amount.yellow());
-    } else {
-        println!("{} {}", label, amount);
+        if count > 0 {
+            println!("{} {}", label.green(), amount.yellow());
+        } else {
+            println!("{} {}", label, amount);
+        }
     }
-}
 
-fn show_count(label: &str, count: usize) {
-    show_amount(label, count, &count.separated_string());
+    fn show_count(&self, label: &str, count: usize) {
+        self.show_amount(label, count, &count.separated_string());
+    }
 }
 
 fn count_models<T: Model>(db: &Database) -> Result<usize> {
@@ -119,40 +131,66 @@ impl Stats {
 
 impl Cmd for Args {
     #[inline]
-    fn run(self, rl: &mut Shell) -> Result<()> {
-        ttl::reap_expired(rl)?;
+    fn run(mut self, rl: &mut Shell) -> Result<()> {
+        if self.all {
+            self.all = false;
+            let mut first = true;
+
+            let config = rl.config();
+            for ws in workspaces::list()? {
+                if first {
+                    first = false;
+                } else if !self.json {
+                    println!();
+                }
+
+                let mut rl = shell::init(&args::Args {
+                    workspace: Some(ws),
+                    subcommand: None,
+                }, config, false)?;
+                self.clone().run(&mut rl)?;
+            }
+        } else {
+            ttl::reap_expired(rl)?;
 
-        let db = rl.db();
+            let db = rl.db();
 
-        let mut stats = Stats::count(rl.workspace().into(), &db)?;
-        if !self.short {
-            stats.add_blob_usage(rl.blobs())?;
-        }
+            let workspace = rl.workspace();
 
-        if self.json {
-            let stats = serde_json::to_string(&stats)?;
-            println!("{}", stats);
-        } else {
-            println!("{:>41}", stats.workspace.bold());
-            show_count("domains", stats.domains);
-            show_count("subdomains", stats.subdomains);
-            show_count("ipaddrs", stats.ipaddrs);
-            show_count("urls", stats.urls);
-            show_count("emails", stats.emails);
-            show_count("phonenumbers", stats.phonenumbers);
-            show_count("devices", stats.devices);
-            show_count("networks", stats.networks);
-            show_count("accounts", stats.accounts);
-            show_count("breaches", stats.breaches);
-            show_count("images", stats.images);
-            show_count("ports", stats.ports);
-            show_count("netblocks", stats.netblocks);
-            show_count("cryptoaddrs", stats.cryptoaddrs);
-            show_count("activity", stats.activity);
-
-            if let Some(blobs) = stats.blobs {
-                show_count("blobs", blobs.count);
-                show_amount("blobs (size)", blobs.total_size as usize, &blobs.total_human_size);
+            if !self.json {
+                // print this early in case counting takes a while
+                println!("{:>41}", workspace.bold());
+            }
+
+            let mut stats = Stats::count(workspace.into(), &db)?;
+            if !self.short {
+                stats.add_blob_usage(rl.blobs())?;
+            }
+
+            if self.json {
+                let stats = serde_json::to_string(&stats)?;
+                println!("{}", stats);
+            } else {
+                self.show_count("domains", stats.domains);
+                self.show_count("subdomains", stats.subdomains);
+                self.show_count("ipaddrs", stats.ipaddrs);
+                self.show_count("urls", stats.urls);
+                self.show_count("emails", stats.emails);
+                self.show_count("phonenumbers", stats.phonenumbers);
+                self.show_count("devices", stats.devices);
+                self.show_count("networks", stats.networks);
+                self.show_count("accounts", stats.accounts);
+                self.show_count("breaches", stats.breaches);
+                self.show_count("images", stats.images);
+                self.show_count("ports", stats.ports);
+                self.show_count("netblocks", stats.netblocks);
+                self.show_count("cryptoaddrs", stats.cryptoaddrs);
+                self.show_count("activity", stats.activity);
+
+                if let Some(blobs) = stats.blobs {
+                    self.show_count("blobs", blobs.count);
+                    self.show_amount("blobs (size)", blobs.total_size as usize, &blobs.total_human_size);
+                }
             }
         }
 

src/cmd/workspace_cmd.rs

@@ -24,7 +24,7 @@ pub struct Args {
     /// Skip confirmation
     #[structopt(short = "f", long = "force")]
     force: bool,
-    workspace: Option<Workspace>,
+    workspaces: Vec<Workspace>,
 }
 
 fn delete(workspace: Workspace, force: bool) -> Result<()> {
@@ -38,13 +38,13 @@ fn delete(workspace: Workspace, force: bool) -> Result<()> {
     Ok(())
 }
 
-fn usage(workspace: Option<Workspace>) -> Result<()> {
-    if let Some(ws) = workspace {
-        println!("{}", ws.usage_human()?);
-    } else {
-        for ws in workspaces::list()? {
-            println!("{:50}: {}", ws.as_str(), ws.usage_human()?);
-        }
+fn usage(mut workspaces: Vec<Workspace>) -> Result<()> {
+    if workspaces.is_empty() {
+        workspaces = workspaces::list()?;
+    }
+
+    for ws in workspaces {
+        println!("{:50}: {}", ws.as_str(), ws.usage_human()?);
     }
 
     Ok(())
@@ -67,30 +67,34 @@ fn list() -> Result<()> {
     Ok(())
 }
 
-fn run(args: Args, rl: Option<&mut Shell>) -> Result<()> {
+fn run(mut args: Args, rl: Option<&mut Shell>) -> Result<()> {
     if args.delete {
-        if let Some(workspace) = args.workspace {
-            if let Some(rl) = rl {
+        if args.workspaces.is_empty() {
+            bail!("--delete requires workspace");
+        }
+
+        for workspace in args.workspaces {
+            if let Some(rl) = &rl {
                 if *rl.db().workspace() == workspace {
                     bail!("Can't delete current workspace")
                 }
             }
-
-            delete(workspace, args.force)
-        } else {
-            bail!("--delete requires workspace")
+            delete(workspace, args.force)?;
         }
     } else if args.usage {
-        usage(args.workspace)
-    } else if let Some(workspace) = args.workspace {
-        if let Some(rl) = rl {
-            change(rl, workspace)
-        } else {
-            Ok(())
-        }
+        usage(args.workspaces)?;
     } else {
-        list()
+        match args.workspaces.len() {
+            0 => list()?,
+            1 => if let Some(rl) = rl {
+                // we've already tested there's one item in the list
+                change(rl, args.workspaces.pop().unwrap())?;
+            },
+            _ => bail!("Only one argument allowed when switching workspaces"),
+        }
     }
+
+    Ok(())
 }
 
 impl Cmd for Args {

src/main.rs

@@ -1,4 +1,4 @@
-use env_logger::{self, Env};
+use env_logger::Env;
 use sn0int::args::{self, Args, SubCommand};
 use sn0int::auth;
 use sn0int::cmd::{self, LiteCmd};