This commit is contained in:
parent
fb7d706dae
commit
819eb7d362
@ -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);
|
||||
|
||||
|
@ -4,8 +4,8 @@ use std::collections::{HashMap, HashSet};
|
||||
use std::time;
|
||||
|
||||
pub(crate) fn gen_events(
|
||||
state: &HashMap<String, Peer>,
|
||||
prev: &HashMap<String, Peer>,
|
||||
state: &HashMap<ECCKey, Peer>,
|
||||
prev: &HashMap<ECCKey, Peer>,
|
||||
listeners: &Vec<Box<EventListener>>,
|
||||
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::<HashSet<&str>>()
|
||||
.chain(prev.keys())
|
||||
.collect::<HashSet<&ECCKey>>()
|
||||
.iter()
|
||||
.map(|p| (*p, (prev.get(*p), state.get(*p))))
|
||||
.collect::<HashMap<&str, (Option<&Peer>, Option<&Peer>)>>()
|
||||
.collect::<HashMap<&ECCKey, (Option<&Peer>, 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) {
|
||||
|
@ -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
|
||||
|
@ -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<Box<EventListener>> {
|
||||
let mut listeners: Vec<Box<EventListener>> = 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<WireguardController> = 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<HashMap<ECCKey, Peer>> = 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::<io::Result<HashMap<_, _>>>()?;
|
||||
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);
|
||||
|
@ -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<ECCKey> {
|
||||
pub fn public_key(&self) -> Option<ECCKey> {
|
||||
//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<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)]
|
||||
impl Hash for Interface {
|
||||
fn hash<H: Hasher>(&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<H: Hasher>(&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<T>) -> 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::<Vec<_>>().join(",")
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
)
|
||||
}
|
||||
}
|
||||
|
@ -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,
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user