Feat/Async-Websocket

sample.env → .env.example

@@ -2,6 +2,9 @@ HOST=misskey.example.com
 SECRET_TOKEN=misskey_token
 
 CONFIG_DIR=./config
+DB_TYPE=redis
+DB_URL=redis://localhost:6379
+
 RUN_SERVER=false
-DB_TYPE=pickle
-DB_URL=ssss
+SERVER_HOST=0.0.0.0
+SERVER_PORT=8000

.gitignore

@@ -1,12 +1,17 @@
+# venv
 env/
+
+# cache
 __pycache__
+.pytest_cache
 
+# private
 *.env
-.pytest_cache
-!sample.env
 data/users.pickle
+ngWords.txt
 
+# vscode
+.vscode/*
 !.vscode/settings.json
-.vscode/
 
-ngWords.txt
+*.rdb

.vscode/settings.json

@@ -1,6 +1,11 @@
 {
     "cSpell.words": [
-        "dotenv"
+        "coloredlogs",
+        "dotenv",
+        "levelname",
+        "misskey",
+        "renote",
+        "websockets"
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+pythonpath = "src"
+testpaths = ["tests"]

requirements.txt

@@ -1,23 +1,39 @@
+annotated-types==0.7.0
 blinker==1.8.2
 certifi==2024.6.2
+cffi==1.16.0
 charset-normalizer==3.3.2
 click==8.1.7
 colorama==0.4.6
 coloredlogs==15.0.1
+cryptography==42.0.8
 Flask==3.0.3
+hiredis==2.3.2
 humanfriendly==10.0
 idna==3.7
 iniconfig==2.0.0
 itsdangerous==2.2.0
 Jinja2==3.1.4
 MarkupSafe==2.1.5
+more-itertools==10.3.0
 packaging==24.0
 pluggy==1.5.0
+pycparser==2.22
+pydantic==2.7.3
+pydantic-settings==2.3.1
+pydantic_core==2.18.4
 pyreadline3==3.4.1
 pytest==8.2.2
 python-dotenv==1.0.1
+python-ulid==1.1.0
 redis==5.0.5
+redis-om==0.3.1
 requests==2.32.3
+types-cffi==1.16.0.20240331
+types-pyOpenSSL==24.1.0.20240425
+types-redis==4.6.0.20240425
+types-setuptools==70.0.0.20240524
+typing_extensions==4.12.2
 urllib3==2.2.1
-websocket-client==1.8.0
+websockets==12.0
 Werkzeug==3.0.3

src/emojis.py

@@ -1,6 +1,6 @@
-import json
 import random
-from typing import Any
+
+from utils import load_from_json_path
 
 
 class ConfigJsonError(Exception):
@@ -9,22 +9,22 @@ class ConfigJsonError(Exception):
 
 class EmojiSet:
     def __init__(self, data: str | dict) -> None:
-        if isinstance(data, str):
-            with open(data) as f:
-                loaded = json.load(f)
-        else:
-            loaded = data
+        loaded = load_from_json_path(data, dict)
         self._check_format(loaded)
 
         self.response_emojis = loaded["triggers"]
         self.others = loaded["others"]
 
-    def _check_format(self, json: Any) -> None:
+    def _check_format(self, json: dict) -> None:
         if not isinstance(json, dict) or sorted(json.keys()) != ["others", "triggers"]:
-            raise ConfigJsonError("response.jsonは{'triggers': [], 'others': []}の形にしてください。")
-
-        if any([tuple(i.keys()) != ("keywords", "emoji") for i in json["triggers"]]):
-            raise ConfigJsonError("response.jsonのトリガーのキーはkeywordsとemojiにしてください。")
+            raise ConfigJsonError(
+                "response.jsonは{'triggers': [], 'others': []}の形にしてください。"
+            )
+
+        if any([sorted(i.keys()) != ["emoji", "keywords"] for i in json["triggers"]]):
+            raise ConfigJsonError(
+                "response.jsonのトリガーのキーはkeywordsとemojiにしてください。"
+            )
 
     def get_response_emoji(self, text: str) -> str:
         for i in self.response_emojis:

src/environs.py

@@ -0,0 +1,33 @@
+from ipaddress import IPv4Address
+from pathlib import Path
+from typing import Literal, Optional
+
+from pydantic import DirectoryPath, Field, model_validator
+from pydantic.networks import IPvAnyAddress, RedisDsn
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class DotenvSettings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
+
+
+class Settings(DotenvSettings):
+    host: str = Field(default=...)
+    secret_token: str = Field(default=...)
+
+    db_type: Literal["redis", "pickle"] = "redis"
+    db_url: Optional[RedisDsn] = None
+
+    config_dir: DirectoryPath = Path("./config")
+    run_server: bool = False
+    server_host: IPvAnyAddress = IPv4Address("0.0.0.0")  # type: ignore
+    server_port: int = 8000
+
+    @model_validator(mode="after")
+    def validate_environ(self):
+        if self.db_url is None and self.db_type == "redis":
+            raise ValueError("DB_URL is required when DB_TYPE is redis")
+
+        if not (self.server_port > 0 and self.server_port < 65536):
+            raise ValueError("SERVER_PORT must be between 1 and 65535")
+        return self

src/keep_alive.py

@@ -1,34 +1,29 @@
-import logging
+import asyncio
 from multiprocessing import Process
-from os import getenv
 
-import coloredlogs
-from dotenv import load_dotenv
 from flask import Flask, jsonify
 
+import environs
 import logging_styles
 import mainbot
 
-
 app = Flask("app")
 
-logger = logging.getLogger(__name__)
-logging_styles.set_default()
-coloredlogs.install(logger=logger)
+logger = logging_styles.getLogger(__name__)
 
 
 @app.get("/")
 def pong():
     return jsonify({"message": "Pong!"})
 
 
-def run_server():
-    app.run(host="0.0.0.0", port=8080)
+def run_server(host: str, port: int):
+    app.run(host=host, port=port)
 
 
 if __name__ == "__main__":
-    load_dotenv()
-    if getenv("RUN_SERVER", False):
-        Process(target=run_server).start()
+    config = environs.Settings()
+    if config.run_server:
+        Process(target=run_server, args=(str(config.server_host), config.server_port)).start()
         logger.info("Web server started!")
-    mainbot.Bot().start_bot()
+    asyncio.run(mainbot.Bot(config).start_bot())

src/logging_styles.py

@@ -19,10 +19,17 @@
     "critical": {"color": "red"},
 }
 
+formatter = coloredlogs.ColoredFormatter(
+    fmt=DEFAULT_LOG_FORMAT,
+    datefmt=DEFAULT_DATE_FORMAT,
+    field_styles=DEFAULT_FIELD_STYLES,
+    level_styles=DEFAULT_LEVEL_STYLES,
+)
 
-def set_default():
-    coloredlogs.DEFAULT_LOG_LEVEL = DEFAULT_LOG_LEVEL
-    coloredlogs.DEFAULT_LOG_FORMAT = DEFAULT_LOG_FORMAT
-    coloredlogs.DEFAULT_DATE_FORMAT = DEFAULT_DATE_FORMAT
-    coloredlogs.DEFAULT_FIELD_STYLES = DEFAULT_FIELD_STYLES
-    coloredlogs.DEFAULT_LEVEL_STYLES = DEFAULT_LEVEL_STYLES
+
+def getLogger(name: str) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.setLevel(DEFAULT_LOG_LEVEL)
+    logger.addHandler(logging.StreamHandler())
+    logger.handlers[0].setFormatter(formatter)
+    return logger

src/mainbot.py

@@ -1,36 +1,34 @@
+import asyncio
 import json
-import logging
-import os
 from threading import Thread
 
-import coloredlogs
-import websocket
+import websockets
 
 import utils
 import logging_styles
 import misskey_api as misskey
+from environs import Settings
+from userdb import UserDB
 from ngwords import NGWords
 from emojis import EmojiSet
 
 
 class Bot:
     counter = utils.Counter(100, lambda: None)
 
-    def __init__(self, restart: bool = True) -> None:
-        logger = logging.getLogger(__name__)
-        logging_styles.set_default()
-        coloredlogs.install(logger=logger)
-        self.logger = logger
+    def __init__(self, settings: Settings, restart: bool = True) -> None:
+        self.logger = logging_styles.getLogger(__name__)
+        self.config = settings
         self._restart = restart
 
-        self.config_dir = utils.config_dir()
+        self.config_dir = self.config.config_dir
 
-        logger.info("Loading response.json...")
-        self.emojis = EmojiSet(os.path.join(self.config_dir, "response.json"))
-        logger.info("Loading ngwords.txt...")
-        self.ngw = NGWords(os.path.join(self.config_dir, "ngwords.txt"))
+        self.logger.info("Loading response.json...")
+        self.emojis = EmojiSet(str(self.config_dir.joinpath("response.json")))
+        self.logger.info("Loading ngwords.txt...")
+        self.ngw = NGWords(str(self.config_dir.joinpath("ngwords.txt")))
 
-        self.db = utils.get_db()
+        self.db = UserDB(str(self.config.db_url))  # TODO: redis以外への対応
 
     # TODO: なんか良い名前に変えたい
     def send_welcome(self, note_id: str, note_text: str) -> None:
@@ -50,10 +48,11 @@ def need(self) -> bool:
             return True
         return False
 
-    def on_message(self, ws, message: str) -> None:
+    async def on_message(self, ws, message: str) -> None:
         note_body = json.loads(message)["body"]["body"]
         note_id = note_body["id"]
         note_text = note_body["text"]
+        user_id = note_body["userId"]
         if note_text is None:
             note_text = ""
 
@@ -64,13 +63,15 @@ def on_message(self, ws, message: str) -> None:
         # Renote不可ならreturn
         return_flg = True
         if self.ngw.match(note_text):
-            self.logger.info(f"Detected NG word. | noteId: {note_id}, \
-                               word: {self.ngw.why(note_text)}")
+            self.logger.info(
+                f"Detected NG word. | noteId: {note_id}, \
+                               word: {self.ngw.why(note_text)}"
+            )
         elif misskey.can_reply(note_body):
             Thread(target=misskey.reply, args=(note_id, "Pong!")).start()
         elif not misskey.can_renote(note_body):
             pass
-        elif note_body["userId"] in set(self.db):
+        elif await self.db.get_user_by_id(user_id):
             self.logger.debug("Skiped api request because it was registered in DB.")
         else:
             return_flg = False
@@ -81,40 +82,53 @@ def on_message(self, ws, message: str) -> None:
         self.logger.debug(
             f"Notes not registered in database. | body: {note_text} , id: {note_id}"
         )
-        user_info = misskey.get_user_info(user_id=note_body["userId"])
+        user_info = misskey.get_user_info(user_id=user_id)
 
         if (notes_count := user_info["notesCount"]) == 1:
             self.send_welcome(note_id, note_text)
         elif notes_count <= 10:  # ノート数が10以下ならRenote出来る可能性
-            notes = misskey.get_user_notes(note_body["userId"], note_id, 10)
+            notes = misskey.get_user_notes(user_id, note_id, 10)
             if all([not misskey.can_renote(note) for note in notes]):
                 self.send_welcome(note_id, note_text)
                 return None
 
         if notes_count > 5:
-            self.db.append(note_body["userId"])
-            if (count := len(self.db)) % 100 == 0:
-                utils.update_db("have_note_user_ids", self.db, False)
-                self.logger.info(f"DataBase Updated. | length: {count}")
+            await self.db.add_user(user_id, note_body["user"]["username"])
+            self.logger.info("DataBase Updated.")
 
-    def on_error(self, ws, error) -> None:
+    async def on_error(self, ws, error) -> None:
         self.logger.warning(str(error))
 
-    def on_close(self, ws, status_code, msg) -> None:
+    async def on_close(self, ws, status_code, msg) -> bool:
         self.logger.error(f"WebSocket closed. | code:{status_code} msg: {msg}")
-        if self._restart:
-            self.start_bot()
+        return self._restart
 
-    def start_bot(self):
+    async def start_bot(self):
         streaming_api = f"wss://{misskey.HOST}/streaming?i={misskey.TOKEN}"
         USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36"  # NOQA
-        MESSAGE = {"type": "connect", "body": {"channel": "hybridTimeline", "id": "1"}}
-        # WebSocketの接続
-        ws = websocket.WebSocketApp(
-            streaming_api,
-            on_message=self.on_message, on_error=self.on_error, on_close=self.on_close,
-            header={"User-Agent": USER_AGENT}
-        )
-        ws.on_open = lambda ws: ws.send(json.dumps(MESSAGE))
-        self.logger.info("Bot was started!")
-        ws.run_forever()
+        CONNECTMSG = {
+            "type": "connect",
+            "body": {"channel": "hybridTimeline", "id": "1"},
+        }
+
+        pong = await self.db.ping()
+        if not pong:
+            raise Exception("DB connection failed.")
+
+        while True:
+            async with websockets.connect(streaming_api, user_agent_header=USER_AGENT) as ws:
+                # self.on_open(ws)
+                self.logger.info("Bot was started!")
+                await ws.send(json.dumps(CONNECTMSG))
+                while True:
+                    try:
+                        msg = await ws.recv()
+                        await self.on_message(ws, str(msg))
+                    except websockets.ConnectionClosed:
+                        await self.on_close(ws, ws.close_code, ws.close_reason)
+                        if not self._restart:
+                            return
+                        break
+                    except Exception as e:
+                        await self.on_error(ws, e)
+            await asyncio.sleep(5)