Repository: kpcyrd/sn0int (language: rust)

Issue:
Allow exposing stdin to scripts
In some cases we want to pass a file to a script. Since we don't expose any functions to access the filesystem this is currently not possible, but we could add an `--stdin` mode that enables some functions to access stdin. This would allow parsing some text based files or output of some tools:
```
cat dhcpd.leases | sn0int run --stdin -f ./dhcpd-parser.lua
```


Discussion / hints:

Fix patch:
diff --git a/docs/reference.rst b/docs/reference.rst
index 571367b7..feceba55 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -404,6 +404,18 @@ Update the label of the progress indicator.
 
     status('ohai')
 
+stdin_readline
+--------------
+
+Read a line from stdin. The final newline is not removed.
+
+.. code-block:: lua
+
+    stdin_readline()
+
+.. note::
+   This only works with `sn0int run --stdin`.
+
 url_join
 --------
 
diff --git a/modules/dev/arp-scan.lua b/modules/dev/arp-scan.lua
new file mode 100644
index 00000000..a31b4178
--- /dev/null
+++ b/modules/dev/arp-scan.lua
@@ -0,0 +1,21 @@
+-- Description: Parse arp-scan output
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+-- sudo arp-scan -qglI wlp3s0
+
+function run()
+    while true do
+        x = stdin_readline()
+        if x == nil then
+            break
+        end
+
+        m = regex_find('(.+)\t(.+)', x)
+        if m ~= nil then
+            ip = m[2]
+            mac = m[3]
+            info(json_encode({ip, mac}))
+        end
+    end
+end
diff --git a/modules/harness/stdin.lua b/modules/harness/stdin.lua
new file mode 100644
index 00000000..9f50b314
--- /dev/null
+++ b/modules/harness/stdin.lua
@@ -0,0 +1,13 @@
+-- Description: Read from stdin
+-- Version: 0.1.0
+-- License: GPL-3.0
+
+function run()
+    while true do
+        x = stdin_readline()
+        if x == nil then
+            break
+        end
+        info(json_encode(x))
+    end
+end
diff --git a/sn0int-registry/sn0int-common/src/metadata.rs b/sn0int-registry/sn0int-common/src/metadata.rs
index b2b3aa2e..cf9ab5b4 100644
--- a/sn0int-registry/sn0int-common/src/metadata.rs
+++ b/sn0int-registry/sn0int-common/src/metadata.rs
@@ -171,6 +171,21 @@ mod tests {
         });
     }
 
