diff --git a/wg-event-gen/src/controller.rs b/wg-event-gen/src/controller.rs index 445bff6..866f198 100644 --- a/wg-event-gen/src/controller.rs +++ b/wg-event-gen/src/controller.rs @@ -1,4 +1,6 @@ -use crate::model::{Base64Backed, ECCKey, Interface, Peer, PeerBuilder, SharedKey, WireguardController}; +use crate::model::{ + Base64Backed, ECCKey, Interface, Peer, PeerBuilder, SharedKey, WireguardController, +}; use std::io::{BufRead, BufReader, Error, ErrorKind, Result, Write}; use std::net::{IpAddr, SocketAddr}; use std::os::unix::net::UnixStream; @@ -46,6 +48,10 @@ impl WireguardController for Userspace { Ok(()) }; match key.as_ref() { + "" => { + //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))?)?; } @@ -64,9 +70,15 @@ impl WireguardController for Userspace { "last_handshake_time_nsec" => { builder.add_last_handshake(Duration::from_nanos(value_as_num()?.into())); } - "persistent_keepalive" => { + "persistent_keepalive_interval" => { builder.persistent_keepalive(Some(Duration::from_secs(value_as_num()?.into()))); } + "rx_bytes" => { + builder.add_traffic((parse_err(value.parse::())?, 0)); + } + "tx_bytes" => { + builder.add_traffic((0, (parse_err(value.parse::())?))); + } "allowed_ip" => { let mut parts = value.split("/").into_iter(); let net = match ( @@ -89,7 +101,11 @@ impl WireguardController for Userspace { format!("Returned error code: {}", code), ))?, }, - _ => Err(Error::new(ErrorKind::InvalidData, "Unknown key"))?, + "listen_port" | "fwmark" | "private_key" => (), //Ignore for now + _ => Err(Error::new( + ErrorKind::InvalidData, + ["Unknown key: \"", &key, "\""].join(""), + ))?, } Ok(peer) } @@ -100,19 +116,19 @@ impl WireguardController for Userspace { match build_peer(builder, line) { Ok(Some(value)) => Some(Some(Ok(value))), Err(err) => { - eprintln!("{:?}", err); - None - }, //TODO: propagate + eprintln!("{:?}", err); + None + } //TODO: propagate _ => Some(None), } }) .flatten(); Ok(Box::new(peers)) } - + fn interface(&mut self) -> Result { let mut stream = UnixStream::connect(&self.0)?; - stream.write_all(b"get=1\n")?; + stream.write_all(b"get=1\n")?; unimplemented!("TODO: return iface") } diff --git a/wg-event-gen/src/main.rs b/wg-event-gen/src/main.rs index 37e5765..087019c 100644 --- a/wg-event-gen/src/main.rs +++ b/wg-event-gen/src/main.rs @@ -8,28 +8,42 @@ mod controller; mod model; mod opts; -use controller::{Userspace}; +use controller::Userspace; use model::WireguardController; use opts::Opts; use std::io; -use structopt::StructOpt; use std::thread::sleep; -use std::time::{Duration,Instant}; +use std::time::{Duration, Instant}; +use structopt::StructOpt; 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); - println!("Polling {} every {:?}", opts.socket.to_str().unwrap(), interval); + let timeout = Duration::from_secs(opts.timeout); + println!( + "Polling {} every {:?}", + opts.socket.to_str().unwrap(), + interval + ); loop { - let now = Instant::now(); - let peers = controller.peers()?; - for peer in peers { - dbg!(peer?); - } - let pause = interval - now.elapsed(); - dbg!(pause); - sleep(if pause > interval/2 { pause } else {interval}); - } + let now = Instant::now(); + let peers = controller.peers()?; + println!("Connected peers:"); + for peer in peers { + let peer = peer?; + dbg!(peer.last_handshake); + if peer.last_handshake.map(|h| h < timeout).unwrap_or(false) { + println!("{:?}", peer.key); + } + } + let pause = interval - now.elapsed(); + dbg!(pause); + sleep(if pause > interval / 2 { + pause + } else { + interval + }); + } Ok(()) } diff --git a/wg-event-gen/src/model.rs b/wg-event-gen/src/model.rs index 8e3b732..d7a0c63 100644 --- a/wg-event-gen/src/model.rs +++ b/wg-event-gen/src/model.rs @@ -1,19 +1,22 @@ use base64::{decode, encode}; use std::error::Error; +use std::hash::Hash; use std::io; use std::net::{IpAddr, SocketAddr}; use std::time::Duration; use std::time::Instant; -#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)] +const KEY_SIZE: usize = 48; //TODO: use VEC instead of array + +#[derive(Debug, PartialEq, Eq, Hash, Clone)] pub enum ECCKey { - PublicKey([u8; 32]), - PrivateKey([u8; 32]), + PublicKey(Vec), + PrivateKey(Vec), } pub trait Base64Backed { - fn from_bytes(bytes: [u8; 32]) -> Self; - fn bytes(&self) -> &[u8; 32]; + fn from_bytes(bytes: Vec) -> Self; + fn bytes(&self) -> &Vec; fn from_base64>(key: I) -> io::Result where Self: Sized, @@ -29,12 +32,17 @@ pub trait Base64Backed { }; /*.map_err(|err| { })?;*/ - let mut bytes = [0; 32]; - if key.len() != 32 { - return Err(io::Error::new(io::ErrorKind::Other, "Mismatched key size")); + if key.len() != KEY_SIZE { + return Err(io::Error::new( + io::ErrorKind::Other, + format!( + "Mismatched key size. Expected: {}, Got {}", + KEY_SIZE, + key.len() + ), + )); } - bytes.copy_from_slice(&key); - Ok(Self::from_bytes(bytes)) + Ok(Self::from_bytes(key)) } fn as_base64(&self) -> io::Result { @@ -43,14 +51,14 @@ pub trait Base64Backed { } impl Base64Backed for ECCKey { - fn bytes(&self) -> &[u8; 32] { + fn bytes(&self) -> &Vec { match self { ECCKey::PublicKey(bytes) => &bytes, ECCKey::PrivateKey(bytes) => &bytes, } } - fn from_bytes(bytes: [u8; 32]) -> ECCKey { + fn from_bytes(bytes: Vec) -> ECCKey { ECCKey::PublicKey(bytes) } } @@ -62,40 +70,42 @@ impl ECCKey { } } -#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)] -pub struct SharedKey([u8; 32]); +#[derive(Debug, PartialEq, Eq, Hash, Clone)] +pub struct SharedKey(Vec); impl Base64Backed for SharedKey { - fn bytes(&self) -> &[u8; 32] { + fn bytes(&self) -> &Vec { &self.0 } - fn from_bytes(bytes: [u8; 32]) -> SharedKey { + fn from_bytes(bytes: Vec) -> SharedKey { SharedKey(bytes) } } #[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)] pub struct Interface { - key: ECCKey, - port: usize, - fwmark: Option + pub key: ECCKey, + pub port: usize, + pub fwmark: Option, } #[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)] pub struct Peer { - key: ECCKey, + pub key: ECCKey, #[builder(default = "None")] - shared_key: Option, + pub shared_key: Option, #[builder(default = "None")] - endpoint: Option, - allowed_ips: Vec<(IpAddr, u8)>, - last_handshake: Option, + pub endpoint: Option, + #[builder(default = "Vec::new()")] + pub allowed_ips: Vec<(IpAddr, u8)>, #[builder(default = "None")] - persistent_keepalive: Option, + pub last_handshake: Option, + #[builder(default = "None")] + pub persistent_keepalive: Option, #[builder(default = "(0u64,0u64)")] - traffic: (u64, u64), + pub traffic: (u64, u64), #[builder(default = "Instant::now()")] - parsed: Instant, + pub parsed: Instant, } impl PeerBuilder { @@ -128,6 +138,15 @@ impl PeerBuilder { .map(|shake| shake.map(|shake| shake + d)); } } + + pub fn add_traffic(&mut self, txrx: (u64, u64)) { + if let Some(ref mut traffic) = &mut self.traffic { + traffic.0 += txrx.0; + traffic.1 += txrx.1; + } else { + self.traffic = Some(txrx); + } + } } pub trait WireguardController {