From 6f7d9d52430ab8a21e17e80afc8fbd20f8339f93 Mon Sep 17 00:00:00 2001
From: Bluefissure <bluefissure@foxmail.com>
Date: Tue, 21 Nov 2023 14:44:22 +0800
Subject: [PATCH] fix: resume in new ws

---
 qqbot/QQBot.py  | 55 +++++++++++++++++++++++++++++++------------------
 qqbot/consts.py |  4 ++++
 qqbot/main.py   | 13 +++---------
 3 files changed, 42 insertions(+), 30 deletions(-)

diff --git a/qqbot/QQBot.py b/qqbot/QQBot.py
index b42af84e..0af29a42 100644
--- a/qqbot/QQBot.py
+++ b/qqbot/QQBot.py
@@ -4,33 +4,51 @@
 import threading
 import time
 import requests
+import websockets
 import inspect
 from collections import defaultdict
-from consts import EVENT_INTENT
+from consts import EVENT_INTENT, ClientState
 
 class QQBot(object):
-    def __init__(self, config, ws) -> None:
+    def __init__(self, config) -> None:
         self.token = None
         self.expiration = 0
         self.config = config
         self.app_id = str(config['app_id'])
-        self.ws = ws
+        self.ws = None
         self._s = 0
         self.username = 'Unknown QQBot'
         self.bot = True
         self.version = None
         self.session_id = None
         self.id = None
-        self.logged_in = False
         self._log = logging.getLogger('QQBot')
         self.http = requests.Session()
         self._heartbeat_thread = None
         self._subscriptions = defaultdict(list)
+        self._state = ClientState.INIT
         self._refresh_token()
 
     def __str__(self) -> str:
         return f'QQBot {self.username}'
 
+    async def run(self):
+        reconnect_count = 0
+        while reconnect_count < 50:
+            try:
+                self._log.info('Connecting to QQ...')
+                async with websockets.connect("wss://api.sgroup.qq.com/websocket/") as websocket:
+                    self.ws = websocket
+                    async for message in websocket:
+                        jdata = json.loads(message)
+                        await self.handle(jdata)
+            except websockets.exceptions.ConnectionClosedError:
+                self._log.info('Connection closed, try reconnecting...')
+                reconnect_count += 1
+            except Exception as e:
+                raise e
+
+
     def _refresh_token(self):
         if time.time() >= self.expiration:
             self._log.info('Refreshing token...')
@@ -73,22 +91,17 @@ def headers(self) -> dict:
     def s(self) -> int:
         return self._s
 
-    @property
-    def log(self):
-        return self._log
-
     @s.setter
     def s(self, value: int):
         self._s = value
 
+    @property
+    def log(self):
+        return self._log
 
     async def _heartbeat(self):
         first_hearbeat = True
-        session_id = self.session_id
         while True:
-            if self.session_id != session_id:
-                self._log.debug('%s session id changed.', self)
-                break
             self._refresh_token()
             await self.ws.send(json.dumps({
                 "op": 1,
@@ -116,8 +129,8 @@ def _update_info(self, ready_data: dict):
         self.bot = user['bot']
         self.username = user['username']
 
-    def _logged_in(self, ready_data: dict):
-        self.logged_in = True
+    def _on_ready(self, ready_data: dict):
+        self._state = ClientState.READY
         self._update_info(ready_data)
         self._spawn_heartbeat()
 
@@ -190,22 +203,24 @@ def reply_channel_message(self, message:dict, content: str, image: str = None):
     async def handle(self, message: dict):
         op = message['op']
         s = message.get('s', -1)
-        if s > -1:
+        if self._state == ClientState.READY and s > -1:
             self.s = s
         if op == 10:
             self._log.debug('QQ says hello.')
-            if not self.logged_in:
+            if self._state == ClientState.INIT:
                 await self.try_login()
+            elif self._state == ClientState.RECONNECT:
+                await self.reconnect()
         elif op == 7:
             self._log.debug('QQ requires reconnect.')
-            await self.reconnect()
+            self._state = ClientState.RECONNECT
         elif op == 0:
             t = message['t']
             if t == 'READY':
-                if not self.logged_in:
-                    self._logged_in(message['d'])
-                    self._log.info('%s logged in.', self)
+                self._on_ready(message['d'])
+                self._log.info('%s is ready.', self)
             elif t == "RESUMED":
+                self._state = ClientState.READY
                 self._log.info('%s resumed.', self)
             elif t in self._subscriptions:
                 for func in self._subscriptions[t]:
diff --git a/qqbot/consts.py b/qqbot/consts.py
index e51bf146..8c8edcd6 100644
--- a/qqbot/consts.py
+++ b/qqbot/consts.py
@@ -52,3 +52,7 @@ class IntentOffset(Enum):
     for vv in v:
         EVENT_INTENT[vv] = k
 
+class ClientState(Enum):
+    INIT = 0
+    READY = 1
+    RECONNECT = 2
\ No newline at end of file
diff --git a/qqbot/main.py b/qqbot/main.py
index d87728fa..62fd38b0 100644
--- a/qqbot/main.py
+++ b/qqbot/main.py
@@ -22,20 +22,13 @@
 with open("config.yaml", "r", encoding='utf-8') as f:
     BOT_CONFIG = yaml.load(f, Loader=yaml.FullLoader)
 
-async def handler(websocket):
-    Q = QQBot(BOT_CONFIG, websocket)
-    Q.subscribe(events)
-    async for message in websocket:
-        jdata = json.loads(message)
-        await Q.handle(jdata)
-
-
 async def main():
     while True:
         try:
             _log.info("=== OtterBot QQ Bot v0.0.0.1 ===")
-            async with websockets.connect("wss://api.sgroup.qq.com/websocket/") as ws:
-                await handler(ws)
+            Q = QQBot(BOT_CONFIG)
+            Q.subscribe(events)
+            await Q.run()
         except Exception as e:
             _log.error(e)
             traceback.print_exc()