Allow filtering targets

src/cmd/mod.rs

@@ -7,4 +7,5 @@ pub mod select_cmd;
 pub mod mod_cmd;
 pub mod noscope_cmd;
 pub mod scope_cmd;
+pub mod target_cmd;
 pub mod quickstart_cmd;

src/cmd/run_cmd.rs

@@ -1,6 +1,6 @@
 use errors::*;
 
-use db::Database;
+use db::{Database, Filter};
 use sn0int_common::metadata::Source;
 use serde::Serialize;
 use serde_json;
@@ -21,10 +21,9 @@ fn prepare_arg<T: Serialize + Model>(x: T) -> Result<(serde_json::Value, Option<
     Ok((arg, Some(pretty)))
 }
 
-fn prepare_args<T: Scopable + Serialize + Model>(db: &Database) -> Result<Vec<(serde_json::Value, Option<String>)>> {
-    db.list::<T>()?
+fn prepare_args<T: Scopable + Serialize + Model>(db: &Database, filter: &Filter) -> Result<Vec<(serde_json::Value, Option<String>)>> {
+    db.filter::<T>(filter)?
         .into_iter()
-        .filter(|x| x.scoped())
         .map(prepare_arg)
         .collect()
 }
@@ -34,12 +33,14 @@ pub fn execute(rl: &mut Readline) -> Result<()> {
         .map(|m| m.to_owned())
         .ok_or_else(|| format_err!("No module selected"))?;
 
+    let filter = rl.scoped_targets();
+
     let args = match module.source() {
-        Some(Source::Domains) => prepare_args::<Domain>(rl.db()),
-        Some(Source::Subdomains) => prepare_args::<Subdomain>(rl.db()),
-        Some(Source::IpAddrs) => prepare_args::<IpAddr>(rl.db()),
-        Some(Source::Urls) => prepare_args::<Url>(rl.db()),
-        Some(Source::Emails) => prepare_args::<Email>(rl.db()),
+        Some(Source::Domains) => prepare_args::<Domain>(rl.db(), &filter),
+        Some(Source::Subdomains) => prepare_args::<Subdomain>(rl.db(), &filter),
+        Some(Source::IpAddrs) => prepare_args::<IpAddr>(rl.db(), &filter),
+        Some(Source::Urls) => prepare_args::<Url>(rl.db(), &filter),
+        Some(Source::Emails) => prepare_args::<Email>(rl.db(), &filter),
         None => Ok(vec![(serde_json::Value::Null, None)]),
     };
 

src/cmd/target_cmd.rs

@@ -0,0 +1,51 @@
+use errors::*;
+
+use db;
+use shell::Readline;
+use sn0int_common::metadata::Source;
+use structopt::StructOpt;
+use models::*;
+
+
+#[derive(Debug, StructOpt)]
+pub struct Args {
+    // TODO: target -p # print current filter
+    // TODO: target -c # clear filter
+
+    filter: Vec<String>,
+}
+
+pub fn run(rl: &mut Readline, args: &[String]) -> Result<()> {
+    let args = Args::from_iter_safe(args)?;
+
+    let source = rl.module()
+        .ok_or_else(|| format_err!("No module selected"))
+        .map(|x| x.source().clone())?;
+
+    if args.filter.is_empty() {
+        match source {
+            Some(Source::Domains) => select::<Domain>(rl)?,
+            Some(Source::Subdomains) => select::<Subdomain>(rl)?,
+            Some(Source::IpAddrs) => select::<IpAddr>(rl)?,
+            Some(Source::Urls) => select::<Url>(rl)?,
+            Some(Source::Emails) => select::<Email>(rl)?,
+            None => bail!("Module doesn't have sources"),
+        }
+    } else {
+        debug!("Setting filter to {:?}", args.filter);
+        let filter = db::Filter::parse_optional(&args.filter)?;
+        rl.set_target(Some(filter));
+    }
+
+    Ok(())
+}
+
+fn select<T: Model + Detailed>(rl: &mut Readline) -> Result<()> {
+    let filter = rl.scoped_targets();
+
+    for obj in rl.db().filter::<T>(&filter)? {
+        println!("{}", obj.detailed(rl.db())?);
+    }
+
+    Ok(())
+}

src/db.rs

@@ -338,6 +338,11 @@ impl Filter {
         &self.query
     }
 
+    pub fn and_scoped(&self) -> Filter {
+        let query = format!("({}) AND unscoped=0", self.query);
+        Filter::new(query)
+    }
+
     pub fn sql(&self) -> SqlLiteral<Bool> {
         sql::<Bool>(&self.query)
     }

src/shell.rs

@@ -6,7 +6,7 @@ use complete::CmdCompleter;
 use config::Config;
 use colored::Colorize;
 use ctrlc;
-use db::Database;
+use db::{self, Database};
 use engine::{Engine, Module};
 use geoip::{GeoIP, AsnDB, Maxmind};
 use rustyline::error::ReadlineError;
@@ -34,6 +34,7 @@ pub enum Command {
     Set,
     Select,
     SwitchDb,
+    Target,
     Update,
     Use,
     Quickstart,
@@ -54,6 +55,7 @@ impl Command {
             Command::Set => "set",
             Command::Select => "select",
             Command::SwitchDb => "switch_db",
+            Command::Target => "target",
             Command::Update => "update",
             Command::Use => "use",
             Command::Quickstart => "quickstart",
@@ -74,6 +76,7 @@ impl Command {
                 Command::Set.as_str(),
                 Command::Select.as_str(),
                 Command::SwitchDb.as_str(),
+                Command::Target.as_str(),
                 Command::Update.as_str(),
                 Command::Use.as_str(),
                 Command::Quickstart.as_str(),
@@ -99,6 +102,7 @@ impl FromStr for Command {
             "set"  => Ok(Command::Set),
             "select" => Ok(Command::Select),
             "switch_db" => Ok(Command::SwitchDb),
+            "target"  => Ok(Command::Target),
             "update" => Ok(Command::Update),
             "use"  => Ok(Command::Use),
             "quickstart"  => Ok(Command::Quickstart),
@@ -151,12 +155,29 @@ impl Readline {
 
     pub fn set_module(&mut self, module: Module) {
         self.prompt.module = Some(module);
+        // TODO: possibly refactor
+        self.prompt.target = None;
     }
 
     pub fn module(&self) -> Option<&Module> {
         self.prompt.module.as_ref()
     }
 
+    pub fn set_target(&mut self, target: Option<db::Filter>) {
+        self.prompt.target = target;
+    }
+
+    pub fn target(&self) -> &Option<db::Filter> {
+        &self.prompt.target
+    }
+
+    pub fn scoped_targets(&self) -> db::Filter {
+        match &self.prompt.target {
+            Some(filter) => filter.and_scoped(),
+            _ => db::Filter::new("unscoped=0"),
+        }
+    }
+
     pub fn db(&self) -> &Database {
         &self.db
     }
@@ -310,6 +331,7 @@ pub fn run_once(rl: &mut Readline) -> Result<bool> {
         Some((Command::Set, _args)) => println!("set"),
         Some((Command::Select, args)) => select_cmd::run(rl, &args)?,
         Some((Command::SwitchDb, args)) => switch_db_cmd::run(rl, &args)?,
+        Some((Command::Target, args)) => target_cmd::run(rl, &args)?,
         Some((Command::Update, _args)) => {
             // TODO
             // worker::spawn("Updating public suffix list");

src/term.rs

@@ -1,3 +1,4 @@
+use db;
 use engine::Module;
 use rand::prelude::*;
 use std::fmt;
@@ -129,13 +130,16 @@ pub fn error(line: &str) {
 pub struct Prompt {
     pub workspace: String,
     pub module: Option<Module>,
+    // TODO: wrapper type that holds module+options
+    pub target: Option<db::Filter>,
 }
 
 impl Prompt {
     pub fn new(workspace: String) -> Prompt {
         Prompt {
             workspace,
             module: None,
+            target: None,
         }
     }
 }