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
This commit is contained in:
əlemi 2022-08-13 04:45:31 +02:00
parent e0aeb920d0
commit 4020bfeddb
No known key found for this signature in database
GPG key ID: BBCBFE5D7244634E

View file

@ -1,52 +1,98 @@
#!/usr/bin/env python #!/usr/bin/env python
import asyncio
import sys import sys
import struct import struct
import logging
from time import sleep from time import sleep
import serial import serial
import psutil import psutil
def cpu_load_serial_driver(device:str, retry_interval:float=5.0): class State:
while True: run: bool
try: device: str
port = serial.Serial(device, baudrate=57600) baudrate: int
avg_usage_to_serial(port) loop: asyncio.AbstractEventLoop
except serial.SerialException as e: packets: asyncio.Queue
print(f"[!] Could not connect to device: {str(e)}", file=sys.stderr) _port: serial.Serial
else: _task: asyncio.Task
port.close()
sleep(retry_interval)
def avg_usage_to_serial(port:serial.Serial): def __init__(self, device:str="/dev/ttyUSB0", baudrate=57600):
net = psutil.net_io_counters() self.run = True
net_tx = net.bytes_sent self.baudrate = baudrate
net_rx = net.bytes_recv self.device = device
port.write(struct.pack("BB", 0, 0)) self.packets = asyncio.Queue()
port.flush() self.port = None
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
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 __name__ == "__main__":
if len(sys.argv) < 2: if len(sys.argv) < 2:
print("[!] No device specified") logging.error("[!] No device specified")
exit(-1) exit(-1)
cpu_load_serial_driver(sys.argv[1]) logging.basicConfig(level=logging.WARNING)
state = State(sys.argv[1])
asyncio.run(state.run_tasks())