fix: crash

main.py

@@ -81,9 +81,28 @@ class History(logging.Handler):
         self.records.clear()
         return copy
 
+class ConnectionContext:
+    def __enter__(self):
+        self.healthcheck()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.healthcheck()
+        return True
+    
+    class Error(Exception):
+        pass
+    
+    @classmethod
+    def healthcheck(cls):
+        if connection.closed is not None:
+            raise cls.Error()
+
 def main(driver: WebDriver, logger = logging.getLogger('main')):
     try:
         http = PoolManager()
+        context = ConnectionContext()
+
         driver.get(str(Path('index.html').resolve()))
         endpoint = server.listeners[0]
         parameters = vars(args)
@@ -109,9 +128,10 @@ def main(driver: WebDriver, logger = logging.getLogger('main')):
         return True
     
     def sleep(seconds: float):
-        try: WebDriverWait(driver, seconds, seconds).until(lambda _: False)
-        except: pass
-        return True
+        with context:
+            try: WebDriverWait(driver, seconds, seconds).until(lambda _: False)
+            except: pass
+            return True
     
     def locate(selector, wait=True, condition=EC.visibility_of_element_located) -> WebElement:
         while True:
@@ -154,8 +174,12 @@ def main(driver: WebDriver, logger = logging.getLogger('main')):
 
     try:
         until(lambda x: 'loginProgress' in x.find_element(By.TAG_NAME, "body").get_attribute('class'), watch=False)
-        locate("input.account").send_keys(parameters.get('account'))
-        locate("input#password").send_keys(parameters.get('password'))
+        account = str(parameters.get('account'))
+        password = str(parameters.get('password'))
+        logger.info('Logging in as %s (%s)', account.split('@', 1).pop(0).capitalize(), account)
+
+        locate("input.account").send_keys(account)
+        locate("input#password").send_keys(password)
         click("input.agree-checkbox")
         click("button.login-btn")
     except Exception as e:
@@ -165,14 +189,14 @@ def main(driver: WebDriver, logger = logging.getLogger('main')):
     logger.info('Waiting for authentication to complete...')
     while True:
         try:
-            locate("#container", False)
+            locate("#container", wait=False)
             logger.info('Done')
             break
         except:
             sleep(3)
 
     ready = False
-    sidebar = locate(".new-layout-left", False)
+    sidebar = locate(".new-layout-left")
     driver.execute_script("arguments[0].remove();", sidebar)
 
     def handle(request: NetworkRequest):
@@ -288,7 +312,6 @@ async def handler(request: ws.WebSocketRequest, logger = logging.getLogger('webs
 
     if connection is None:
         connection = websocket
-        logger.info('Connection established')
     else:
         await websocket.aclose(code=1000, reason="Non-singular connection prohibited")
         return