Merge pull request #187 from kpcyrd/add-stdin

Allow streaming some structs from stdin into db
This commit is contained in:
kpcyrd
2020-11-01 20:11:59 +00:00
committed by GitHub
2 changed files with 166 additions and 54 deletions

66
Cargo.lock generated
View File

@@ -2,9 +2,9 @@
# It is not intended for manual editing.
[[package]]
name = "addr2line"
version = "0.13.0"
version = "0.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b6a2d3371669ab3ca9797670853d61402b03d0b4b9ebf33d677dfa720203072"
checksum = "7c0929d69e78dd9bf5408269919fcbcaeb2e35e5d43e5815517cdc6a8e11a423"
dependencies = [
"gimli",
]
@@ -103,9 +103,9 @@ checksum = "cdb031dd78e28731d87d56cc8ffef4a8f36ca26c38fe2de700543e627f8a464a"
[[package]]
name = "backtrace"
version = "0.3.53"
version = "0.3.54"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "707b586e0e2f247cbde68cdd2c3ce69ea7b7be43e1c5b426e37c9319c4b9838e"
checksum = "2baad346b2d4e94a24347adeee9c7a93f412ee94b9cc26e5b59dea23848e9f28"
dependencies = [
"addr2line",
"cfg-if 1.0.0",
@@ -804,29 +804,27 @@ dependencies = [
[[package]]
name = "data-encoding"
version = "2.3.0"
version = "2.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4d0e2d24e5ee3b23a01de38eefdcd978907890701f08ffffd4cb457ca4ee8d6"
checksum = "993a608597367c6377b258c25d7120740f00ed23a2252b729b1932dd7866f908"
[[package]]
name = "data-encoding-macro"
version = "0.1.8"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "de6489dde5128f5ab2f71f88f8807a237cecf08d96dc7ca4be64e0730dc7d961"
checksum = "2749f3e10bdb761d61b9486511e57bc613a0e8b5b630f299cad65734fc0f794d"
dependencies = [
"data-encoding",
"data-encoding-macro-internal",
"proc-macro-hack",
]
[[package]]
name = "data-encoding-macro-internal"
version = "0.1.8"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8d2d6daefd5f1d4b74a891a5d2ab7dccba028d423107c074232a0c5dc0d40a9e"
checksum = "f0f83e699727abca3c56e187945f303389590305ab2f0185ea445aa66e8d5f2a"
dependencies = [
"data-encoding",
"proc-macro-hack",
"syn 1.0.48",
]
@@ -1406,9 +1404,9 @@ dependencies = [
[[package]]
name = "gimli"
version = "0.22.0"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aaf91faf136cb47367fa430cd46e37a788775e7fa104f8b4bcb3861dc389b724"
checksum = "f6503fe142514ca4799d4c26297c4248239fe8838d827db6bd6065c6ed29a6ce"
[[package]]
name = "glob"
@@ -2074,9 +2072,9 @@ dependencies = [
[[package]]
name = "memchr"
version = "2.3.3"
version = "2.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3728d817d99e5ac407411fa471ff9800a778d88a24685968b36824eaf4bee400"
checksum = "0ee1c47aaa256ecabcaea351eae4a9b01ef39ed810004e298d2511ed284b1525"
[[package]]
name = "memoffset"
@@ -2375,9 +2373,9 @@ dependencies = [
[[package]]
name = "num-integer"
version = "0.1.43"
version = "0.1.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8d59457e662d541ba17869cf51cf177c0b5f0cbf476c66bdc90bf1edac4f875b"
checksum = "d2cc698a63b549a70bc047073d2949cce27cd1c7b0a4a862d08a8031bc2801db"
dependencies = [
"autocfg 1.0.1",
"num-traits",
@@ -2385,9 +2383,9 @@ dependencies = [
[[package]]
name = "num-iter"
version = "0.1.41"
version = "0.1.42"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7a6e6b7c748f995c4c29c5f5ae0248536e04a5739927c74ec0fa564805094b9f"
checksum = "b2021c8337a54d21aca0d59a92577a029af9431cb59b909b03252b9c164fad59"
dependencies = [
"autocfg 1.0.1",
"num-integer",
@@ -2396,9 +2394,9 @@ dependencies = [
[[package]]
name = "num-rational"
version = "0.3.0"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5b4d7360f362cfb50dde8143501e6940b22f644be75a4cc90b2d81968908138"
checksum = "e5fa6d5f418879385b213d905f7cf5bf4aa553d4c380f0152d1d4f2749186fa9"
dependencies = [
"autocfg 1.0.1",
"num-integer",
@@ -2407,9 +2405,9 @@ dependencies = [
[[package]]
name = "num-traits"
version = "0.2.12"
version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac267bcc07f48ee5f8935ab0d24f316fb722d7a1292e2913f0cc196b29ffd611"
checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290"
dependencies = [
"autocfg 1.0.1",
]
@@ -2444,9 +2442,9 @@ dependencies = [
[[package]]
name = "object"
version = "0.21.1"
version = "0.22.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37fd5004feb2ce328a52b0b3d01dbf4ffff72583493900ed15f22d4111c51693"
checksum = "8d3b63360ec3cb337817c2dbd47ab4a0f170d285d8e5a2064600f3def1402397"
[[package]]
name = "once_cell"
@@ -2860,9 +2858,9 @@ dependencies = [
[[package]]
name = "proc-macro-hack"
version = "0.5.18"
version = "0.5.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "99c605b9a0adc77b7211c6b1f722dcb613d68d66859a44f3d485a6da332b0598"
checksum = "dbf0c48bc1d91375ae5c3cd81e3722dff1abcf81a30960240640d223f59fe0e5"
[[package]]
name = "proc-macro2"
@@ -3150,9 +3148,9 @@ dependencies = [
[[package]]
name = "regex"
version = "1.4.1"
version = "1.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8963b85b8ce3074fecffde43b4b0dded83ce2f367dc8d363afc56679f3ee820b"
checksum = "38cf2c13ed4745de91a5eb834e11c00bcc3709e773173b2ce4c56c9fbde04b9c"
dependencies = [
"aho-corasick",
"memchr",
@@ -3162,9 +3160,9 @@ dependencies = [
[[package]]
name = "regex-syntax"
version = "0.6.20"
version = "0.6.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8cab7a364d15cde1e505267766a2d3c4e22a843e1a601f0fa7564c0f82ced11c"
checksum = "3b181ba2dcf07aaccad5448e8ead58db5b742cf85dfe035e2227f137a539a189"
[[package]]
name = "remove_dir_all"
@@ -3988,9 +3986,9 @@ checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
[[package]]
name = "state"
version = "0.4.1"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7345c971d1ef21ffdbd103a75990a15eb03604fc8b8852ca8cb418ee1a099028"
checksum = "3015a7d0a5fd5105c91c3710d42f9ccf0abfb287d62206484dcc67f9569a6483"
[[package]]
name = "static_assertions"

View File

@@ -2,6 +2,7 @@ use crate::errors::*;
use crate::blobs::Blob;
use crate::cmd::Cmd;
use crate::db::DbChange;
use crate::gfx;
use crate::models::*;
use crate::shell::Shell;
@@ -9,7 +10,9 @@ use structopt::StructOpt;
use structopt::clap::AppSettings;
use crate::utils;
use crate::term;
use std::fmt::Debug;
use std::fs;
use std::io::{self, BufRead};
use std::net;
use std::net::SocketAddr;
use ipnetwork;
@@ -24,6 +27,9 @@ pub struct Args {
/// Do not actually insert into database
#[structopt(short="n", long="dry-run")]
dry_run: bool,
/// Stream structs from stdin line by line
#[structopt(long)]
stdin: bool,
}
#[derive(Debug, StructOpt)]
@@ -74,25 +80,68 @@ pub enum Target {
impl Cmd for Args {
fn run(self, rl: &mut Shell) -> Result<()> {
match self.subcommand {
Target::Domain(args) => args.insert(rl, self.dry_run),
Target::Subdomain(args) => args.insert(rl, self.dry_run),
Target::IpAddr(args) => args.insert(rl, self.dry_run),
Target::Url(args) => args.insert(rl, self.dry_run),
Target::Email(args) => args.insert(rl, self.dry_run),
Target::PhoneNumber(args) => args.insert(rl, self.dry_run),
Target::Device(args) => args.insert(rl, self.dry_run),
Target::Network(args) => args.insert(rl, self.dry_run),
Target::Account(args) => args.insert(rl, self.dry_run),
Target::Breach(args) => args.insert(rl, self.dry_run),
Target::Image(args) => args.insert(rl, self.dry_run),
Target::Netblock(args) => args.insert(rl, self.dry_run),
Target::Port(args) => args.insert(rl, self.dry_run),
Target::CryptoAddr(args) => args.insert(rl, self.dry_run),
if self.stdin {
match self.subcommand {
Target::Domain(_) => stream_stdin::<AddDomain>(rl),
Target::Subdomain(_) => stream_stdin::<AddSubdomain>(rl),
Target::IpAddr(_) => stream_stdin::<AddIpAddr>(rl),
Target::Url(_) => stream_stdin::<AddUrl>(rl),
Target::Email(_) => stream_stdin::<AddEmail>(rl),
Target::PhoneNumber(_) => stream_stdin::<AddPhoneNumber>(rl),
Target::Port(_) => stream_stdin::<AddPort>(rl),
Target::CryptoAddr(_) => stream_stdin::<AddCryptoAddr>(rl),
Target::Device(_) | Target::Network(_) | Target::Account(_) | Target::Breach(_) | Target::Image(_) | Target::Netblock(_) => {
bail!("This struct doesn't support streaming from stdin")
},
}
} else {
match self.subcommand {
Target::Domain(args) => args.insert(rl, self.dry_run),
Target::Subdomain(args) => args.insert(rl, self.dry_run),
Target::IpAddr(args) => args.insert(rl, self.dry_run),
Target::Url(args) => args.insert(rl, self.dry_run),
Target::Email(args) => args.insert(rl, self.dry_run),
Target::PhoneNumber(args) => args.insert(rl, self.dry_run),
Target::Device(args) => args.insert(rl, self.dry_run),
Target::Network(args) => args.insert(rl, self.dry_run),
Target::Account(args) => args.insert(rl, self.dry_run),
Target::Breach(args) => args.insert(rl, self.dry_run),
Target::Image(args) => args.insert(rl, self.dry_run),
Target::Netblock(args) => args.insert(rl, self.dry_run),
Target::Port(args) => args.insert(rl, self.dry_run),
Target::CryptoAddr(args) => args.insert(rl, self.dry_run),
}
}
}
}
fn stream_stdin<T: InsertFromString>(rl: &mut Shell) -> Result<()> {
let stdin = io::stdin();
for line in stdin.lock().lines() {
let line = line?;
match T::from_string(rl, line) {
Ok(insert) => {
debug!("Received {:?}", insert);
let result = rl.db().insert_generic(insert.clone())?;
if let Some((DbChange::Insert, _)) = result {
if let Ok(value) = insert.value(rl.db()) {
let log = format!("Adding {} {:?}", insert.family(), value);
term::success(&log);
}
}
},
Err(err) => term::error(&format!("Error: {:#}", err)),
}
}
Ok(())
}
trait InsertFromString {
// TODO: some implementations run insert_struct instead of just transforming
fn from_string(rl: &mut Shell, x: String) -> Result<Insert>;
}
trait IntoInsert: Sized {
fn into_insert(self, rl: &mut Shell) -> Result<Insert>;
@@ -116,7 +165,12 @@ impl IntoInsert for AddDomain {
Some(domain) => domain,
_ => utils::question("Domain")?,
};
Self::from_string(rl, domain)
}
}
impl InsertFromString for AddDomain {
fn from_string(rl: &mut Shell, domain: String) -> Result<Insert> {
// ensure input is a valid domain
let dns_name = rl.psl()?.parse_dns_name(&domain)
.map_err(|e| format_err!("Failed to parse domain: {}", e))?;
@@ -143,7 +197,12 @@ impl IntoInsert for AddSubdomain {
Some(subdomain) => subdomain,
_ => utils::question("Subdomain")?,
};
Self::from_string(rl, subdomain)
}
}
impl InsertFromString for AddSubdomain {
fn from_string(rl: &mut Shell, subdomain: String) -> Result<Insert> {
let dns_name = rl.psl()?.parse_dns_name(&subdomain)
.map_err(|e| format_err!("Failed to parse dns_name: {}", e))?;
@@ -170,7 +229,8 @@ pub struct AddIpAddr {
}
impl IntoInsert for AddIpAddr {
fn into_insert(self, _rl: &mut Shell) -> Result<Insert> {
fn into_insert(self, rl: &mut Shell) -> Result<Insert> {
// TODO: change this to Option<String> so we don't need to parse back and forth
let ipaddr = match self.ipaddr {
Some(ipaddr) => ipaddr,
_ => {
@@ -178,6 +238,13 @@ impl IntoInsert for AddIpAddr {
ipaddr.parse()?
},
};
Self::from_string(rl, ipaddr.to_string())
}
}
impl InsertFromString for AddIpAddr {
fn from_string(_rl: &mut Shell, ipaddr: String) -> Result<Insert> {
let ipaddr = ipaddr.parse()?;
let family = match ipaddr {
net::IpAddr::V4(_) => "4",
@@ -214,7 +281,12 @@ impl IntoInsert for AddUrl {
Some(url) => url,
_ => utils::question("URL")?,
};
Self::from_string(rl, url)
}
}
impl InsertFromString for AddUrl {
fn from_string(rl: &mut Shell, url: String) -> Result<Insert> {
let parts = url::Url::parse(&url)?;
let subdomain = parts.domain()
.ok_or_else(|| format_err!("url doesn't have a domain host"))?;
@@ -258,12 +330,18 @@ pub struct AddEmail {
}
impl IntoInsert for AddEmail {
fn into_insert(self, _rl: &mut Shell) -> Result<Insert> {
fn into_insert(self, rl: &mut Shell) -> Result<Insert> {
let email = match self.email {
Some(email) => email,
_ => utils::question("Email")?,
};
Self::from_string(rl, email)
}
}
impl InsertFromString for AddEmail {
fn from_string(_rl: &mut Shell, email: String) -> Result<Insert> {
// TODO: consider doing basic validation
Ok(Insert::Email(NewEmail {
value: email,
displayname: None,
@@ -290,6 +368,7 @@ impl IntoInsert for AddPhoneNumber {
},
};
// TODO: consider doing basic validation
Ok(Insert::PhoneNumber(NewPhoneNumber {
value: phonenumber,
name: name,
@@ -307,6 +386,26 @@ impl IntoInsert for AddPhoneNumber {
}
}
impl InsertFromString for AddPhoneNumber {
fn from_string(_rl: &mut Shell, phonenumber: String) -> Result<Insert> {
// TODO: consider doing basic validation
Ok(Insert::PhoneNumber(NewPhoneNumber {
value: phonenumber,
name: None,
valid: None,
last_online: None,
country: None,
carrier: None,
line: None,
is_ported: None,
last_ported: None,
caller_name: None,
caller_type: None,
unscoped: false,
}))
}
}
#[derive(Debug, StructOpt)]
pub struct AddDevice {
mac: Option<String>,
@@ -575,12 +674,24 @@ impl IntoInsert for AddPort {
};
let addr = if let Some(addr) = self.addr {
// TODO: change this to a regular string so we don't parse->fmt->parse
addr
} else {
let addr = utils::question("IP:Port")?;
addr.parse()?
};
Self::from_string(rl, format!("{}/{}", protocol, addr))
}
}
impl InsertFromString for AddPort {
fn from_string(rl: &mut Shell, value: String) -> Result<Insert> {
let idx = value.find('/')
.ok_or_else(|| format_err!("Could not locate `/` delimiter for protocol/ip:port"))?;
let protocol = value[..idx].to_string();
let addr = value[idx+1..].parse::<SocketAddr>()?;
let family = match addr.ip() {
net::IpAddr::V4(_) => "4",
net::IpAddr::V6(_) => "6",
@@ -606,8 +717,6 @@ impl IntoInsert for AddPort {
_ => bail!("IpAddr is out out of scope"),
};
let value = format!("{}/{}", protocol, addr);
Ok(Insert::Port(NewPort {
ip_addr_id,
value,
@@ -631,13 +740,18 @@ pub struct AddCryptoAddr {
}
impl IntoInsert for AddCryptoAddr {
fn into_insert(self, _rl: &mut Shell) -> Result<Insert> {
fn into_insert(self, rl: &mut Shell) -> Result<Insert> {
let address = if let Some(address) = self.address {
address
} else {
utils::question("Address")?
};
Self::from_string(rl, address)
}
}
impl InsertFromString for AddCryptoAddr {
fn from_string(_rl: &mut Shell, address: String) -> Result<Insert> {
Ok(Insert::CryptoAddr(NewCryptoAddr {
value: address,
currency: None,