init

1 files changed, 212 insertions, 0 deletions

diff --git a/src/setup.rs b/src/setup.rs
new file mode 100644
index 0000000..e7d454c
--- /dev/null
+++ b/src/setup.rs
@@ -0,0 +1,212 @@
+use std::net::{IpAddr, SocketAddr};
+
+use ipnet::IpNet;
+use netlink_packet_wireguard::{
+    constants::{AF_INET, AF_INET6, WGDEVICE_F_REPLACE_PEERS, WGPEER_F_REPLACE_ALLOWEDIPS},
+    nlas::{WgAllowedIp, WgAllowedIpAttrs, WgDeviceAttrs, WgPeer, WgPeerAttrs},
+    Wireguard, WireguardCmd,
+};
+
+use super::Key;
+
+#[derive(Debug)]
+pub struct PeerDescriptor {
+    pub(super) public_key: Key,
+    pub(super) preshared_key: Option<Key>,
+    pub(super) endpoint: Option<SocketAddr>,
+    pub(super) keepalive: Option<u16>,
+    pub(super) allowed_ips: Option<Vec<IpNet>>,
+}
+
+impl PeerDescriptor {
+    pub fn new(public_key: Key) -> Self {
+        Self {
+            public_key,
+            preshared_key: None,
+            endpoint: None,
+            keepalive: None,
+            allowed_ips: None,
+        }
+    }
+
+    pub fn preshared_key_optional(mut self, preshared_key: Option<Key>) -> Self {
+        self.preshared_key = preshared_key;
+        self
+    }
+
+    pub fn preshared_key(mut self, preshared_key: Key) -> Self {
+        self.preshared_key = Some(preshared_key);
+        self
+    }
+
+    pub fn endpoint_optional(mut self, endpoint: Option<SocketAddr>) -> Self {
+        self.endpoint = endpoint;
+        self
+    }
+
+    pub fn endpoint(mut self, endpoint: SocketAddr) -> Self {
+        self.endpoint = Some(endpoint);
+        self
+    }
+
+    pub fn keepalive_optional(mut self, keepalive: Option<u16>) -> Self {
+        self.keepalive = keepalive;
+        self
+    }
+
+    pub fn keepalive(mut self, keepalive: u16) -> Self {
+        self.keepalive = Some(keepalive);
+        self
+    }
+
+    pub fn allowed_ip_optional(self, allowed_ip: Option<IpNet>) -> Self {
+        if let Some(allowed_ip) = allowed_ip {
+            self.allowed_ip(allowed_ip)
+        } else {
+            self
+        }
+    }
+
+    pub fn allowed_ip(mut self, allowed_ip: IpNet) -> Self {
+        let mut allowed_ips = self.allowed_ips.take().unwrap_or_default();
+        allowed_ips.push(allowed_ip);
+        self.allowed_ips = Some(allowed_ips);
+        self
+    }
+
+    pub fn allowed_ips_optional(self, allowed_ips: Option<Vec<IpNet>>) -> Self {
+        if let Some(allowed_ips) = allowed_ips {
+            self.allowed_ips(allowed_ips)
+        } else {
+            self
+        }
+    }
+
+    pub fn allowed_ips(mut self, allowed_ips: Vec<IpNet>) -> Self {
+        self.allowed_ips = Some(allowed_ips);
+        self
+    }
+
+    pub(super) fn into_wireguard(self) -> WgPeer {
+        let mut nlas = Vec::new();
+        nlas.push(WgPeerAttrs::PublicKey(self.public_key.into_array()));
+        nlas.extend(
+            self.preshared_key
+                .map(|key| WgPeerAttrs::PresharedKey(key.into_array())),
+        );
+        nlas.extend(self.endpoint.map(WgPeerAttrs::Endpoint));
+        nlas.extend(self.keepalive.map(WgPeerAttrs::PersistentKeepalive));
+        nlas.extend(self.allowed_ips.map(|allowed_ips| {
+            WgPeerAttrs::AllowedIps(allowed_ips.into_iter().map(ipnet_to_wg).collect())
+        }));
+        nlas.push(WgPeerAttrs::Flags(WGPEER_F_REPLACE_ALLOWEDIPS));
+        WgPeer(nlas)
+    }
+}
+
+#[derive(Debug)]
+pub struct DeviceDescriptor {
+    pub(super) addresses: Vec<IpNet>,
+    pub(super) private_key: Option<Key>,
+    pub(super) listen_port: Option<u16>,
+    pub(super) fwmark: Option<u32>,
+    pub(super) peers: Option<Vec<PeerDescriptor>>,
+}
+
+impl Default for DeviceDescriptor {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl DeviceDescriptor {
+    pub fn new() -> Self {
+        Self {
+            addresses: Vec::default(),
+            private_key: None,
+            listen_port: None,
+            fwmark: None,
+            peers: None,
+        }
+    }
+
+    pub fn address(mut self, address: IpNet) -> Self {
+        self.addresses.push(address);
+        self
+    }
+
+    pub fn addresses(mut self, addresses: impl IntoIterator<Item = IpNet>) -> Self {
+        self.addresses.extend(addresses);
+        self
+    }
+
+    pub fn private_key(mut self, key: Key) -> Self {
+        self.private_key = Some(key);
+        self
+    }
+
+    pub fn listen_port(mut self, port: u16) -> Self {
+        self.listen_port = Some(port);
+        self
+    }
+
+    pub fn listen_port_optional(mut self, port: Option<u16>) -> Self {
+        self.listen_port = port;
+        self
+    }
+
+    pub fn fwmark(mut self, fwmark: u32) -> Self {
+        self.fwmark = Some(fwmark);
+        self
+    }
+
+    pub fn peer(mut self, peer: PeerDescriptor) -> Self {
+        let mut p = self.peers.take().unwrap_or_default();
+        p.push(peer);
+        self.peers = Some(p);
+        self
+    }
+
+    pub fn peers(mut self, peers: impl IntoIterator<Item = PeerDescriptor>) -> Self {
+        let mut p = self.peers.take().unwrap_or_default();
+        p.extend(peers);
+        self.peers = Some(p);
+        self
+    }
+
+    pub(super) fn into_wireguard(self, device_name: String) -> Wireguard {
+        let mut nlas = Vec::new();
+        nlas.push(WgDeviceAttrs::IfName(device_name));
+        nlas.extend(
+            self.private_key
+                .map(|key| WgDeviceAttrs::PrivateKey(key.into_array())),
+        );
+        nlas.extend(self.listen_port.map(WgDeviceAttrs::ListenPort));
+        nlas.extend(self.fwmark.map(WgDeviceAttrs::Fwmark));
+        nlas.extend(self.peers.map(|peers| {
+            WgDeviceAttrs::Peers(
+                peers
+                    .into_iter()
+                    .map(PeerDescriptor::into_wireguard)
+                    .collect(),
+            )
+        }));
+        nlas.push(WgDeviceAttrs::Flags(WGDEVICE_F_REPLACE_PEERS));
+
+        Wireguard {
+            cmd: WireguardCmd::SetDevice,
+            nlas,
+        }
+    }
+}
+
+fn ipnet_to_wg(net: IpNet) -> WgAllowedIp {
+    let mut nlas = Vec::default();
+    nlas.push(WgAllowedIpAttrs::Cidr(net.prefix_len()));
+    nlas.push(WgAllowedIpAttrs::IpAddr(net.addr()));
+    match net.addr() {
+        IpAddr::V4(_) => nlas.push(WgAllowedIpAttrs::Family(AF_INET)),
+        IpAddr::V6(_) => nlas.push(WgAllowedIpAttrs::Family(AF_INET6)),
+    }
+    WgAllowedIp(nlas)
+}