path: root/src/daemon.rs

use atomic_write_file::AtomicWriteFile;
use base64::prelude::*;
use core::net::SocketAddr;
use dbus::{channel::MatchingReceiver, message::MatchRule};
use dbus_crossroads::{Context, Crossroads, MethodErr};
use defguard_wireguard_rs::{
    host::Peer, key::Key, net::IpAddrMask, InterfaceConfiguration, WGApi, WireguardInterfaceApi,
};
use log::{debug, error, info, warn};
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::{
    collections::{BTreeSet, HashMap},
    fs::File,
    io::{ErrorKind, Read, Write},
    marker::PhantomData,
    net::ToSocketAddrs,
    ops::DerefMut,
    sync::Arc,
    time::SystemTime,
};
use thiserror::Error;
use tokio::{
    net::TcpListener,
    runtime::Builder,
    signal::unix::{signal, SignalKind},
    sync::{broadcast, RwLock},
    task,
};
use xdg::BaseDirectories;

use std::str::FromStr;

#[derive(Debug, Error)]
pub enum DaemonError {
    #[error("{0}")]
    Io(#[from] std::io::Error),

    #[error("{0}")]
    XdgBase(#[from] xdg::BaseDirectoriesError),

    // TODO hier wärs nett zu unterscheiden was decoded wurde
    #[error("{0}")]
    Decoding(#[from] serde_json::Error),

    #[error("{0}")]
    IpMaskParse(#[from] defguard_wireguard_rs::net::IpAddrParseError),

    #[error("{0}")]
    WgInterfaceError(#[from] defguard_wireguard_rs::error::WireguardInterfaceError),

    #[error("{0}")]
    DbusError(#[from] dbus::Error),
}

#[derive(Serialize, Deserialize, Clone)]
enum Endpoint {
    Ip(SocketAddr),
    Domain(String, u16),
}

// subset of defguard_wireguard_rs::host::Peer, with hostname added
#[derive(Serialize, Deserialize)]
struct PeerConfig {
    psk: Option<Key>,
    ips: Vec<(IpAddrMask, Option<String>)>,
    // if false: the hostnames are kept around for sharing, but we personally do not use them
    use_hostnames: bool,
    endpoint: Option<Endpoint>,

    last_changed: SystemTime,
    known_to: Vec<usize>,

    mäsch_endpoint: SocketAddr,
}

fn default_wg_port() -> u16 {
    51820
}

#[derive(Serialize, Deserialize)]
struct Network {
    privkey: String,

    // this really should be a different type, but this is what defguard takes...
    address: String,

    #[serde(default = "default_wg_port")]
    listen_port: u16,

    peers: HashMap<Key, PeerConfig>,

    mäsch_port: u16,
}

#[derive(Serialize, Deserialize, Default)]
struct Config {
    networks: HashMap<String, Network>,
}

// TODO das überschreibt änderungen an /etc/hosts während der runtime :(
struct State {
    conf: Config,
    nw_handles: HashMap<String, (WGApi, task::JoinHandle<()>)>,
    hostfile: Option<(String, BTreeSet<String>)>,
}

impl Drop for State {
    fn drop(&mut self) {
        for (wg_api, _) in self.nw_handles.values() {
            let _ = wg_api.remove_interface();
        }
    }
}

pub fn daemon() -> Result<(), DaemonError> {
    let config_path = BaseDirectories::with_prefix("mäsch")?.place_state_file("daemon.json")?;
    let config: Config = match File::open(config_path) {
        Ok(f) => serde_json::from_reader(f)?,
        Err(e) => match e.kind() {
            ErrorKind::NotFound => Config::default(),
            _ => Err(e)?,
        },
    };
    info!("read config");

    let hostfile = match File::open("/etc/hosts") {
        Ok(mut f) => {
            let mut r = String::new();
            f.read_to_string(&mut r)?;

            let seen_hostnames: BTreeSet<String> = r
                .lines()
                .map(|l| {
                    l.split_whitespace()
                        .take_while(|dom| dom.chars().next().unwrap() != '#')
                        .skip(1)
                })
                .flatten()
                .map(|dom| dom.to_owned())
                .collect();

            Some((r, seen_hostnames))
        }
        Err(e) => {
            warn!("failed to read /etc/hosts: {e}");
            None
        }
    };

    let state = Arc::new(RwLock::new(State {
        conf: config,
        nw_handles: HashMap::new(),
        hostfile,
    }));

    let rt = Builder::new_current_thread().enable_all().build()?;
    rt.block_on(run_networks(state))?;

    Ok(())
}

async fn run_networks(state: Arc<RwLock<State>>) -> Result<(), DaemonError> {
    let mut state_rw_guard = state.write().await;
    let state_rw = state_rw_guard.deref_mut();

    // load existing configurations
    for (name, nw) in &state_rw.conf.networks {
        let wg_api = add_network(
            &mut state_rw.hostfile,
            name.clone(),
            nw.privkey.clone(),
            nw.address.clone(),
            nw.listen_port,
            &nw.peers,
        )
        .await?;

        let addr = IpAddrMask::from_str(&nw.address)?.ip;
        let h = task::spawn(print_error(run_network(
            state.clone(),
            TcpListener::bind((addr, nw.mäsch_port)).await?,
            name.clone(),
        )));

        state_rw.nw_handles.insert(name.clone(), (wg_api, h));

        debug!("loaded configuration for {0}", name);
    }
    info!("loaded all existing configurations");
    drop(state_rw_guard);

    // set up dbus interface
    let mut cr = Crossroads::new();
    let state_ref = state.clone();
    let if_token = cr.register("de.a.maesch", move |b| {
        b.signal::<(String, String), _>("Proposal", ("network", "peer_data"));
        b.method_with_cr_async(
            "AddNetwork",
            ("name", "key", "ip", "listen_port", "maesch_port"),
            ("success",),
            move |ctx, _, args: (String, String, String, u16, u16)| {
                debug!("Received AddNetwork");
                handle_add_network(ctx, state_ref.clone(), args)
            },
        );
    });
    cr.insert("/de/a/maesch", &[if_token], ());

    // drive dbus interface
    let (res, c) = dbus_tokio::connection::new_system_sync()?;
    cr.set_async_support(Some((
        c.clone(),
        Box::new(|x| {
            tokio::spawn(x);
        }),
    )));
    let _ = tokio::spawn(print_error(async {
        res.await;
        Result::<!, &'static str>::Err("lost connection to dbus!")
    }));
    let receive_token = c.start_receive(
        MatchRule::new_method_call(),
        Box::new(move |msg, conn| {
            cr.handle_message(msg, conn).unwrap();
            true
        }),
    );
    c.request_name("de.a.maesch", false, true, false).await?;

    // wait for SIGTERM/SIGINT
    let mut sigterm_fut = signal(SignalKind::terminate())?;
    let mut sigint_fut = signal(SignalKind::interrupt())?;
    let mut sighup_fut = signal(SignalKind::hangup())?;

    tokio::select! {
        _ = sigterm_fut.recv() => info!("Received SIGTERM"),
        _ = sigint_fut.recv() => info!("Received SIGINT"),
        _ = sighup_fut.recv() => info!("Received SIGHUP"),
    };

    // clean exit
    c.stop_receive(receive_token);
    let mut state_rw_guard = state.write().await;
    for (_, (wg_api, h)) in state_rw_guard.nw_handles.drain() {
        let _ = wg_api.remove_interface();
        h.abort(); // could also join the handles... don't think that would do too much, though
    }

    Ok(())
}

// TODO also take peers
async fn handle_add_network(
    mut ctx: Context,
    state: Arc<RwLock<State>>,
    (name, may_key, may_ip, may_lp, may_mp): (String, String, String, u16, u16),
) -> PhantomData<(bool,)> {
    // NOTE: this is kinda stupid: we convert to a string later anyways, as thats what
    // defguard_wg takes...
    let key = Key::new(match may_key.as_str() {
        "" => rand::thread_rng().gen(),
        _ => match BASE64_STANDARD.decode(may_key) {
            Ok(v) if v.len() == 32 => v.try_into().unwrap(),
            _ => {
                warn!("AddNetwork with bad key");
                return ctx.reply(Err(MethodErr::invalid_arg("bad key")));
            }
        },
    });

    // we store the ip as the original string, but should validate it regardless
    let (ip, ip_string) = match may_ip.as_str() {
        "" => todo!(),
        _ => match IpAddrMask::from_str(&may_ip) {
            Ok(ip_mask) => (ip_mask.ip, may_ip),
            Err(_) => {
                warn!("AddNetwork with bad ip");
                return ctx.reply(Err(MethodErr::invalid_arg("invalid ip")));
            }
        },
    };

    let lp = if may_lp == 0 { 25565 } else { may_lp };
    let mp = if may_mp == 0 { 51820 } else { may_mp };

    let mut state_rw_guard = state.write().await;
    let state_rw = state_rw_guard.deref_mut();

    let wg_api = match add_network(
        &mut state_rw.hostfile,
        name.clone(),
        key.to_string(),
        ip_string,
        lp,
        &HashMap::new(),
    )
    .await
    {
        Ok(wg_api) => wg_api,
        Err(e) => {
            warn!("AddNetwork couldn't add network: {e}");
            return ctx.reply(Err(MethodErr::failed(&e)));
        }
    };

    // TODO ins wg_api
    let listener = match TcpListener::bind((ip, mp)).await {
        Ok(l) => l,
        Err(e) => {
            let _ = wg_api.remove_interface();
            warn!("AddNetwork couldn't start listener: {e}");
            return ctx.reply(Err(MethodErr::failed(&e)));
        }
    };
    let h = task::spawn(print_error(run_network(
        state.clone(),
        listener,
        name.clone(),
    )));

    state_rw.nw_handles.insert(name.clone(), (wg_api, h));

    // TODO save new config

    ctx.reply(Ok((true,)))
}

async fn add_network(
    hostfile: &mut Option<(String, BTreeSet<String>)>,
    name: String,
    privkey: String,
    address: String,
    port: u16,
    peers: &HashMap<Key, PeerConfig>,
) -> Result<WGApi, DaemonError> {
    let wg = WGApi::new(name.clone(), false)?;
    let defguard_peers = peers
        .iter()
        .map(|(peer_key, p)| Peer {
            public_key: peer_key.clone(),
            preshared_key: p.psk.clone(),
            protocol_version: None,
            endpoint: p
                .endpoint
                .clone()
                .map(|e| match e {
                    Endpoint::Ip(ep) => Some(ep),
                    Endpoint::Domain(s, p) => (s, p)
                        .to_socket_addrs()
                        .ok()
                        .map(|mut it| it.next())
                        .flatten(),
                })
                .flatten(),
            last_handshake: None,
            tx_bytes: 0,
            rx_bytes: 0,
            persistent_keepalive_interval: None,
            allowed_ips: p.ips.iter().map(|(ip_mask, _)| ip_mask.clone()).collect(),
        })
        .collect();
    wg.create_interface()?;
    wg.configure_interface(&InterfaceConfiguration {
        name: name.clone(),
        prvkey: privkey,
        address: address,
        port: port as u32,
        peers: defguard_peers,
    })?;

    if let Some((hosts_str, hosts)) = hostfile {
        peers
            .values()
            .map(|peer| {
                if peer.use_hostnames {
                    peer.ips
                        .iter()
                        .map(|(mask, may_dom)| {
                            if let Some(dom) = may_dom
                                && hosts.insert(dom.clone())
                            {
                                hosts_str.push_str(&format!("{}", mask.ip));
                                hosts_str.push('\t');
                                hosts_str.push_str(&dom);
                                hosts_str.push('\n');
                            }
                        })
                        .count();
                }
            })
            .count();
    }

    if let Some((hosts_str, _)) = hostfile {
        debug!("writing hosts file: {hosts_str}");

        let mut f = AtomicWriteFile::open("/etc/hosts")?;
        f.write(hosts_str.as_bytes())?;
        f.commit()?;
    }

    Ok(wg)
}

async fn print_error<E: std::fmt::Display, O, F: std::future::Future<Output = Result<O, E>>>(
    f: F,
) -> () {
    match f.await {
        Err(e) => error!("oh no: {e}"),
        _ => (),
    };
}

async fn run_network(
    state: Arc<RwLock<State>>,
    sock: TcpListener,
    nw_name: String,
) -> Result<(), DaemonError> {
    Ok(())
}