From 4020bfeddb13ee86d7f51158190d8cb9a39a614b Mon Sep 17 00:00:00 2001
From: alemidev <me@alemi.dev>
Date: Sat, 13 Aug 2022 04:45:31 +0200
Subject: [PATCH] feat: reworked script with async loop

now allows async tasks with separate timers for different kind of
packets: net and cpu leds are fast, display redraw is slow
---
 script/cpu-net-monitor.py | 120 ++++++++++++++++++++++++++------------
 1 file changed, 83 insertions(+), 37 deletions(-)

diff --git a/script/cpu-net-monitor.py b/script/cpu-net-monitor.py
index a507889..4234448 100755
--- a/script/cpu-net-monitor.py
+++ b/script/cpu-net-monitor.py
@@ -1,52 +1,98 @@
 #!/usr/bin/env python
+import asyncio
 import sys
 import struct
+import logging
 from time import sleep
 
 import serial
 import psutil
 
-def cpu_load_serial_driver(device:str, retry_interval:float=5.0):
-	while True:
-		try:
-			port = serial.Serial(device, baudrate=57600)
-			avg_usage_to_serial(port)
-		except serial.SerialException as e:
-			print(f"[!] Could not connect to device: {str(e)}", file=sys.stderr)
-		else:
-			port.close()
-		sleep(retry_interval)
+class State:
+	run: bool
+	device: str
+	baudrate: int
+	loop: asyncio.AbstractEventLoop
+	packets: asyncio.Queue
+	_port: serial.Serial
+	_task: asyncio.Task
 
-def avg_usage_to_serial(port:serial.Serial):
-	net = psutil.net_io_counters()
-	net_tx = net.bytes_sent
-	net_rx = net.bytes_recv
-	port.write(struct.pack("BB", 0, 0))
-	port.flush()
-	while True:
-		# Map float [0:100] to int [0:255], square it to put more values in the lower end, where led is more sensible
-		load = [ int(((x/100) **2) * 255) for x in psutil.cpu_percent(0.05, percpu=True) ] # mypy whines but percpu returns a list
-		try:
-			port.write(struct.pack("BBBBBB", 1, 4, *load))
-			port.flush()
-		except serial.SerialException as e:
-			print(f"[!] Failed writing payload to device: {str(e)}", file=sys.stderr)
-			break
-		net = psutil.net_io_counters()
-		try:
-			port.write(struct.pack("BBBB", 2, 2, int(net.bytes_sent > net_tx), int(net.bytes_recv > net_rx)))
-			port.flush()
-		except serial.SerialException as e:
-			print(f"[!] Failed writing payload to device: {str(e)}", file=sys.stderr)
-			break
-		net_rx = net.bytes_recv
-		net_tx = net.bytes_sent
+	def __init__(self, device:str="/dev/ttyUSB0", baudrate=57600):
+		self.run = True
+		self.baudrate = baudrate
+		self.device = device
+		self.packets = asyncio.Queue()
+		self.port = None
 
+	async def run_port_manager(self):
+		while self.run:
+			logging.debug("[*] Connecting to device '%s'", self.device)
+			try:
+				self.port = await self._loop.run_in_executor(None, serial.Serial, self.device, self.baudrate)
+				await self._loop.run_in_executor(None, self.port.write, struct.pack("BB", 0, 0))
+				while self.run:
+					pkt = await self.packets.get()
+					logging.debug("[>] Dispatching packet [ %s ]", str(pkt))
+					await self._loop.run_in_executor(None, self.port.write, pkt)
+			except serial.SerialException as e:
+				logging.error("[!] Error operating with device: %s", str(e))
+			except Exception as e:
+				logging.exception("unhandled exception")
+				self.run = False
+			finally:
+				if self.port:
+					self.port.close()
+
+	async def display_polling(self):
+		rx = 0
+		tx = 0
+		while self.run:
+			cpu_report = await self._loop.run_in_executor(None, psutil.cpu_percent, 1, True)
+			load = [ int(((x/100) **2) * 255) for x in cpu_report ] # mypy whines but percpu returns a list
+			net = psutil.net_io_counters(pernic=True)
+			d_rx = sum(v.bytes_recv for k, v in net.items() if k != "lo")
+			d_tx = sum(v.bytes_sent for k, v in net.items() if k != "lo")
+			await self.packets.put(struct.pack("BBBBBBBB", 3, 6, *load, min(int((d_tx - tx) / 1000), 255), min(int((d_rx - rx) / 1000), 255)))
+			rx = d_rx
+			tx = d_tx
+
+	async def cpu_load_leds(self):
+		while self.run:
+			cpu_report = await self._loop.run_in_executor(None, psutil.cpu_percent, 0.05, True)
+			load = [ int(((x/100) **2) * 255) for x in cpu_report ] # mypy whines but percpu returns a list
+			logging.info("CPU [%d|%d|%d|%d]", *load)
+			await self.packets.put(struct.pack("BBBBBB", 1, 4, *load))
+
+	async def net_traffic_leds(self):
+		rx = 0
+		tx = 0
+		while self.run:
+			net = psutil.net_io_counters(pernic=True)
+			d_rx = sum(v.bytes_recv for k, v in net.items() if k != "lo")
+			d_tx = sum(v.bytes_sent for k, v in net.items() if k != "lo")
+			logging.info("NET [TX %d | %d RX]", d_tx - tx, d_rx - rx)
+			await self.packets.put(struct.pack("BBBB", 2, 2, min(d_tx - tx, 255), min(d_rx - rx, 255)))
+			rx = d_rx
+			tx = d_tx
+			await asyncio.sleep(0.01)
+
+	async def run_tasks(self):
+		self._loop = asyncio.get_event_loop()
+		port_manager = self._loop.create_task(self.run_port_manager())
+		cpu_leds = self._loop.create_task(self.cpu_load_leds())
+		net_leds = self._loop.create_task(self.net_traffic_leds())
+		display = self._loop.create_task(self.display_polling())
+		await asyncio.gather(port_manager, cpu_leds, net_leds, display)
 
 if __name__ == "__main__":
 	if len(sys.argv) < 2:
-		print("[!] No device specified")
+		logging.error("[!] No device specified")
 		exit(-1)
-	
-	cpu_load_serial_driver(sys.argv[1])
 
+	logging.basicConfig(level=logging.WARNING)
+
+	state = State(sys.argv[1])
+
+	asyncio.run(state.run_tasks())
+
+