add authentication

add built-in cleanup
general refactor

cef_3M/util.py

@@ -2,6 +2,8 @@ import asyncio
 import hashlib
 import json
 import re
+import socket
+import traceback
 
 import MySQLdb
 
@@ -14,9 +16,9 @@ from .sql import SessionMaker, AlertEndpoints, ergoQueryFetchOne
 
 from fastapi import UploadFile
 
-
 safeName = re.compile(r"[^\w\d\.-]")
 
+
 # If this gets too out of hand, put an async breakpoint to allow other things to be handled while the hash occurs
 async def SHA256(f: UploadFile) -> str:
     sha = hashlib.sha256()
@@ -25,6 +27,7 @@ async def SHA256(f: UploadFile) -> str:
     await f.seek(0)
     return sha.hexdigest()
 
+
 minioClient = Minio(
     config.MINIO_INTERNAL_ADDR,
     secure=False,  # you will probably not have SSL
@@ -32,7 +35,8 @@ minioClient = Minio(
     secret_key=config.MINIO_SECRET_KEY,
 )
 
-redis = Redis(host='localhost', port=6379, db=0, protocol=3)
+redis = Redis(host=config.REDIS_ADDR, port=6379, db=0, protocol=3)
+
 
 class ErgoClient:
     def __init__(self):
@@ -40,34 +44,19 @@ class ErgoClient:
         self.writer = None
         asyncio.get_running_loop().create_task(self.init())
 
-    @staticmethod
-    def retry(f):
-        async def wrapper(self, *args, **kwargs):
-            i = 30
-            while i:
-                try:
-                    return await f(self, *args, **kwargs)
-                except RuntimeError:
-                    self.init()
-                    i -= 1
-                    print("Couldn't connect")
-        return wrapper
-
-    @retry
     async def init(self):
         self.reader, self.writer = await asyncio.open_connection(config.ERGO_ADDR, config.ERGO_PORT)
-        await asyncio.get_running_loop().create_task(self.readEvents())
+        await asyncio.create_task(self.readEvents())
 
-    @retry
     async def readEvents(self):
         while 1:
             rawLine = await self.reader.readline()
-            if not rawLine: break
+            if not rawLine:
+                break
             line = rawLine.decode("utf8").strip().split()
             if line[0] == "MENTION":
                 await self.handleMention(line[1], line[2], line[3])
 
-
     async def handleMention(self, username: str, channel: str, msgid: str):
         session = SessionMaker()
         for target in session.query(AlertEndpoints).filter(AlertEndpoints.username == username):
@@ -89,19 +78,28 @@ class ErgoClient:
             await pusher.send_async(encoded)
         session.close()
 
-
-    @retry
     async def write(self, msg):
-
-        self.writer.write(msg+b"\n")
+        if self.writer is None:
+            for _ in range(30):
+                await asyncio.sleep(1)
+                if self.writer:
+                    break
+        self.writer.write(msg + b"\n")
         await self.writer.drain()
 
     async def broadcastAs(self, user, *message):
         await self.write(f"BROADCASTAS {user} {' '.join(message)}".encode("utf8"))
 
+    async def fullyRemoveUser(self, user):
+        await self.write(f"FULLYREMOVE {user}".encode("utf8"))
 
     async def broadcastTo(self, user, *message):
         await self.write(f"BROADCASTTO {user} {' '.join(message)}".encode("utf8"))
 
 
 ergo = ErgoClient()
+
+privilegedIps = set()
+for host in config.PRIVILEGED_HOSTS:
+    for addr in [x[-1][0] for x in socket.getaddrinfo("localhost", 0)]:
+        privilegedIps.add(addr)