wg-rs: add cleaner library structs and methods

3 files changed, 296 insertions, 77 deletions

diff --git a/wg-rs/src/lib.rs b/wg-rs/src/lib.rs
index 2885dc2..b0313b1 100644
--- a/wg-rs/src/lib.rs
+++ b/wg-rs/src/lib.rs
@@ -1,106 +1,207 @@
+mod raw;
 use std::ffi::CStr;
+use std::net::{IpAddr, SocketAddr};
 
-pub fn list_device_names() -> Vec<String> {
-    let mut interfaces = Vec::new();
-    let dev_names = unsafe { wg_sys::wg_list_device_names() };
-
-    let mut offset = 0;
-    loop {
-        let devn = unsafe { CStr::from_ptr(dev_names.add(offset)) };
-        let name = devn.to_str().unwrap();
-        if name.is_empty() {
-            break;
-        }
-        offset += name.len() + 1;
-        interfaces.push(name.to_string());
-    }
+use raw::wg_get_device;
 
-    interfaces
+#[derive(Debug)]
+pub struct WgDevice {
+    pub name: String,
+    pub ifindex: u32,
+    pub private_key: String,
+    pub public_key: String,
+    pub fwmark: u32,
+    pub listen_port: u16,
+    pub peers: Vec<WgPeer>,
+    pub flags: WgDeviceFlags,
 }
 
-pub fn generate_private_key() -> [u8; 32] {
-    let mut key = [0u8; 32];
-    unsafe {
-        wg_sys::wg_generate_private_key(key.as_mut_ptr());
-    }
-    key
+#[allow(non_snake_case)]
+#[derive(Debug)]
+pub struct WgDeviceFlags {
+    pub ReplacePeers: bool,
+    pub HasPrivateKey: bool,
+    pub HasPublicKey: bool,
+    pub HasListenPort: bool,
+    pub HasFwmark: bool,
 }
 
-pub fn generate_public_key(mut private_key: [u8; 32]) -> [u8; 32] {
-    let mut public_key = [0u8; 32];
-    unsafe {
-        wg_sys::wg_generate_public_key(public_key.as_mut_ptr(), private_key.as_mut_ptr());
+impl From<wg_sys::wg_device> for WgDevice {
+    fn from(device: wg_sys::wg_device) -> Self {
+        let mut dev = Self {
+            name: CStr::from_bytes_until_nul(
+                &device.name.iter().map(|&c| c as u8).collect::<Vec<u8>>(),
+            )
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .to_string(),
+            ifindex: device.ifindex,
+            private_key: raw::wg_key_to_base64(device.private_key),
+            public_key: raw::wg_key_to_base64(device.public_key),
+            fwmark: device.fwmark,
+            listen_port: device.listen_port,
+            peers: vec![],
+            flags: WgDeviceFlags {
+                ReplacePeers: (device.flags.0 & wg_sys::wg_device_flags::WGDEVICE_REPLACE_PEERS.0)
+                    != 0,
+                HasPrivateKey: (device.flags.0
+                    & wg_sys::wg_device_flags::WGDEVICE_HAS_PRIVATE_KEY.0)
+                    != 0,
+                HasPublicKey: (device.flags.0 & wg_sys::wg_device_flags::WGDEVICE_HAS_PUBLIC_KEY.0)
+                    != 0,
+                HasListenPort: (device.flags.0
+                    & wg_sys::wg_device_flags::WGDEVICE_HAS_LISTEN_PORT.0)
+                    != 0,
+                HasFwmark: (device.flags.0 & wg_sys::wg_device_flags::WGDEVICE_HAS_FWMARK.0) != 0,
+            },
+        };
+
+        let mut peer_opt = unsafe { device.first_peer.as_ref() };
+
+        while let Some(peer) = peer_opt {
+            dev.peers.push(WgPeer::from(peer));
+            peer_opt = unsafe { peer.next_peer.as_ref() };
+        }
+
+        dev
     }
-    public_key
 }
 
-pub fn generate_preshared_key() -> [u8; 32] {
-    let mut key = [0u8; 32];
-    unsafe {
-        wg_sys::wg_generate_preshared_key(key.as_mut_ptr());
-    }
-    key
+#[derive(Debug)]
+pub struct WgPeer {
+    pub public_key: String,
+    pub preshared_key: String,
+    pub endpoint: SocketAddr,
+    pub allowed_ips: Vec<WgAllowedIp>,
+    pub rx_bytes: u64,
+    pub tx_bytes: u64,
+    pub persistent_keepalive: u16,
+    pub last_handshake_time: (u64, u64),
+    pub flags: WgPeerFlags,
+}
+#[allow(non_snake_case)]
+#[derive(Debug)]
+pub struct WgPeerFlags {
+    pub RemoveMe: bool,
+    pub ReplaceAllowedIps: bool,
+    pub HasPublicKey: bool,
+    pub HasPresharedKey: bool,
+    pub HasPersistentKeepalive: bool,
 }
 
-pub fn key_to_base64(mut key: [u8; 32]) -> String {
-    let mut encoded = [0i8; 45];
+impl From<&wg_sys::wg_peer> for WgPeer {
+    fn from(peer: &wg_sys::wg_peer) -> Self {
+        let endpoint = unsafe {
+            match peer.endpoint.addr.as_ref().sa_family as i32 {
+                wg_sys::AF_INET => SocketAddr::new(
+                    IpAddr::V4(peer.endpoint.addr4.as_ref().sin_addr.s_addr.to_be().into()),
+                    peer.endpoint.addr4.as_ref().sin_port.to_be(),
+                ),
+                wg_sys::AF_INET6 => SocketAddr::new(
+                    IpAddr::V6(peer.endpoint.addr6.as_ref().sin6_addr.s6_addr.into()),
+                    peer.endpoint.addr6.as_ref().sin6_port.to_be(),
+                ),
+                _ => unimplemented!("Unsupported family"),
+            }
+        };
 
-    unsafe {
-        wg_sys::wg_key_to_base64(encoded.as_mut_ptr(), key.as_mut_ptr());
-        CStr::from_ptr(encoded.as_ptr())
-            .to_str()
-            .unwrap()
-            .to_string()
-    }
-}
+        let mut wgpeer = Self {
+            public_key: raw::wg_key_to_base64(peer.public_key),
+            preshared_key: raw::wg_key_to_base64(peer.preshared_key),
+            rx_bytes: peer.rx_bytes,
+            tx_bytes: peer.tx_bytes,
+            persistent_keepalive: peer.persistent_keepalive_interval,
+            endpoint,
+            last_handshake_time: (
+                peer.last_handshake_time.tv_sec as u64,
+                peer.last_handshake_time.tv_nsec as u64,
+            ),
+            allowed_ips: vec![],
+            flags: WgPeerFlags {
+                RemoveMe: (peer.flags.0 & wg_sys::wg_peer_flags::WGPEER_REMOVE_ME.0) != 0,
+                ReplaceAllowedIps: (peer.flags.0
+                    & wg_sys::wg_peer_flags::WGPEER_REPLACE_ALLOWEDIPS.0)
+                    != 0,
+                HasPublicKey: (peer.flags.0 & wg_sys::wg_peer_flags::WGPEER_HAS_PUBLIC_KEY.0) != 0,
+                HasPresharedKey: (peer.flags.0 & wg_sys::wg_peer_flags::WGPEER_HAS_PRESHARED_KEY.0)
+                    != 0,
+                HasPersistentKeepalive: (peer.flags.0
+                    & wg_sys::wg_peer_flags::WGPEER_HAS_PERSISTENT_KEEPALIVE_INTERVAL.0)
+                    != 0,
+            },
+        };
+
+        let mut aip_opt = unsafe { peer.first_allowedip.as_ref() };
 
-pub fn key_from_base64(encoded: &str) -> Result<[u8; 32], String> {
-    let mut key = [0u8; 32];
-    let cstr = std::ffi::CString::new(encoded).unwrap();
-    let res = unsafe { wg_sys::wg_key_from_base64(key.as_mut_ptr(), cstr.into_raw()) };
-    if res == 0 {
-        Ok(key)
-    } else {
-        Err("Invalid base64 string".to_string())
+        while let Some(aip) = aip_opt {
+            wgpeer.allowed_ips.push(WgAllowedIp::from(aip));
+            aip_opt = unsafe { aip.next_allowedip.as_ref() };
+        }
+
+        wgpeer
     }
 }
 
-pub fn get_device(name: &str) -> Result<wg_sys::wg_device, String> {
-    let cstr = std::ffi::CString::new(name).unwrap();
-    let device = Box::new(wg_sys::wg_device {
-        name: [0i8; 16],
-        ifindex: 0,
-        flags: wg_sys::wg_device_flags(0),
-        listen_port: 0,
-        fwmark: 0,
-        public_key: [0u8; 32],
-        private_key: [0u8; 32],
-        first_peer: std::ptr::null_mut(),
-        last_peer: std::ptr::null_mut(),
-    });
-    let devptr = &mut Box::into_raw(device);
-    let res = unsafe { wg_sys::wg_get_device(devptr, cstr.into_raw()) };
-    if res == 0 {
+#[derive(Debug)]
+pub struct WgAllowedIp {
+    pub ip: IpAddr,
+    pub cidr: u8,
+}
+impl From<&wg_sys::wg_allowedip> for WgAllowedIp {
+    fn from(allowed_ip: &wg_sys::wg_allowedip) -> Self {
         unsafe {
-            Ok(**devptr)
+            match allowed_ip.family as i32 {
+                wg_sys::AF_INET => Self {
+                    ip: IpAddr::V4(
+                        allowed_ip
+                            .__bindgen_anon_1
+                            .ip4
+                            .as_ref()
+                            .s_addr
+                            .to_be()
+                            .into(),
+                    ),
+                    cidr: allowed_ip.cidr,
+                },
+                wg_sys::AF_INET6 => Self {
+                    ip: IpAddr::V6(allowed_ip.__bindgen_anon_1.ip6.as_ref().s6_addr.into()),
+                    cidr: allowed_ip.cidr,
+                },
+                _ => unimplemented!("Unsupported family"),
+            }
         }
-    } else {
-        Err(format!("error getting device (res = {res})"))
     }
 }
 
+pub fn list_device_names() -> Vec<String> {
+    raw::wg_list_device_names()
+}
+
+pub fn generate_private_key() -> String {
+    raw::wg_key_to_base64(raw::wg_generate_private_key())
+}
+
+pub fn generate_public_key(private_key: &str) -> String {
+    let private_key = raw::wg_key_from_base64(private_key).unwrap();
+    raw::wg_key_to_base64(raw::wg_generate_public_key(private_key))
+}
+
+pub fn generate_preshared_key() -> String {
+    raw::wg_key_to_base64(raw::wg_generate_preshared_key())
+}
+
+pub fn get_device(name: &str) -> Result<WgDevice, String> {
+    let dev = wg_get_device(name).unwrap();
+    Ok(WgDevice::from(dev))
+}
+
 #[cfg(test)]
-mod tests {
+mod test {
     use super::*;
 
     #[test]
-    fn base64_translations() {
-        let key = generate_private_key();
-        assert!(key_from_base64(&key_to_base64(key)).unwrap() == key);
-    }
-
-    #[test]
     fn test_get_device() {
-        dbg!(get_device("wg0"));
+        dbg!(get_device("wg0").unwrap());
     }
 }
diff --git a/wg-rs/src/raw.rs b/wg-rs/src/raw.rs
new file mode 100644
index 0000000..320936a
--- /dev/null
+++ b/wg-rs/src/raw.rs
@@ -0,0 +1,116 @@
+use std::ffi::CStr;
+
+pub fn wg_list_device_names() -> Vec<String> {
+    let mut interfaces = Vec::new();
+    let dev_names = unsafe { wg_sys::wg_list_device_names() };
+
+    let mut offset = 0;
+    loop {
+        let devn = unsafe { CStr::from_ptr(dev_names.add(offset)) };
+        let name = devn.to_str().unwrap();
+        if name.is_empty() {
+            break;
+        }
+        offset += name.len() + 1;
+        interfaces.push(name.to_string());
+    }
+
+    interfaces
+}
+
+pub fn wg_generate_private_key() -> [u8; 32] {
+    let mut key = [0u8; 32];
+    unsafe {
+        wg_sys::wg_generate_private_key(key.as_mut_ptr());
+    }
+    key
+}
+
+pub fn wg_generate_public_key(mut private_key: [u8; 32]) -> [u8; 32] {
+    let mut public_key = [0u8; 32];
+    unsafe {
+        wg_sys::wg_generate_public_key(public_key.as_mut_ptr(), private_key.as_mut_ptr());
+    }
+    public_key
+}
+
+pub fn wg_generate_preshared_key() -> [u8; 32] {
+    let mut key = [0u8; 32];
+    unsafe {
+        wg_sys::wg_generate_preshared_key(key.as_mut_ptr());
+    }
+    key
+}
+
+pub fn wg_key_to_base64(mut key: [u8; 32]) -> String {
+    let mut encoded = [0i8; 45];
+
+    unsafe {
+        wg_sys::wg_key_to_base64(encoded.as_mut_ptr(), key.as_mut_ptr());
+        CStr::from_ptr(encoded.as_ptr())
+            .to_str()
+            .unwrap()
+            .to_string()
+    }
+}
+
+pub fn wg_key_from_base64(encoded: &str) -> Result<[u8; 32], String> {
+    let mut key = [0u8; 32];
+    let cstr = std::ffi::CString::new(encoded).unwrap();
+    let res = unsafe { wg_sys::wg_key_from_base64(key.as_mut_ptr(), cstr.into_raw()) };
+    if res == 0 {
+        Ok(key)
+    } else {
+        Err("Invalid base64 string".to_string())
+    }
+}
+
+pub fn wg_get_device(name: &str) -> Result<wg_sys::wg_device, String> {
+    let cstr = std::ffi::CString::new(name).unwrap();
+    let device = Box::new(wg_sys::wg_device {
+        name: [0i8; 16],
+        ifindex: 0,
+        flags: wg_sys::wg_device_flags(0),
+        listen_port: 0,
+        fwmark: 0,
+        public_key: [0u8; 32],
+        private_key: [0u8; 32],
+        first_peer: std::ptr::null_mut(),
+        last_peer: std::ptr::null_mut(),
+    });
+    let devptr = &mut Box::into_raw(device);
+    let res = unsafe { wg_sys::wg_get_device(devptr, cstr.into_raw()) };
+    if res == 0 {
+        unsafe { Ok(**devptr) }
+    } else {
+        Err(format!("error getting device (res = {res})"))
+    }
+}
+
+pub fn wg_set_device(mut device: wg_sys::wg_device) -> Result<(), String> {
+    let res = unsafe { wg_sys::wg_set_device(&mut device) };
+    if res == 0 {
+        Ok(())
+    } else {
+        Err(format!("error setting device (res = {res})"))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn base64_translations() {
+        let key = wg_generate_private_key();
+        assert!(wg_key_from_base64(&wg_key_to_base64(key)).unwrap() == key);
+    }
+
+    #[test]
+    fn test_get_device() {
+        let _ = dbg!(wg_get_device("wg0").unwrap());
+        // dev.private_key = wg_generate_private_key();
+        // dev.public_key = wg_generate_public_key(dev.private_key);
+        // wg_set_device(dbg!(dev)).unwrap();
+    }
+}
diff --git a/wg-sys/src/lib.rs b/wg-sys/src/lib.rs
index ca85511..2cf9dd9 100644
--- a/wg-sys/src/lib.rs
+++ b/wg-sys/src/lib.rs
@@ -5,3 +5,5 @@
 #![allow(clippy::missing_safety_doc)]
 
 include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+
+pub use libc::{AF_INET, AF_INET6};
\ No newline at end of file