From e07591b38b5b7793d09705322098396a8dfa960b Mon Sep 17 00:00:00 2001
From: agatha <agatha@juggalol.com>
Date: Sun, 15 Sep 2024 18:02:25 -0400
Subject: [PATCH] use matrix bot

---
 .gitignore        |  2 +-
 matrix.py         | 89 +++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt  |  2 ++
 stockbot-buyvm.py | 62 ++++++++++++++++++++-------------
 4 files changed, 129 insertions(+), 26 deletions(-)
 create mode 100644 matrix.py

diff --git a/.gitignore b/.gitignore
index c1c906f..ab45957 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,4 +4,4 @@ venv/
 __pycache__/
 *.py[cod]
 
-config.py
+config.json
diff --git a/matrix.py b/matrix.py
new file mode 100644
index 0000000..28286a1
--- /dev/null
+++ b/matrix.py
@@ -0,0 +1,89 @@
+import markdown
+from loguru import logger
+from nio import AsyncClient, LoginResponse
+
+
+class MatrixBot:
+    def __init__(self, config: dict):
+        self.config = config
+
+        self.client = AsyncClient(
+            homeserver=self.config['homeserver'],
+            user=self.config['username']
+        )
+        self.logged_in = False
+
+    async def ensure_logged_in(self):
+        if not self.logged_in:
+            try:
+                response = await self.client.login(password=self.config['password'])
+                if isinstance(response, LoginResponse):
+                    self.logged_in = True
+                    logger.info(f"Logged in as {self.config['username']}")
+                else:
+                    logger.error(f"Failed to login as {self.config['username']}: {response}")
+                    logger.error("Closing nio session")
+                    await self.client.close()
+            except Exception as e:
+                logger.error(f"Exception during login: {e}")
+                await self.client.close()
+                raise
+
+    async def send_message(self, message: str):
+        await self.ensure_logged_in()
+
+        if not self.logged_in:
+            logger.error("Unable to send message, login failed")
+            return
+
+        try:
+            await self.client.room_send(
+                room_id=self.config['room_id'],
+                message_type="m.room.message",
+                content={
+                    "msgtype": "m.text",
+                    "body": message
+                }
+            )
+            logger.info("Message sent")
+        except Exception as e:
+            logger.error(f"Exception during sending message: {e}")
+            raise
+
+    async def send_markdown(self, message: str):
+        await self.ensure_logged_in()
+
+        if not self.logged_in:
+            logger.error("Unable to send message, login failed")
+            return
+
+        try:
+            # Convert message to markdown
+            html = markdown.markdown(message)
+
+            # Send markdown formatted message
+            await self.client.room_send(
+                room_id=self.config['room_id'],
+                message_type="m.room.message",
+                content={
+                    "msgtype": "m.text",
+                    "body": message,
+                    "format": "org.matrix.custom.html",
+                    "formatted_body": html
+                }
+            )
+            logger.info("Markdown message sent")
+        except Exception as e:
+            logger.error(f"Exception during sending markdown message: {e}")
+            raise
+
+    async def close(self):
+        if self.logged_in:
+            try:
+                await self.client.logout()
+                self.logged_in = False
+                logger.info(f"Logged out from {self.config['homeserver']}")
+            except Exception as e:
+                logger.error(f"Exception during logout: {e}")
+            finally:
+                await self.client.close()  # Ensure the client is closed
diff --git a/requirements.txt b/requirements.txt
index 8e0da18..c4a916b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,5 @@
 beautifulsoup4
 requests
 loguru
+matrix-nio
+markdown
diff --git a/stockbot-buyvm.py b/stockbot-buyvm.py
index 95c185b..4c07fdc 100644
--- a/stockbot-buyvm.py
+++ b/stockbot-buyvm.py
@@ -1,9 +1,10 @@
 """buyvm stock checker"""
+import json
+import asyncio
 import requests
 from bs4 import BeautifulSoup
 from loguru import logger
-
-from config import DISCORD_WEBHOOK
+from matrix import MatrixBot
 
 BASE_URL = 'https://my.frantech.ca/'
 URLS = [
@@ -14,13 +15,6 @@ URLS = [
 ]
 
 
-def send_notification(payload):
-    try:
-        requests.post(DISCORD_WEBHOOK, json=payload)
-    except requests.RequestException as e:
-        logger.error(f'error sending notification: {str(e)}')
-
-
 def get_url(url):
     try:
         response = requests.get(url)
@@ -58,8 +52,20 @@ def get_packages(html):
     return packages
 
 
-def main():
+def load_config(filename):
+    with open(filename) as f:
+        return json.load(f)
+
+
+async def main():
     logger.info('checking buyvm stocks')
+
+    # load configuration
+    config = load_config('config.json')
+
+    # initialize bot
+    bot = MatrixBot(config['matrix'])
+
     for url in URLS:
         html = get_url(url)
 
@@ -70,21 +76,27 @@ def main():
         for package in packages:
             if package['qty'] > 0:
                 logger.info(f"{package['name']}: {package['qty']} in stock")
-                send_notification({
-                    "username": "stockbot-buyvm",
-                    "embeds": [
-                        {
-                            "author": {
-                                "name": "BuyVM",
-                            },
-                            "title": package['name'],
-                            "url": package['url'],
-                            "description": f"{package['qty']} in stock now!"
-                        }
-                    ],
-                    "content": "STOCK ALERT"
-                })
+                await bot.send_message(f"{package['name']}: {package['qty']} in stock")
+
+    await bot.close()
+
+
+def main_with_shutdown():
+    loop = asyncio.get_event_loop()
+    main_task = loop.create_task(main())
+
+    try:
+        loop.run_until_complete(main_task)
+    except asyncio.CancelledError:
+        logger.info("Main task has been cancelled.")
+    finally:
+        pending_tasks = [t for t in asyncio.all_tasks(loop) if not t.done()]
+        if pending_tasks:
+            loop.run_until_complete(asyncio.gather(*pending_tasks, return_exceptions=True))
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.close()
+        logger.info("Event loop closed.")
 
 
 if __name__ == '__main__':
-    main()
+    main_with_shutdown()
\ No newline at end of file