From 0064f0daa98d60c1bd40aece4c898d3e6b8b24ea Mon Sep 17 00:00:00 2001
From: alemidev <me@alemi.dev>
Date: Fri, 30 Sep 2022 02:06:50 +0200
Subject: [PATCH] feat: initial implementation keeps track of network per
 session

---
 Cargo.toml  |   9 +++++
 src/main.rs | 109 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 Cargo.toml
 create mode 100644 src/main.rs

diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..835a7aa
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "server-monitor"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+rocket = { version = "0.5.0-rc.2", features = ["json"] }
+systemstat = "0.2"
+serde = "1"
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..1e9c86e
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,109 @@
+#[macro_use] extern crate rocket;
+
+use std::collections::HashMap;
+
+use rocket::State;
+use rocket::serde::{Serialize, json::Json};
+use rocket::tokio::sync::Mutex;
+use systemstat::{System, Platform};
+
+#[derive(Serialize)]
+struct SystemInfoView {
+	session: String,
+	mem: f32,
+	cpu: f32,
+	tx: u64,
+	rx: u64,
+	disk: u64,
+}
+
+struct NetDeltaTracker {
+	tx: u64,
+	rx: u64,
+}
+
+impl Default for NetDeltaTracker {
+	fn default() -> Self {
+		NetDeltaTracker { tx: 0, rx: 0 }
+	}
+}
+
+impl NetDeltaTracker {
+	pub fn delta(&mut self, tx: u64, rx: u64) -> (u64, u64) {
+		let out = (tx - self.tx, rx - self.rx);
+		self.tx = tx;
+		self.rx = rx;
+		return out;
+	}
+}
+
+struct NetDeltaState {
+	pub store: Mutex<HashMap<String, NetDeltaTracker>>,
+}
+
+#[get("/<session>")]
+async fn stats_delta(session: String, state: &State<NetDeltaState>) -> Json<SystemInfoView> {
+	let sys = System::new();
+	let (delta_tx, delta_rx);
+
+	let (net_tx, net_rx) = match sys.network_stats(&"enp5s0".to_string()) {
+		Ok(stats) => (stats.tx_bytes.as_u64(), stats.rx_bytes.as_u64()),
+		Err(x) => (0, 0),
+	};
+
+	{
+		let mut store = state.inner().store.lock().await;
+		if !store.contains_key(&session) {
+			store.insert(session.clone(), NetDeltaTracker::default());
+		}
+
+		let net_tracker = store.get_mut(&session).unwrap();
+
+		(delta_tx, delta_rx) = net_tracker.delta(net_tx, net_rx);
+	}
+
+	Json(
+		SystemInfoView {
+			session,
+			mem: match sys.memory() {
+				Ok(mem) => { 1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32)) },
+				Err(e) => { -1.0 },
+			},
+			cpu: 0.0,
+			tx: delta_tx,
+			rx: delta_rx,
+			disk: 0,
+		}
+	)
+}
+
+#[get("/")]
+async fn stats_total() -> Json<SystemInfoView> {
+	let sys = System::new();
+
+	let (net_tx, net_rx) = match sys.network_stats(&"enp5s0".to_string()) {
+		Ok(stats) => (stats.tx_bytes.as_u64(), stats.rx_bytes.as_u64()),
+		Err(x) => (0, 0),
+	};
+
+	Json(
+		SystemInfoView {
+			session: "<TOTAL>".to_string(),
+			mem: match sys.memory() {
+				Ok(mem) => { 1.0 - ((mem.free.as_u64() as f32) / (mem.total.as_u64() as f32)) },
+				Err(e) => { -1.0 },
+			},
+			cpu: 0.0,
+			tx: net_tx,
+			rx: net_rx,
+			disk: 0,
+		}
+	)
+}
+
+#[launch]
+fn rocket() -> _ {
+	rocket::build()
+		.manage(NetDeltaState { store: Mutex::new(HashMap::new()) })
+		.mount("/", routes![stats_delta, stats_total])
+}