track events

This commit is contained in:
shim_ 2019-01-09 22:22:49 +01:00
parent f6c1c2771e
commit c58c970cad
6 changed files with 246 additions and 0 deletions

1
.dockerignore Normal file
View File

@ -0,0 +1 @@
*/target

1
.gitignore vendored Normal file
View File

@ -0,0 +1 @@
*/target

View File

@ -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

14
wg-event-gen/Cargo.lock generated Normal file
View File

@ -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"

8
wg-event-gen/Cargo.toml Normal file
View File

@ -0,0 +1,8 @@
[package]
name = "wg-event-gen"
version = "0.1.0"
authors = ["shimun <wg.shimun@shimun.net>"]
edition = "2018"
[dependencies]
strum = "0.13.0"

214
wg-event-gen/src/main.rs Normal file
View File

@ -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<KV>),
Peer(Vec<KV>),
}
impl State {
fn kv(&self) -> &Vec<KV> {
match self {
State::Interface(kv) => kv,
State::Peer(kv) => kv,
}
}
fn kv_mut(&mut self) -> &mut Vec<KV> {
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<SocketAddr> {
self.kv()
.iter()
.filter(|(key, _)| key == &"endpoint")
.map(|(_, value)| value.parse::<SocketAddr>().unwrap())
.next()
}
pub fn last_handshake(&self) -> Option<u64> {
self.kv()
.iter()
.filter(|(key, _)| key == &"last_handshake_time_nsec")
.map(|(_, value)| value.parse::<u64>().unwrap())
.next()
}
pub fn push(&mut self, key: String, value: String) {
self.kv_mut().push((key, value));
}
pub fn delta(&self, other: Self) -> Vec<KV> {
let kv = self.kv();
other
.kv()
.iter()
.filter(|pair| !kv.contains(pair))
.map(|p| p.clone())
.collect::<Vec<KV>>()
}
}
struct Socket {
pub path: PathBuf,
}
impl Socket {
pub fn get(&self) -> Result<Vec<State>> {
let mut stream = UnixStream::connect(&self.path)?;
stream.write_all(b"get=1\n")?;
let mut state: Vec<State> = 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::<String>();
let value = iter.collect::<String>();
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<HashMap<String, State>> {
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::<u64>()
.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<HashMap<String, State>> = 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!("<path> does not exist");
}
}