From c58c970cad4338afa02ad228d38c5d2068f43ff3 Mon Sep 17 00:00:00 2001 From: shim_ <> Date: Wed, 9 Jan 2019 22:22:49 +0100 Subject: [PATCH] track events --- .dockerignore | 1 + .gitignore | 1 + Dockerfile | 8 ++ wg-event-gen/Cargo.lock | 14 +++ wg-event-gen/Cargo.toml | 8 ++ wg-event-gen/src/main.rs | 214 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 246 insertions(+) create mode 100644 .dockerignore create mode 100644 .gitignore create mode 100644 wg-event-gen/Cargo.lock create mode 100644 wg-event-gen/Cargo.toml create mode 100644 wg-event-gen/src/main.rs diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..c910d58 --- /dev/null +++ b/.dockerignore @@ -0,0 +1 @@ +*/target diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c910d58 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +*/target diff --git a/Dockerfile b/Dockerfile index a626122..8e41d54 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,3 +1,11 @@ +FROM rust:1.31.1 AS eventbuild + +COPY wg-event-gen/ /build + +WORKDIR /build + +RUN cargo build --release + FROM golang AS build COPY wireguard-go /go/src/wireguard diff --git a/wg-event-gen/Cargo.lock b/wg-event-gen/Cargo.lock new file mode 100644 index 0000000..b1bf710 --- /dev/null +++ b/wg-event-gen/Cargo.lock @@ -0,0 +1,14 @@ +[[package]] +name = "strum" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" + +[[package]] +name = "wg-event-gen" +version = "0.1.0" +dependencies = [ + "strum 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[metadata] +"checksum strum 0.13.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f6b3fc98c482ff9bb37a6db6a6491218c4c82bec368bd5682033e5b96b969143" diff --git a/wg-event-gen/Cargo.toml b/wg-event-gen/Cargo.toml new file mode 100644 index 0000000..f27cfa5 --- /dev/null +++ b/wg-event-gen/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "wg-event-gen" +version = "0.1.0" +authors = ["shimun "] +edition = "2018" + +[dependencies] +strum = "0.13.0" diff --git a/wg-event-gen/src/main.rs b/wg-event-gen/src/main.rs new file mode 100644 index 0000000..e09ea50 --- /dev/null +++ b/wg-event-gen/src/main.rs @@ -0,0 +1,214 @@ +use std::collections::HashMap; +use std::env; +use std::io::prelude::*; +use std::io::{BufRead, BufReader, Error, ErrorKind, Result}; +use std::net::SocketAddr; +use std::os::unix::net::UnixStream; +use std::path::PathBuf; +use std::{thread, time}; + +pub type KV = (String, String); + +#[derive(Debug, PartialEq, Eq, Hash)] +enum State { + Interface(Vec), + Peer(Vec), +} + +impl State { + fn kv(&self) -> &Vec { + match self { + State::Interface(kv) => kv, + State::Peer(kv) => kv, + } + } + + fn kv_mut(&mut self) -> &mut Vec { + match self { + State::Interface(kv) => kv, + State::Peer(kv) => kv, + } + } + + pub fn id<'a>(&'a self) -> Option<&'a String> { + self.kv() + .iter() + .filter(|(key, _)| key == &"private_key" || key == &"public_key") + .map(|(_, value)| value) + .next() + } + + pub fn addr(&self) -> Option { + self.kv() + .iter() + .filter(|(key, _)| key == &"endpoint") + .map(|(_, value)| value.parse::().unwrap()) + .next() + } + + pub fn last_handshake(&self) -> Option { + self.kv() + .iter() + .filter(|(key, _)| key == &"last_handshake_time_nsec") + .map(|(_, value)| value.parse::().unwrap()) + .next() + } + + pub fn push(&mut self, key: String, value: String) { + self.kv_mut().push((key, value)); + } + + pub fn delta(&self, other: Self) -> Vec { + let kv = self.kv(); + other + .kv() + .iter() + .filter(|pair| !kv.contains(pair)) + .map(|p| p.clone()) + .collect::>() + } +} + +struct Socket { + pub path: PathBuf, +} + +impl Socket { + pub fn get(&self) -> Result> { + let mut stream = UnixStream::connect(&self.path)?; + stream.write_all(b"get=1\n")?; + let mut state: Vec = vec![]; + let mut cur = State::Interface(Vec::with_capacity(0)); + for line in BufReader::new(stream).lines() { + let line = line?; + let mut iter = line.chars(); + let key = iter.by_ref().take_while(|c| c != &'=').collect::(); + let value = iter.collect::(); + match key.as_ref() { + "errno" if value != "0" => Err(Error::new( + ErrorKind::Other, + format!("Socket said error: {}", value), + ))?, + "public_key" | "private_key" => { + state.push(cur); + cur = if key == "private_key" { + State::Interface(Vec::with_capacity(3)) + } else { + State::Peer(Vec::with_capacity(5)) + }; + cur.push(key, value); + } + _ => cur.push(key, value), + } + } + Ok(state) + } + + pub fn get_by_id(&self) -> Result> { + let state = self.get()?; + let mut ided = HashMap::new(); + for s in state { + if let Some(id) = s.id() { + ided.insert(id.clone(), s); + } + } + Ok(ided) + } +} + +trait EventListener { + fn added<'a>(&self, peer: &'a State) { + self.connected(peer); + } + + fn connected<'a>(&self, peer: &'a State); + + fn disconnected<'a>(&self, peer: &'a State); + + fn removed<'a>(&self, peer: &'a State) { + self.disconnected(peer) + } + + fn roaming<'a>(&self, peer: &'a State, previous_addr: SocketAddr); +} + +struct LogListener; + +impl EventListener for LogListener { + fn connected<'a>(&self, peer: &'a State) { + println!("{} connected!", peer.id().unwrap()); + } + + fn disconnected<'a>(&self, peer: &'a State) { + println!("{} disconnected!", peer.id().unwrap()); + } + + fn roaming<'a>(&self, peer: &'a State, previous_addr: SocketAddr) { + println!( + "{} roamed {} -> {}!", + peer.id().unwrap(), + previous_addr, + peer.addr().unwrap() + ); + } +} + +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 = args + .next() + .map(|i| { + i.parse::() + .expect("[interval] has to be a positive int") + }) + .unwrap_or(1000); + let listeners = vec![LogListener]; + + let timeout: u64 = 3 * 1000; + + if let Some(path) = path { + let sock = Socket { path }; + let mut prev_state: Option> = None; + loop { + let state = match sock.get_by_id() { + Ok(state) => state, + Err(err) => { + eprintln!("Failed to read from socket: {}", err); + continue; + } + }; + if let Some(prev_state) = prev_state { + for (peer, state) in state.iter() { + if let Some(p_state) = prev_state.get(peer) { + if let (Some(addr), Some(p_addr)) = (state.addr(), p_state.addr()) { + if addr != p_addr { + listeners.iter().for_each(|l| l.roaming(state, p_addr)); + } + } + } else { + listeners.iter().for_each(|l| l.connected(state)); + } + if let Some(shake) = state.last_handshake() { + if (shake / 1000) > timeout && shake / 1000 < timeout + interval { + listeners.iter().for_each(|l| l.disconnected(state)); + } + } + } + prev_state + .iter() + .filter(|(k, _)| !state.contains_key(k.clone())) + .for_each(|(_, state)| listeners.iter().for_each(|l| l.disconnected(state))); + } + state.keys().for_each(|k| print!("{}, ", k)); println!(""); + prev_state = Some(state); + thread::sleep(time::Duration::from_millis(interval)); + } + } else { + println!(" does not exist"); + } +}