Merge branch 'structop'

wg-event-gen/src/gen.rs

@@ -17,7 +17,7 @@ pub(crate) fn gen_events(
             .chain(prev.keys().map(String::as_ref))
             .collect::<HashSet<&str>>()
             .iter()
-            .map(|p| (p.to_owned(), (prev.get(*p), state.get(*p))))
+            .map(|p| (*p, (prev.get(*p), state.get(*p))))
             .collect::<HashMap<&str, (Option<&Peer>, Option<&Peer>)>>()
     };
     for (_id, (prev, cur)) in side_by_side {
@@ -29,6 +29,12 @@ pub(crate) fn gen_events(
                     _ => true,
                 };
 
+                if let (Some(prev_addr), Some(cur_addr)) = (prev.endpoint, cur.endpoint) {
+                    if prev_addr != cur_addr {
+                        listeners.roaming(&cur, prev_addr);
+                    }
+                }
+
                 if !timedout(&prev) && timedout(&cur) {
                     listeners.disconnected(&cur);
                     continue;
@@ -37,12 +43,6 @@ pub(crate) fn gen_events(
                 if timedout(&prev) && !timedout(&cur) {
                     listeners.connected(&cur);
                 }
-
-                if prev.endpoint != cur.endpoint {
-                    if let (Some(prev_addr), Some(_)) = (prev.endpoint, cur.endpoint) {
-                        listeners.roaming(&cur, prev_addr);
-                    }
-                }
             }
             (None, Some(cur)) => listeners.added(&cur),
             (Some(prev), None) => listeners.removed(&prev),

wg-event-gen/src/main.rs

@@ -1,10 +1,15 @@
+#[macro_use]
+extern crate structopt;
+
 mod gen;
 mod listener;
+mod opts;
 
 use listener::*;
 
 use base64;
 use hex;
+use opts::Opts;
 use std::collections::HashMap;
 use std::env;
 use std::fmt;
@@ -15,6 +20,8 @@ use std::os::unix::net::UnixStream;
 use std::path::PathBuf;
 use std::thread;
 use std::time::Duration;
+use structopt::StructOpt;
+
 use time;
 
 pub type KV = (String, String);
@@ -215,58 +222,33 @@ impl Socket {
 }
 
 fn main() {
-    let mut args = env::args();
-    args.next(); //Ignore program name
-    let path = args
-        .next()
-        .map(PathBuf::from)
-        .filter(|p| p.as_path().exists());
-    let interval = Duration::from_millis(
-        args.next()
-            .map(|i| {
-                i.parse::<u64>()
-                    .expect("[interval] has to be a positive int")
-            })
-            .unwrap_or(1000),
-    );
+    let opts = Opts::from_args();
+
+    let timeout = Duration::from_secs(opts.timeout);
+    let interval = Duration::from_secs(opts.poll);
+    let events = opts.events;
+    let path = opts.socket;
+
     let mut listeners: Vec<Box<EventListener>> = vec![Box::new(LogListener)];
 
-    let events: PathBuf = "/etc/wireguard/events.sh".into();
-
-    if events.exists() {
+    if let Some(events) = events {
         listeners.push(Box::new(ScriptListener::new(events)))
     }
 
-    let timeout = env::vars()
-        .collect::<HashMap<String, String>>()
-        .get("WG_EVENT_GEN_TIMEOUT")
-        .map(|timeout| {
-            Duration::from_secs(
-                timeout
-                    .parse::<u64>()
-                    .expect(&format!("Can't parse {} as timeout", timeout)),
-            )
-        })
-        .unwrap_or(Duration::from_secs(30));
-
-    if let Some(path) = path {
-        let sock = Socket { path };
-        let mut prev_state: Option<HashMap<String, Peer>> = None;
-        loop {
-            let state = match sock.get_peers() {
-                Ok(state) => state,
-                Err(err) => {
-                    eprintln!("Failed to read from socket: {}", err);
-                    continue;
-                }
-            };
-            if let Some(prev_state) = prev_state {
-                gen::gen_events(&state, &prev_state, &listeners, timeout, interval);
+    let sock = Socket { path };
+    let mut prev_state: Option<HashMap<String, Peer>> = None;
+    loop {
+        let state = match sock.get_peers() {
+            Ok(state) => state,
+            Err(err) => {
+                eprintln!("Failed to read from socket: {}", err);
+                continue;
             }
-            prev_state = Some(state);
-            thread::sleep(interval);
+        };
+        if let Some(prev_state) = prev_state {
+            gen::gen_events(&state, &prev_state, &listeners, timeout, interval);
         }
-    } else {
-        println!("<path> does not exist");
+        prev_state = Some(state);
+        thread::sleep(interval);
     }
 }

wg-event-gen/src/opts.rs

@@ -0,0 +1,33 @@
+use std::path::PathBuf;
+use structopt::StructOpt;
+
+#[derive(StructOpt, Debug)]
+#[structopt(name = "event-gen")]
+pub struct Opts {
+    #[structopt(
+        short = "t",
+        long = "timeout",
+        default_value = "120",
+        env = "WG_EVENT_TIMEOUT"
+    )]
+    pub timeout: u64,
+
+    #[structopt(
+        short = "p",
+        long = "poll-interval",
+        default_value = "3000",
+        env = "WG_EVENT_INTERVAL"
+    )]
+    pub poll: u64,
+
+    #[structopt(
+        short = "e",
+        long = "event-handler",
+        parse(from_os_str),
+        env = "WG_EVENT_HANDLER"
+    )]
+    pub events: Option<PathBuf>,
+
+    #[structopt(name = "SOCKET", parse(from_os_str), env = "WG_EVENT_SOCKET")]
+    pub socket: PathBuf,
+}