297 lines
9.0 KiB
Rust
297 lines
9.0 KiB
Rust
mod gen;
|
|
mod listener;
|
|
|
|
use crate::gen::*;
|
|
use crate::listener::*;
|
|
|
|
use base64;
|
|
use hex;
|
|
use std::collections::HashMap;
|
|
use std::env;
|
|
use std::fmt;
|
|
use std::io::prelude::*;
|
|
use std::io::{BufRead, BufReader, Error, ErrorKind, Result};
|
|
use std::net::{IpAddr, SocketAddr};
|
|
use std::os::unix::net::UnixStream;
|
|
use std::path::PathBuf;
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
use time;
|
|
|
|
pub type KV = (String, String);
|
|
|
|
#[derive(Debug, PartialEq, Eq, Hash)]
|
|
enum State {
|
|
Interface(Vec<KV>),
|
|
Peer(Vec<KV>),
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
|
|
pub struct Peer {
|
|
public_key: String,
|
|
endpoint: Option<SocketAddr>,
|
|
allowed_ips: Vec<(IpAddr, u8)>,
|
|
last_handshake: Option<Duration>,
|
|
persistent_keepalive: Option<Duration>,
|
|
traffic: (u64, u64),
|
|
parsed: time::Timespec,
|
|
}
|
|
|
|
impl Peer {
|
|
fn from_kv(entries: &Vec<KV>) -> Result<Peer> {
|
|
let key = match entries
|
|
.iter()
|
|
.filter(|(key, _)| key == &"public_key")
|
|
.map(|(_, value)| value)
|
|
.next()
|
|
{
|
|
Some(key) => key,
|
|
None => return Err(Error::new(ErrorKind::Other, "Peer is missing key")),
|
|
};
|
|
Ok(Peer {
|
|
public_key: base64::encode(&hex::decode(key).unwrap()),
|
|
endpoint: entries
|
|
.iter()
|
|
.filter(|(key, _)| key == &"endpoint")
|
|
.map(|(_, value)| value.parse::<SocketAddr>().unwrap())
|
|
.next(),
|
|
allowed_ips: entries
|
|
.iter()
|
|
.filter(|(key, _)| key == &"allowed_ip")
|
|
.map(|(_, value)| {
|
|
let mut parts = value.split("/").into_iter();
|
|
match (
|
|
parts.next().and_then(|addr| addr.parse::<IpAddr>().ok()),
|
|
parts.next().and_then(|mask| mask.parse::<u8>().ok()),
|
|
) {
|
|
(Some(addr), Some(mask)) => Some((addr, mask)),
|
|
(Some(addr), None) if addr.is_ipv6() => Some((addr, 128)),
|
|
(Some(addr), None) => Some((addr, 32)),
|
|
_ => None,
|
|
}
|
|
})
|
|
.filter_map(|net| net)
|
|
.collect::<Vec<(IpAddr, u8)>>(),
|
|
last_handshake: entries
|
|
.iter()
|
|
.filter_map(|(key, value)| {
|
|
let value = || value.parse::<u64>().unwrap();
|
|
match key.as_ref() {
|
|
"last_handshake_time_sec" if value() != 0 => {
|
|
Some(Duration::new(value(), 0))
|
|
}
|
|
"last_handshake_time_nsec" if value() != 0 => {
|
|
Some(Duration::from_nanos(value()))
|
|
}
|
|
_ => None,
|
|
}
|
|
})
|
|
.fold(None, |acc, add| {
|
|
if let Some(dur) = acc {
|
|
Some(dur + add)
|
|
} else {
|
|
Some(add)
|
|
}
|
|
}),
|
|
persistent_keepalive: entries
|
|
.iter()
|
|
.filter(|(key, _)| key == &"persistent_keepalive")
|
|
.map(|(_, value)| Duration::from_secs(value.parse::<u64>().unwrap()))
|
|
.next(),
|
|
traffic: (0, 0),
|
|
parsed: time::get_time(),
|
|
})
|
|
}
|
|
|
|
pub fn last_handshake_rel(&self) -> Option<Duration> {
|
|
let time = self.parsed;
|
|
Some(Duration::new(time.sec as u64, time.nsec as u32) - self.last_handshake?)
|
|
}
|
|
}
|
|
|
|
impl State {
|
|
pub 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<String> {
|
|
self.kv()
|
|
.iter()
|
|
.filter(|(key, _)| key == &"private_key" || key == &"public_key")
|
|
.map(|(_, value)| base64::encode(&hex::decode(&value).unwrap()))
|
|
.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>>()
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for State {
|
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
|
for (k, v) in self.kv() {
|
|
write!(f, "({:10}= {})", k, v)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl fmt::Display for Peer {
|
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
|
write!(f, "{:?}", self)
|
|
// write!(f, "peer {}\nshake {} ago\naddr {}\nkeepalive {}\n", self.public_key, self.last_handshake.map(|d|d.to_string()).unwrap_or("-"), self.endpoint.map(|d|d.to_string()).unwrap_or("-"), self.persistent_keepalive.map(|d|d.to_string()).unwrap_or("-"))
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
pub fn get_peers(&self) -> Result<HashMap<String, Peer>> {
|
|
let by_id = self.get_by_id()?;
|
|
Ok(by_id
|
|
.iter()
|
|
.filter_map(|(id, state)| {
|
|
Peer::from_kv(state.kv())
|
|
.ok()
|
|
.map(|peer| (id.to_owned(), peer))
|
|
})
|
|
.collect())
|
|
}
|
|
}
|
|
|
|
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 mut listeners: Vec<Box<EventListener>> = vec![Box::new(LogListener)];
|
|
|
|
let events: PathBuf = "/etc/wireguard/events.sh".into();
|
|
|
|
if events.exists() {
|
|
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);
|
|
}
|
|
prev_state = Some(state);
|
|
thread::sleep(interval);
|
|
}
|
|
} else {
|
|
println!("<path> does not exist");
|
|
}
|
|
}
|