+    #[test]
+    fn verify_no_source() {
+        let metadata = Metadata::from_str(r#"-- Description: Hello world, this is my description
+-- Version: 1.0.0
+-- License: WTFPL
+
+"#).expect("parse");
+        assert_eq!(metadata, Metadata {
+            description: "Hello world, this is my description".to_string(),
+            version: "1.0.0".to_string(),
+            license: License::WTFPL,
+            source: None,
+        });
+    }
+
     #[test]
     fn verify_require_license() {
         let metadata = Metadata::from_str(r#"-- Description: Hello world, this is my description
diff --git a/src/args.rs b/src/args.rs
index 0c332400..b34ebaa0 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -51,13 +51,21 @@ pub enum SubCommand {
 
 #[derive(Debug, StructOpt)]
 pub struct Run {
+    /// Execute a module that has been installed
     pub module: Option<String>,
     #[structopt(short="f", long="file", conflicts_with="module")]
+    /// Run a module from a path
     pub file: Option<String>,
     #[structopt(short="j", long="threads", default_value="1")]
+    /// Run modules concurrently
     pub threads: usize,
     #[structopt(short="v", long="verbose", parse(from_occurrences))]
+    /// Verbose logging, once to print inserts even if they don't add new
+    /// data, twice to activate the debug() function
     pub verbose: u64,
+    #[structopt(long="stdin")]
+    /// Expose stdin to modules
+    pub stdin: bool,
 }
 
 #[derive(Debug, StructOpt)]
diff --git a/src/cmd/run_cmd.rs b/src/cmd/run_cmd.rs
index 5f78aec7..b8c11023 100644
--- a/src/cmd/run_cmd.rs
+++ b/src/cmd/run_cmd.rs
@@ -32,7 +32,7 @@ fn prepare_args<T: Scopable + Serialize + Model>(db: &Database, filter: &Filter)
         .collect()
 }
 
-pub fn execute(rl: &mut Readline, threads: usize, verbose: u64) -> Result<()> {
+pub fn execute(rl: &mut Readline, threads: usize, verbose: u64, has_stdin: bool) -> Result<()> {
     let module = rl.module()
         .map(|m| m.to_owned())
         .ok_or_else(|| format_err!("No module selected"))?;
@@ -49,7 +49,7 @@ pub fn execute(rl: &mut Readline, threads: usize, verbose: u64) -> Result<()> {
     }?;
 
     rl.signal_register().catch_ctrl();
-    worker::spawn(rl, &module, args, threads, verbose);
+    worker::spawn(rl, &module, args, threads, verbose, has_stdin);
     rl.signal_register().reset_ctrlc();
 
     term::info(&format!("Finished {}", module.canonical()));
@@ -59,5 +59,5 @@ pub fn execute(rl: &mut Readline, threads: usize, verbose: u64) -> Result<()> {
 
 pub fn run(rl: &mut Readline, args: &[String]) -> Result<()> {
     let args = Args::from_iter_safe(args)?;
-    execute(rl, args.threads, args.verbose)
+    execute(rl, args.threads, args.verbose, false)
 }
diff --git a/src/engine/ctx.rs b/src/engine/ctx.rs
index 11eb4ffd..9945f853 100644
--- a/src/engine/ctx.rs
+++ b/src/engine/ctx.rs
@@ -13,7 +13,7 @@ use std::result;
 use std::sync::{Arc, Mutex};
 use chrootable_https::dns::Resolver;
 use web::{HttpSession, HttpRequest, RequestOptions};
-use worker::{Event, LogEvent, DatabaseEvent};
+use worker::{Event, LogEvent, DatabaseEvent, StdioEvent};
 
 
 pub trait State {
@@ -73,6 +73,13 @@ pub trait State {
         reply.map_err(|err| format_err!("Failed to update database: {:?}", err))
     }
 
+    fn stdin_readline(&self) -> Result<Option<String>> {
+        self.send(&Event::Stdio(StdioEvent {}));
+        let reply = self.recv()?;
+        let reply: result::Result<Option<String>, String> = serde_json::from_value(reply)?;
+        reply.map_err(|err| format_err!("Failed to read stdin: {:?}", err))
+    }
+
     fn dns_config(&self) -> Arc<Resolver>;
 
     fn psl(&self) -> Arc<Psl>;
@@ -231,6 +238,7 @@ fn ctx<'a>(env: Environment) -> (hlua::Lua<'a>, Arc<LuaState>) {
     runtime::regex_find_all(&mut lua, state.clone());
     runtime::sleep(&mut lua, state.clone());
     runtime::status(&mut lua, state.clone());
+    runtime::stdin_readline(&mut lua, state.clone());
     runtime::url_join(&mut lua, state.clone());
     runtime::url_parse(&mut lua, state.clone());
     runtime::utf8_decode(&mut lua, state.clone());
diff --git a/src/engine/isolation.rs b/src/engine/isolation.rs
index 76db634c..c4c52818 100644
--- a/src/engine/isolation.rs
+++ b/src/engine/isolation.rs
@@ -4,11 +4,11 @@ use engine::{Environment, Module, Reporter};
 use geoip::{GeoIP, AsnDB};
 use psl::Psl;
 use serde_json;
-use worker::{Event, Event2, LogEvent, DatabaseEvent, ExitEvent, EventSender};
+use worker::{Event, Event2, LogEvent, ExitEvent, EventSender, EventWithCallback};
 
 use std::env;
 use std::io::prelude::*;
-use std::io::{self, BufReader, BufRead, Stdin, Stdout};
+use std::io::{self, BufReader, BufRead, stdin, Stdin, Stdout};
 use std::sync::{mpsc, Arc, Mutex};
 use std::process::{Command, Child, Stdio, ChildStdin, ChildStdout};
 
@@ -23,7 +23,6 @@ pub struct StartCommand {
 
 impl StartCommand {
     pub fn new(verbose: u64, dns_config: Resolver, module: Module, arg: serde_json::Value) -> StartCommand {
-        // TODO: compress psl
         StartCommand {
             verbose,
             dns_config,
@@ -77,6 +76,13 @@ impl Supervisor {
         Ok(())
     }
 
+    pub fn send_struct<T: serde::Serialize>(&mut self, value: T, tx: &EventSender) {
+        let value = serde_json::to_value(value).expect("Failed to serialize reply");
+        if let Err(_) = self.send(&value) {
+            tx.send(Event2::Log(LogEvent::Error("Failed to send to child".into())));
+        }
+    }
+
     pub fn recv(&mut self) -> Result<Event> {
         let mut line = String::new();
         let len = self.stdout.read_line(&mut line)?;
@@ -97,15 +103,14 @@ impl Supervisor {
         }
     }
 
-    pub fn send_event_callback(&mut self, event: DatabaseEvent, tx: &EventSender) {
+    pub fn send_event_callback<T: EventWithCallback>(&mut self, event: T, tx: &EventSender)
+        where <T as EventWithCallback>::Payload: serde::Serialize
+    {
         let (tx2, rx2) = mpsc::channel();
-        tx.send(Event2::Database((event, tx2)));
+        tx.send(event.with_callback(tx2));
         let reply = rx2.recv().unwrap();
 
-        let value = serde_json::to_value(reply).expect("Failed to serialize reply");
-        if let Err(_) = self.send(&value) {
-            tx.send(Event2::Log(LogEvent::Error("Failed to send to child".into())));
-        }
+        self.send_struct(reply, tx);
     }
 }
 
@@ -152,9 +157,15 @@ impl Reporter for StdioReporter {
     }
 }
 
-pub fn spawn_module(module: Module, tx: &EventSender, arg: serde_json::Value, verbose: u64) -> Result<()> {
+pub fn spawn_module(module: Module, tx: &EventSender, arg: serde_json::Value, verbose: u64, has_stdin: bool) -> Result<()> {
     let dns_config = Resolver::from_system()?;
 
+    let mut reader = if has_stdin {
+        Some(BufReader::new(stdin()))
+    } else {
+        None
+    };
+
     let mut supervisor = Supervisor::setup(&module)?;
     supervisor.send_start(&StartCommand::new(verbose, dns_config, module, arg))?;
 
@@ -162,6 +173,7 @@ pub fn spawn_module(module: Module, tx: &EventSender, arg: serde_json::Value, ve
         match supervisor.recv()? {
             Event::Log(event) => tx.send(Event2::Log(event)),
             Event::Database(object) => supervisor.send_event_callback(object, &tx),
+            Event::Stdio(object) => object.apply(&mut supervisor, tx, &mut reader),
             Event::Exit(event) => {
                 if let ExitEvent::Err(err) = event {
                     tx.send(Event2::Log(LogEvent::Error(err)));
diff --git a/src/main.rs b/src/main.rs
index e70189aa..f9bc60a6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -42,7 +42,7 @@ fn run_run(gargs: &Args, args: &args::Run, config: Config) -> Result<()> {
         bail!("At least one module or file need to be provided");
     }
 
-    cmd::run_cmd::execute(&mut rl, args.threads, args.verbose)
+    cmd::run_cmd::execute(&mut rl, args.threads, args.verbose, args.stdin)
 }
 
 fn run_sandbox() -> Result<()> {
diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs
index fa748642..6d6bdb01 100644
--- a/src/runtime/mod.rs
+++ b/src/runtime/mod.rs
@@ -18,6 +18,7 @@ import_fns!(print);
 import_fns!(psl);
 import_fns!(regex);
 import_fns!(sleep);
+import_fns!(stdio);
 import_fns!(url);
 import_fns!(utf8);
 import_fns!(x509);
diff --git a/src/runtime/stdio.rs b/src/runtime/stdio.rs
new file mode 100644
index 00000000..1fb4fe17
--- /dev/null
+++ b/src/runtime/stdio.rs
@@ -0,0 +1,13 @@
+use errors::*;
+
+use engine::ctx::State;
+use hlua;
+use std::sync::Arc;
+
+
+pub fn stdin_readline(lua: &mut hlua::Lua, state: Arc<State>) {
+    lua.set("stdin_readline", hlua::function0(move || -> Result<Option<String>> {
+        state.stdin_readline()
+            .map_err(|e| state.set_error(e))
+    }))
+}
diff --git a/src/worker.rs b/src/worker.rs
index ff750217..5a097830 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -3,6 +3,7 @@ use errors::*;
 use channel;
 use db::{Database, DbChange, Family};
 use engine::{self, Module};
+use engine::isolation::Supervisor;
 use models::*;
 use serde_json;
 use shell::Readline;
@@ -10,6 +11,7 @@ use std::result;
 use std::sync::{mpsc, Arc, Mutex};
 use std::time::Duration;
 use std::thread;
+use std::io::{Stdin, BufRead, BufReader};
 use term::{Spinner, StackedSpinners, SpinLogger};
 use threadpool::ThreadPool;
 
@@ -20,6 +22,7 @@ type DbSender = mpsc::Sender<result::Result<Option<i32>, String>>;
 pub enum Event {
     Log(LogEvent),
     Database(DatabaseEvent),
+    Stdio(StdioEvent),
     Exit(ExitEvent),
 }
 
@@ -65,6 +68,12 @@ impl EventSender {
     }
 }
 
+pub trait EventWithCallback {
+    type Payload;
+
+    fn with_callback(self, tx: mpsc::Sender<result::Result<Self::Payload, String>>) -> Event2;
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub enum ExitEvent {
     Ok,
@@ -97,6 +106,14 @@ pub enum DatabaseEvent {
     Update((String, Update)),
 }
 
+impl EventWithCallback for DatabaseEvent {
+    type Payload = Option<i32>;
+
+    fn with_callback(self, tx: mpsc::Sender<result::Result<Self::Payload, String>>) -> Event2 {
+        Event2::Database((self, tx))
+    }
+}
+
 impl DatabaseEvent {
     pub fn apply<T: SpinLogger>(self, tx: DbSender, spinner: &mut T, db: &Database, verbose: u64) {
         match self {
@@ -164,7 +181,35 @@ impl DatabaseEvent {
     }
 }
 
-pub fn spawn(rl: &mut Readline, module: &Module, args: Vec<(serde_json::Value, Option<String>)>, threads: usize, verbose: u64) {
+#[derive(Debug, Serialize, Deserialize)]
+pub struct StdioEvent {
+}
+
+impl StdioEvent {
+    fn read(reader: &mut Option<BufReader<Stdin>>) -> Result<Option<String>> {
+        if let Some(ref mut reader) = reader {
+            let mut line = String::new();
+            let len = reader.read_line(&mut line)?;
+            debug!("stdin: {:?}", line);
+
+            if len > 0 {
+                Ok(Some(line))
+            } else {
+                Ok(None)
+            }
+        } else {
+            bail!("stdin is unavailable");
+        }
+    }
+
+    pub fn apply(self, supervisor: &mut Supervisor, tx: &EventSender, reader: &mut Option<BufReader<Stdin>>) {
+        let reply = Self::read(reader)
+            .map_err(|e| e.to_string());
+        supervisor.send_struct(reply, tx);
+    }
+}
+
+pub fn spawn(rl: &mut Readline, module: &Module, args: Vec<(serde_json::Value, Option<String>)>, threads: usize, verbose: u64, has_stdin: bool) {
     // This function hangs if args is empty, so return early if that's the case
     if args.is_empty() {
         return;
@@ -194,7 +239,7 @@ pub fn spawn(rl: &mut Readline, module: &Module, args: Vec<(serde_json::Value, O
             }
 
             tx.send(Event2::Start);
-            let event = match engine::isolation::spawn_module(module, &tx, arg, verbose) {
+            let event = match engine::isolation::spawn_module(module, &tx, arg, verbose, has_stdin) {
                 Ok(_) => ExitEvent::Ok,
                 Err(err) => ExitEvent::Err(err.to_string()),
             };
@@ -262,6 +307,7 @@ pub fn spawn_fn<F, T>(label: &str, f: F, clear: bool) -> Result<T>
                 recv(rx) -> msg => match msg.ok() {
                     Some(Event::Log(log)) => log.apply(&mut *spinner),
                     Some(Event::Database(_)) => (),
+                    Some(Event::Stdio(_)) => (),
                     // TODO: refactor
                     Some(Event::Exit(ExitEvent::Ok)) => break,
                     Some(Event::Exit(ExitEvent::Err(error))) => spinner.error(&error),


Test patch: