From da637a2ba3258ad19695d1e61a994ea4c5ecb6aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A4=9C=E5=9D=82=E9=9B=85?= <23130178+ShadowRZ@users.noreply.github.com> Date: Wed, 1 Feb 2023 20:48:51 +0800 Subject: [PATCH] Switch to TOML --- .gitignore | 2 +- SETUP.md | 24 +++++----- nyx_bot/callbacks.py | 18 ++------ nyx_bot/config.py | 42 ++++++++++++----- nyx_bot/main.py | 7 +-- nyx_bot/message_responses.py | 11 ++--- nyx_bot/storage.py | 2 +- nyx_bot/utils.py | 22 +++++++-- pyproject.toml | 1 + requirements.txt | 3 +- sample.config.yaml => sample.config.toml | 59 +++++++++++++----------- 11 files changed, 107 insertions(+), 84 deletions(-) rename sample.config.yaml => sample.config.toml (62%) diff --git a/.gitignore b/.gitignore index a5b429c..55dae58 100644 --- a/.gitignore +++ b/.gitignore @@ -20,7 +20,7 @@ build/ dist/ # Config file -config.yaml +nyx_bot.toml # Log files *.log diff --git a/SETUP.md b/SETUP.md index aa2c029..c9bf54e 100644 --- a/SETUP.md +++ b/SETUP.md @@ -43,7 +43,7 @@ You can install [libolm](https://gitlab.matrix.org/matrix-org/olm) from source, or alternatively, check your system's package manager. Version `3.0.0` or greater is required. -**(Optional) postgres development headers** +#### (Optional) postgres development headers By default, the bot uses SQLite as its storage backend. This is fine for a few hundred users, but if you plan to support a much higher volume of requests, you @@ -100,6 +100,8 @@ command to install postgres dependencies alongside those that are necessary: ``` pip install -e ".[postgres]" +# Using requirements.txt +pip install -r requirements-postgres.txt ``` ## Install fonts @@ -121,30 +123,30 @@ You can also use the Python one located in `cutword/nyx_bot-cutword.py`. ## Configuration -Copy the sample configuration file to a new `config.yaml` file. +Copy the sample configuration file to a new `nyx_bot.toml` file. ``` -cp sample.config.yaml config.yaml +cp sample.config.toml nyx_bot.toml ``` Edit the config file. The `matrix` section must be modified at least. -#### (Optional) Set up a Postgres database +### (Optional) Set up a Postgres database -Create a postgres user and database for matrix-reminder-bot: +Create a postgres user and database for it: ``` -sudo -u postgresql psql createuser nio-template -W # prompts for a password -sudo -u postgresql psql createdb -O nio-template nio-template +sudo -u postgresql psql createuser nyx_bot -W # prompts for a password +sudo -u postgresql psql createdb -O nyx_bot nyx_bot ``` Edit the `storage.database` config option, replacing the `sqlite://...` string with `postgres://...`. The syntax is: ``` -database: "postgres://username:password@localhost/dbname?sslmode=disable" +database = "postgres://username:password@localhost/dbname?sslmode=disable" ``` -See also the comments in `sample.config.yaml`. +See also the comments in `sample.config.toml`. ## Running @@ -171,11 +173,11 @@ it comes time to modifying the code for your own purposes, you are expected to replace every instance of "nyx-bot" and its variances with your own project's name. -By default, the bot will run with the config file at `./config.yaml`. However, an +By default, the bot will run with the config file at `./nyx_bot.toml`. However, an alternative relative or absolute filepath can be specified after the command: ``` -nyx-bot other-config.yaml +nyx-bot other-config.toml ``` ## Final steps diff --git a/nyx_bot/callbacks.py b/nyx_bot/callbacks.py index 489ab8e..4c230dd 100644 --- a/nyx_bot/callbacks.py +++ b/nyx_bot/callbacks.py @@ -22,6 +22,7 @@ from nyx_bot.utils import ( get_reply_to, is_bot_event, make_datetime, + should_record_message_content, strip_beginning_quote, ) @@ -38,9 +39,7 @@ class Callbacks: """ self.client = client self.config = config - self.disable_jerryxiao_for = config.disable_jerryxiao_for - self.disable_randomdraw_for = config.disable_randomdraw_for - self.record_message_content_for = config.record_message_content_for + self.room_features = config.room_features self.command_prefix = config.command_prefix self.replace_map = {} @@ -71,7 +70,7 @@ class Callbacks: # Record this message. timestamp = make_datetime(event.server_timestamp) external_url = get_external_url(event) - if room.room_id not in self.record_message_content_for: + if not should_record_message_content(self.room_features, room.room_id): include_text = False MatrixMessage.update_message( room, event, external_url, timestamp, event_replace, include_text @@ -98,21 +97,14 @@ class Callbacks: # Record this message. timestamp = make_datetime(event.server_timestamp) external_url = get_external_url(event) - if room.room_id not in self.record_message_content_for: + if not should_record_message_content(self.room_features, room.room_id): include_text = False MatrixMessage.update_message( room, event, external_url, timestamp, event_replace, include_text ) # General message listener message = Message( - self.client, - self.config, - msg, - room, - event, - reply_to, - self.disable_jerryxiao_for, - self.disable_randomdraw_for, + self.client, self.config, msg, room, event, reply_to, self.room_features ) try: await message.process() diff --git a/nyx_bot/config.py b/nyx_bot/config.py index 7d1fba2..6ab5f00 100644 --- a/nyx_bot/config.py +++ b/nyx_bot/config.py @@ -2,13 +2,17 @@ import logging import os import re import sys +from collections import defaultdict from typing import Any, List, Optional -import yaml +try: + import tomllib +except ModuleNotFoundError: + import tomli as tomllib from nyx_bot.errors import ConfigError -logger = logging.getLogger() +logger = logging.getLogger(__name__) logging.getLogger("peewee").setLevel( logging.INFO ) # Prevent debug messages from peewee lib @@ -24,7 +28,7 @@ class Config: # Load in the config file at the given filepath with open(filepath) as file_stream: - self.config_dict = yaml.safe_load(file_stream.read()) + self.config_dict = tomllib.loads(file_stream.read()) # Parse and validate config options self._parse_config_values() @@ -108,15 +112,29 @@ class Config: self.command_prefix = self._get_cfg(["command_prefix"], default="!c") + " " - self.disable_jerryxiao_for = self._get_cfg( - ["disable_jerryxiao_for"], [], required=False - ) - self.disable_randomdraw_for = self._get_cfg( - ["disable_randomdraw_for"], [], required=False - ) - self.record_message_content_for = self._get_cfg( - ["record_message_content_for"], [], required=False - ) + room_features_dict = self._get_cfg(["room_features"]) + room_features_default = { + "jerryxiao": False, + "randomdraw": False, + "record_messages": False, + } + if "jerryxiao" in room_features_dict: + room_features_default["jerryxiao"] = room_features_dict["jerryxiao"] + del room_features_dict["jerryxiao"] + if "randomdraw" in room_features_dict: + room_features_default["randomdraw"] = room_features_dict["randomdraw"] + del room_features_dict["randomdraw"] + if "record_messages" in room_features_dict: + room_features_default["record_messages"] = room_features_dict[ + "record_messages" + ] + del room_features_dict["record_messages"] + self.room_features = defaultdict(lambda: room_features_default) + for k, v in room_features_dict.items(): + if isinstance(v, dict): + features_override = v + self.room_features[k] = room_features_default | features_override + self.encryption = self._get_cfg(["encryption"], False, required=False) def _get_cfg( diff --git a/nyx_bot/main.py b/nyx_bot/main.py index 1da887c..59b6100 100644 --- a/nyx_bot/main.py +++ b/nyx_bot/main.py @@ -1,5 +1,4 @@ #!/usr/bin/env python3 -import asyncio import logging import os.path import sys @@ -44,7 +43,7 @@ async def main(): if len(sys.argv) > 1: config_path = sys.argv[1] else: - config_path = "config.yaml" + config_path = "nyx_bot.toml" # Read the parsed config file and create a Config object config = Config(config_path) @@ -140,7 +139,3 @@ async def main(): finally: # Make sure to close the client connection on disconnect await client.close() - - -# Run the main function in an asyncio event loop -asyncio.get_event_loop().run_until_complete(main()) diff --git a/nyx_bot/message_responses.py b/nyx_bot/message_responses.py index 885640f..72b29e2 100644 --- a/nyx_bot/message_responses.py +++ b/nyx_bot/message_responses.py @@ -8,6 +8,7 @@ from nyx_bot.chat_functions import gen_result_randomdraw, send_text_to_room from nyx_bot.config import Config from nyx_bot.jerryxiao import send_jerryxiao from nyx_bot.trpg_dicer import get_trpg_dice_result +from nyx_bot.utils import should_enable_jerryxiao, should_enable_randomdraw logger = logging.getLogger(__name__) @@ -21,8 +22,7 @@ class Message: room: MatrixRoom, event: RoomMessageText, reply_to: str, - disable_jerryxiao_for, - disable_randomdraw_for, + room_features, ): """Initialize a new Message @@ -43,8 +43,7 @@ class Message: self.room = room self.event = event self.reply_to = reply_to - self.disable_jerryxiao_for = disable_jerryxiao_for - self.disable_randomdraw_for = disable_randomdraw_for + self.room_features = room_features async def process(self) -> None: """Process and possibly respond to the message""" @@ -61,7 +60,7 @@ class Message: await self._trpg_dicer(query) async def _randomdraw(self, query: str, prob: bool) -> None: - if self.room.room_id in self.disable_randomdraw_for: + if not should_enable_randomdraw(self.room_features, self.room.room_id): return sender = self.event.sender msg = gen_result_randomdraw( @@ -95,7 +94,7 @@ class Message: async def _jerryxiao(self) -> None: """Performs features similar to the bot created by Jerry Xiao""" - if self.room.room_id in self.disable_jerryxiao_for: + if not should_enable_jerryxiao(self.room_features, self.room.room_id): return msg = self.message_content diff --git a/nyx_bot/storage.py b/nyx_bot/storage.py index d895590..3df2607 100644 --- a/nyx_bot/storage.py +++ b/nyx_bot/storage.py @@ -39,7 +39,7 @@ class MatrixMessage(Model): external_url: Optional[str], timestamp: datetime, event_replace: Optional[str] = None, - include_text: Optional[str] = False, + include_text: Optional[bool] = False, ): message_db_item = MatrixMessage.get_or_none( (MatrixMessage.room_id == room.room_id) diff --git a/nyx_bot/utils.py b/nyx_bot/utils.py index 2021347..5fbed8d 100644 --- a/nyx_bot/utils.py +++ b/nyx_bot/utils.py @@ -3,7 +3,7 @@ from datetime import datetime from html.parser import HTMLParser from io import BytesIO, StringIO from random import Random -from typing import Optional, Tuple +from typing import Dict, Optional, Tuple from urllib.parse import unquote, urlparse from zlib import crc32 @@ -56,7 +56,7 @@ async def get_body( async def get_formatted_body( - client: AsyncClient, room: MatrixRoom, event_id: str, replace_map: str + client: AsyncClient, room: MatrixRoom, event_id: str, replace_map: Dict[str, str] ) -> Optional[str]: if event_id not in replace_map: target_response = await client.room_get_event(room.room_id, event_id) @@ -143,12 +143,12 @@ def parse_matrixdotto_link(link: str): # Named Room type_ = "room_named" identifier = f"#{parsed.fragment}" - return (type_, identifier, None) + return type_, identifier, None elif len(paths) == 3: # Must be an event ID room = unquote(paths[1]) event_id = unquote(paths[2]) - return ("event", room, event_id) + return "event", room, event_id async def make_single_quote_image( @@ -271,4 +271,16 @@ async def parse_wordcloud_args( else: raise NyxBotRuntimeError("Argument is not valid.") - return (sender, days) + return sender, days + + +def should_record_message_content(room_features, room_id: str) -> bool: + return room_features[room_id]["record_messages"] + + +def should_enable_jerryxiao(room_features, room_id: str) -> bool: + return room_features[room_id]["jerryxiao"] + + +def should_enable_randomdraw(room_features, room_id: str) -> bool: + return room_features[room_id]["randomdraw"] diff --git a/pyproject.toml b/pyproject.toml index a191718..4f20edc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "matrix-nio>=0.10.0", "Markdown>=3.1.1", "PyYAML>=5.1.2", + "tomli >= 1.1.0; python_version < \"3.11\"", "Wand", "python-magic", "peewee", diff --git a/requirements.txt b/requirements.txt index 1026ea8..233a48f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ -matrix-nio[e2e]>=0.10.0 +matrix-nio>=0.10.0 Markdown>=3.1.1 PyYAML>=5.1.2 +tomli>=1.1.0; python_version < "3.11" Wand python-magic peewee diff --git a/sample.config.yaml b/sample.config.toml similarity index 62% rename from sample.config.yaml rename to sample.config.toml index defa90e..d324658 100644 --- a/sample.config.yaml +++ b/sample.config.toml @@ -3,62 +3,65 @@ # Default values are shown # The string to prefix messages with to talk to the bot in group chats -command_prefix: "!c" +command_prefix= "!c" # Enable E2EE room support ? -encryption: False +encryption = false # Options for connecting to the bot's Matrix account -matrix: +[matrix] # The Matrix User ID of the bot account - user_id: "@bot:example.com" + user_id = "@bot:example.com" # Matrix account password (optional if access token used) - user_password: "" + user_password = "" # Matrix account access token (optional if password used) # Use the provided get_access_token.py script to get one - #user_token: "" + #user_token = "" # The URL of the homeserver to connect to - homeserver_url: https://example.com + homeserver_url = "https://example.com" # The device ID that is **non pre-existing** device # If this device ID already exists, messages will be dropped silently in encrypted rooms - device_id: ABCDEFGHIJ + device_id = "ABCDEFGHIJ" # What to name the logged in device - device_name: my-project-name + device_name = "my-project-name" -storage: +[storage] # The database connection string # For SQLite3, this would look like: # database: "sqlite://bot.db" # For Postgres, this would look like: # database: "postgres://username:password@localhost/dbname?sslmode=disable" - database: "sqlite://bot.db" + database = "sqlite://bot.db" # The path to a directory for internal bot storage # containing encryption keys, sync tokens, etc. - store_path: "./store" + store_path = "./store" # Logging setup -logging: +[logging] # Logging level # Allowed levels are 'INFO', 'WARNING', 'ERROR', 'DEBUG' where DEBUG is most verbose - level: INFO + level = "INFO" # Configure logging to a file - file_logging: + [logging.file_logging] # Whether logging to a file is enabled - enabled: false + enabled = false # The path to the file to log to. May be relative or absolute - filepath: bot.log + filepath = "bot.log" # Configure logging to the console output - console_logging: + [logging.console_logging] # Whether logging to the console is enabled - enabled: true + enabled = true -# Disable Jerry Xiao like feature for the following room ID -disable_jerryxiao_for: - - '!XXXXXXXXXXXXXXXXXX:example.com' +# Room features switch. +# These are the default that can be overriden by subkeys. +[room_features] +jerryxiao = false # Controls Jerry Xiao like feature. +randomdraw = false +record_messages = false # Controls recording message content. -disable_randomdraw_for: - - '!XXXXXXXXXXXXXXXXXX:example.com' - -# Record message content for the following rooms -record_message_content_for: - - '!XXXXXXXXXXXXXXXXXX:example.com' +# # Toogle this room's room features. +# # You don't need to specify all of them. +# [room_features."!XXXXXXXXXXXXXXXXXX:example.com"] +# jerryxiao = false +# randomdraw = false +# record_messages = false \ No newline at end of file