From 781b36cc2e6bbefe2601bfcb0c18f6638d620d80 Mon Sep 17 00:00:00 2001
From: Alex Manning <mail@alex-m.co.uk>
Date: Sat, 12 Feb 2022 10:41:07 +0000
Subject: [PATCH] Add memory and disk space.

---
 load.py     | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++
 maubot.yaml |  6 ++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/load.py b/load.py
index 0f55f02..7a285d4 100644
--- a/load.py
+++ b/load.py
@@ -1,10 +1,23 @@
+import json
 import os
 
+import psutil
+import psutil._common as psc
 from maubot import MessageEvent, Plugin
 from maubot.handlers import command
 from mautrix.types import Format, MessageType, TextMessageEventContent
 
 
+def to_human(data):
+    output = {}
+    for name in data._fields:
+        value = getattr(data, name)
+        if name != "percent":
+            value = psc.bytes2human(value)
+        output[name] = value
+    return output
+
+
 class LoadBot(Plugin):
     @command.new("load", help="Get load average.")
     @command.argument("message", pass_raw=True)
@@ -17,3 +30,41 @@ class LoadBot(Plugin):
             formatted_body=f"{load}",
         )
         await evt.respond(content)
+
+    @command.new("mem", help="Get system memory usage.")
+    @command.argument("message", pass_raw=True)
+    async def memory_handler(self, evt: MessageEvent, message: str) -> None:
+        rawmem = psutil.virtual_memory()
+        rawswap = psutil.swap_memory()
+        mem = {}
+
+        mem["virt"] = to_human(rawmem)
+        mem["swap"] = to_human(rawswap)
+
+        mem = json.dumps(mem, indent=2)
+        content = TextMessageEventContent(
+            msgtype=MessageType.NOTICE,
+            format=Format.HTML,
+            body=f"{mem}",
+            formatted_body=f"{mem}",
+        )
+        await evt.respond(content)
+
+    @command.new("disk", help="Get system disk usage.")
+    @command.argument("message", pass_raw=True)
+    async def disk_handler(self, evt: MessageEvent, message: str) -> None:
+        result = {}
+        parts = psutil.disk_partitions()
+
+        for part in parts:
+            usage = psutil.disk_usage(part.mountpoint)
+            result[str(part.mountpoint)] = to_human(usage)
+
+        result = json.dumps(result, indent=2)
+        content = TextMessageEventContent(
+            msgtype=MessageType.NOTICE,
+            format=Format.HTML,
+            body=f"{result}",
+            formatted_body=f"{result}",
+        )
+        await evt.respond(content)
diff --git a/maubot.yaml b/maubot.yaml
index 88a9a76..1e776cb 100644
--- a/maubot.yaml
+++ b/maubot.yaml
@@ -1,7 +1,9 @@
 maubot: 0.1.0
 id: uk.a09.maubot.load
-version: 0.0.1
-license: MIT
+version: 0.0.2
+license: AGPL-3.0-or-later
 modules:
   - load
+dependencies:
+  - psutil
 main_class: LoadBot