From 819eb7d36233124d935968c263adbb967b6b221e Mon Sep 17 00:00:00 2001 From: Shimun Date: Sun, 31 Mar 2019 18:04:01 +0200 Subject: [PATCH] integrate old code --- wg-event-gen/src/controller.rs | 6 +++--- wg-event-gen/src/gen.rs | 22 +++++++++++---------- wg-event-gen/src/listener.rs | 15 ++++++++------- wg-event-gen/src/main.rs | 35 ++++++++++++++++++++++++++++++++-- wg-event-gen/src/model.rs | 30 ++++++++++++++++++++++------- wg-event-gen/src/opts.rs | 3 +++ 6 files changed, 82 insertions(+), 29 deletions(-) diff --git a/wg-event-gen/src/controller.rs b/wg-event-gen/src/controller.rs index 866f198..fdb6d88 100644 --- a/wg-event-gen/src/controller.rs +++ b/wg-event-gen/src/controller.rs @@ -5,7 +5,7 @@ use std::io::{BufRead, BufReader, Error, ErrorKind, Result, Write}; use std::net::{IpAddr, SocketAddr}; use std::os::unix::net::UnixStream; use std::path::PathBuf; -use std::time::Duration; +use std::time::{Duration, Instant}; pub struct Userspace(PathBuf); @@ -49,8 +49,8 @@ impl WireguardController for Userspace { }; match key.as_ref() { "" => { - //Empty line means end of data - build_peer(&mut peer, builder); //TODO: handle possible actual error case + //Empty line means end of data + build_peer(&mut peer, builder); //TODO: handle possible actual error case } "public_key" => { add_key(&mut peer, parse_err(ECCKey::from_base64(value))?)?; diff --git a/wg-event-gen/src/gen.rs b/wg-event-gen/src/gen.rs index 1e24d6a..1663fb5 100644 --- a/wg-event-gen/src/gen.rs +++ b/wg-event-gen/src/gen.rs @@ -4,8 +4,8 @@ use std::collections::{HashMap, HashSet}; use std::time; pub(crate) fn gen_events( - state: &HashMap, - prev: &HashMap, + state: &HashMap, + prev: &HashMap, listeners: &Vec>, timeout: time::Duration, poll_interval: time::Duration, @@ -13,20 +13,22 @@ pub(crate) fn gen_events( let side_by_side = { state .keys() - .map(String::as_ref) - .chain(prev.keys().map(String::as_ref)) - .collect::>() + .chain(prev.keys()) + .collect::>() .iter() .map(|p| (*p, (prev.get(*p), state.get(*p)))) - .collect::, Option<&Peer>)>>() + .collect::, Option<&Peer>)>>() }; for (_id, (prev, cur)) in side_by_side { match (prev, cur) { (Some(prev), Some(cur)) => { - let timedout = |peer: &Peer| match peer.last_handshake_rel() { - Some(shake) if shake > timeout && shake + poll_interval < timeout => true, - Some(_) => false, - _ => true, + let timedout = |peer: &Peer| { + peer.last_handshake + .map(|shake| { + shake.elapsed().unwrap() > timeout + || shake.elapsed().unwrap() + poll_interval < timeout + }) + .unwrap_or(true) }; if let (Some(prev_addr), Some(cur_addr)) = (prev.endpoint, cur.endpoint) { diff --git a/wg-event-gen/src/listener.rs b/wg-event-gen/src/listener.rs index 6faa331..7b0f994 100644 --- a/wg-event-gen/src/listener.rs +++ b/wg-event-gen/src/listener.rs @@ -3,6 +3,7 @@ use std::net::SocketAddr; use std::path::PathBuf; use std::process::Command; use std::thread; +use std::time::SystemTime; pub trait EventListener { fn added<'a>(&self, peer: &'a Peer) { @@ -50,25 +51,25 @@ pub struct LogListener; impl EventListener for LogListener { fn connected<'a>(&self, peer: &'a Peer) { - println!("{} connected!", peer.public_key); + println!("{} connected!", peer.key); } fn disconnected<'a>(&self, peer: &'a Peer) { - println!("{} disconnected!", peer.public_key); + println!("{} disconnected!", peer.key); } fn added<'a>(&self, peer: &'a Peer) { - println!("{} added!", peer.public_key); + println!("{} added!", peer.key); } fn removed<'a>(&self, peer: &'a Peer) { - println!("{} removed!", peer.public_key); + println!("{} removed!", peer.key); } fn roaming<'a>(&self, peer: &'a Peer, previous_addr: SocketAddr) { println!( "{} roamed {} -> {}!", - peer.public_key, + peer.key, previous_addr, peer.endpoint.unwrap() ); @@ -87,7 +88,7 @@ impl ScriptListener { fn peer_props<'a>(&self, peer: &'a Peer) -> String { format!( "{id} {allowed_ips} {endpoint} {last_handshake} {persistent_keepalive} {traffic}", - id = peer.public_key, + id = peer.key, allowed_ips = peer .allowed_ips .iter() @@ -100,7 +101,7 @@ impl ScriptListener { .unwrap_or("0".to_owned()), last_handshake = peer .last_handshake - .map(|s| s.as_secs() as i64) + .map(|s| s.duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs() as i64) .unwrap_or(-1), persistent_keepalive = peer .persistent_keepalive diff --git a/wg-event-gen/src/main.rs b/wg-event-gen/src/main.rs index 6118bcf..4712636 100644 --- a/wg-event-gen/src/main.rs +++ b/wg-event-gen/src/main.rs @@ -5,31 +5,51 @@ extern crate structopt; extern crate derive_builder; mod controller; +mod gen; +mod listener; mod model; mod opts; +use crate::gen::gen_events; +use crate::listener::*; +use crate::model::{ECCKey, Peer}; use controller::Userspace; use model::WireguardController; use opts::Opts; +use std::collections::HashMap; use std::io; use std::thread::sleep; use std::time::{Duration, Instant}; use structopt::StructOpt; +fn listeners(opts: &Opts) -> Vec> { + let mut listeners: Vec> = Vec::with_capacity(2); + if let Some(events) = opts.events.clone() { + listeners.push(Box::new(ScriptListener::new(events))) + } + if opts.log { + listeners.push(Box::new(LogListener)); + } + listeners +} + fn main() -> io::Result<()> { let opts = Opts::from_args(); let mut controller: Box = Box::new(Userspace::new(opts.socket.clone())); let interval = Duration::from_millis(opts.poll); let timeout = Duration::from_secs(opts.timeout); + let listeners = listeners(&opts); + println!( "Polling {} every {:?}", opts.socket.to_str().unwrap(), interval ); + let mut peers_last: Option> = None; loop { let now = Instant::now(); let peers = controller.peers()?; - println!("Connected peers:"); + /*println!("Connected peers:"); for peer in peers { let peer = peer?; if peer @@ -37,8 +57,19 @@ fn main() -> io::Result<()> { .map(|h| h.elapsed().unwrap() < timeout) .unwrap_or(false) { - println!("{}", peer); + println!("/\\{:?} {}",(timeout - peer.last_handshake.unwrap().elapsed().unwrap()), peer); } + }*/ + let peers = peers + .map(|peer| peer.map(|peer_ok| (peer_ok.key.clone(), peer_ok))) + .collect::>>()?; + if let Some(ref mut peers_last) = peers_last { + + gen_events(&peers, &peers_last, &listeners, timeout, interval); + + *peers_last = peers; + } else { + peers_last = Some(peers); } let pause = interval - now.elapsed(); dbg!(interval - pause); diff --git a/wg-event-gen/src/model.rs b/wg-event-gen/src/model.rs index 4b40b8a..6d95dde 100644 --- a/wg-event-gen/src/model.rs +++ b/wg-event-gen/src/model.rs @@ -1,7 +1,7 @@ use base64::{decode, encode}; use std::error::Error; use std::fmt; -use std::hash::Hash; +use std::hash::{Hash, Hasher}; use std::io; use std::net::{IpAddr, SocketAddr}; use std::time::Instant; @@ -71,7 +71,7 @@ impl Base64Backed for ECCKey { } impl ECCKey { - pub fn extract_public_key(&self) -> Option { + pub fn public_key(&self) -> Option { //TODO: Determine whether Self is a private key and only the return public part Some(self.clone()) } @@ -95,14 +95,20 @@ impl Base64Backed for SharedKey { } } -#[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)] +#[derive(Debug, Builder, PartialEq, Eq, Clone)] pub struct Interface { pub key: ECCKey, pub port: usize, pub fwmark: Option, } -#[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)] +impl Hash for Interface { + fn hash(&self, state: &mut H) { + self.key.public_key().hash(state); + } +} + +#[derive(Debug, Builder, PartialEq, Eq, Clone)] pub struct Peer { pub key: ECCKey, #[builder(default = "None")] @@ -121,10 +127,18 @@ pub struct Peer { pub parsed: Instant, } +impl Hash for Peer { + fn hash(&self, state: &mut H) { + self.key.hash(state); + } +} + impl fmt::Display for Peer { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn dis_opt<'a, T: fmt::Display + 'a>(opt: &Option) -> String { - opt.as_ref().map(|s| s.to_string()).unwrap_or(" ".to_string()) + opt.as_ref() + .map(|s| s.to_string()) + .unwrap_or(" ".to_string()) } write!( f, @@ -132,9 +146,11 @@ impl fmt::Display for Peer { self.key, dis_opt(&self.shared_key), dis_opt(&self.endpoint), - self.allowed_ips.iter() + self.allowed_ips + .iter() .map(|(ip, sub)| format!(" {}/{}", ip, sub)) - .collect::>().join(",") + .collect::>() + .join(",") ) } } diff --git a/wg-event-gen/src/opts.rs b/wg-event-gen/src/opts.rs index 69c2020..55c4182 100644 --- a/wg-event-gen/src/opts.rs +++ b/wg-event-gen/src/opts.rs @@ -31,6 +31,9 @@ pub struct Opts { #[structopt(short = "I", long = "ignore-socket-err", env = "WG_IGNORE_SOCKET_ERR")] pub ignore_socket_errors: bool, + #[structopt(short = "l", long = "log")] + pub log: bool, + #[structopt(name = "SOCKET", parse(from_os_str), env = "WG_EVENT_SOCKET")] pub socket: PathBuf, }