_ Git - dnsseed-rust/blob - src/reader.rs

   1 use std::sync::Arc;
   2 use std::sync::atomic::Ordering;
   3 use std::io::BufReader;
   4 use std::net::{IpAddr, SocketAddr};
   5 use std::time::Instant;
   6
   7 use tokio::prelude::*;
   8 use tokio::io::{stdin, lines};
   9
  10 use crate::printer::Printer;
  11 use crate::datastore::{Store, AddressState, U64Setting, RegexSetting};
  12 use crate::bgp_client::BGPClient;
  13
  14 use crate::{START_SHUTDOWN, scan_node};
  15
  16 use regex::Regex;
  17
  18 pub fn read(store: &'static Store, printer: &'static Printer, bgp_client: Arc<BGPClient>) {
  19         tokio::spawn(lines(BufReader::new(stdin())).for_each(move |line| {
  20                 macro_rules! err {
  21                         () => { {
  22                                 printer.add_line(format!("Unparsable input: \"{}\"", line), true);
  23                                 return future::ok(());
  24                         } }
  25                 }
  26                 let mut line_iter = line.split(' ');
  27                 macro_rules! get_next_chunk {
  28                         () => { {
  29                                 match line_iter.next() {
  30                                         Some(c) => c,
  31                                         None => err!(),
  32                                 }
  33                         } }
  34                 }
  35                 macro_rules! try_parse_next_chunk {
  36                         ($type: ty) => { {
  37                                 match get_next_chunk!().parse::<$type>() {
  38                                         Ok(res) => res,
  39                                         Err(_) => err!(),
  40                                 }
  41                         } }
  42                 }
  43                 match get_next_chunk!() {
  44                         "c" => store.set_u64(U64Setting::ConnsPerSec, try_parse_next_chunk!(u64)),
  45                         "t" => store.set_u64(U64Setting::RunTimeout, try_parse_next_chunk!(u64)),
  46                         "v" => store.set_u64(U64Setting::MinProtocolVersion, try_parse_next_chunk!(u64)),
  47                         "w" => store.set_u64(U64Setting::WasGoodTimeout, try_parse_next_chunk!(u64)),
  48                         "s" => {
  49                                 if line.len() < 3 || !line.starts_with("s ") {
  50                                         err!();
  51                                 }
  52                                 store.set_regex(RegexSetting::SubverRegex, match line[2..].parse::<Regex>() {
  53                                         Ok(res) => res,
  54                                         Err(_) => err!(),
  55                                 });
  56                         },
  57                         "a" => scan_node(Instant::now(), try_parse_next_chunk!(SocketAddr), true),
  58                         "b" => {
  59                                 let ip = try_parse_next_chunk!(IpAddr);
  60                                 printer.add_line(format!("ASN for {} is {}", ip, bgp_client.get_asn(ip)), false);
  61                         },
  62                         "r" => {
  63                                 match AddressState::from_num(try_parse_next_chunk!(u8)) {
  64                                         Some(state) => store.set_u64(U64Setting::RescanInterval(state), try_parse_next_chunk!(u64)),
  65                                         None => err!(),
  66                                 }
  67                         },
  68                         "q" => {
  69                                 START_SHUTDOWN.store(true, Ordering::SeqCst);
  70                                 return future::err(std::io::Error::new(std::io::ErrorKind::Other, ""));
  71                         },
  72                         _ => err!(),
  73                 }
  74                 future::ok(())
  75         }).then(move |_| {
  76                 printer.add_line("Shutting down...".to_string(), true);
  77                 future::ok(())
  78         }));
  79 }