FIX: Fix permissions loss on community servers with tag plugins (#102)

* add tag removal for community servers

* add retrieval of host steamid3

* using steamid3 to check for permissions inside clear command

* use steamid64 instead of username for chats

* fix tests for util.text & commands.clear

* fix command controller tests

* fix winreg import

* refactor StatsData

* merge StatsData into LobbyManager

* remove ENABLE_STATS option

* update lobby_manager

* fix map property of lobby_manager

* refactor lobby manager

* status command triggers are now based on logs instead of time

* minor improvements

* fix logging for stats, add option to disabling stats

* fix tests

* add a few LobbyManager tests

* update bans, add bans tests

* fix bug with line parsing

* fix typos

* fix openai command prefix removal

* add send status command on start until response is received

* add more tests

Author: dborodin836

config.ini

@@ -94,8 +94,8 @@ TOS_VIOLATION = 0
 CUSTOM_PROMPT =
 
 [STATS]
-ENABLE_STATS=0
 STEAM_WEBAPI_KEY=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
+ENABLE_STATS_LOGS=False
 
 [CUSTOM-MODEL-GENERAL]
 ENABLE_SOFT_LIMIT_FOR_CUSTOM_MODEL=True
@@ -123,6 +123,11 @@ CUSTOM_MODEL_HOST=127.0.0.1:5000
 ; CUSTOM_MODEL_SETTINGS = {"echo": true, "max_tokens": 200, "temperature": 2.0, "stop": ["FOO", "BAR"], "logit_bias": {"50256": -100} }
 CUSTOM_MODEL_SETTINGS =
 
+[PERMISSIONS]
+; Set to True if the function should fall back to checking the username for admin privileges when the user's SteamID3
+; does not match the host's SteamID3.
+FALLBACK_TO_USERNAME=True
+
 [EXPERIMENTAL]
 ; Unlike the standard 'Fire-and-Forget' queue which sends messages without acknowledging
 ; their delivery, the confirmable queue adds reliability by ensuring that each message is

config.py

@@ -32,12 +32,13 @@ def has_value(cls, value):
 class Config(BaseModel):
     APP_VERSION: str = "1.3.0"
     HOST_USERNAME: str = ""
+    HOST_STEAMID3: str = "[U:X:XXXXXXX]"
     TOS_VIOLATION: bool
+    FALLBACK_TO_USERNAME: bool
 
     TF2_LOGFILE_PATH: str
     OPENAI_API_KEY: str
 
-    ENABLE_STATS: bool
     STEAM_WEBAPI_KEY: str
     DISABLE_KEYBOARD_BINDINGS: bool
     GPT4_COMMAND: str
@@ -55,6 +56,7 @@ class Config(BaseModel):
     RTD_COMMAND: str
     GLOBAL_CHAT_COMMAND: str
     GPT4_ADMIN_ONLY: bool
+    ENABLE_STATS_LOGS: bool
 
     CUSTOM_PROMPT: str
 
@@ -91,9 +93,6 @@ def api_key_pattern_match(cls, v):
 
     @validator("STEAM_WEBAPI_KEY")
     def steam_webapi_key_pattern_match(cls, v, values):
-        if not values["ENABLE_STATS"]:
-            return v
-
         if not re.fullmatch(WEB_API_KEY_RE_PATTERN, v):
             buffered_fail_message(
                 "STEAM WEB API key not set or invalid!", type_="BOTH", level="ERROR"

main.py

@@ -3,10 +3,8 @@
 # This is required due to config used in imported modules
 init_config()
 
-import contextlib
 import sys
 import threading
-import time
 import tkinter as tk
 
 from pynput import keyboard
@@ -18,24 +16,16 @@
 from modules.gui.log_window import LogWindow, RedirectStdoutToLogWindow
 from modules.logs import get_logger, setup_loggers
 from modules.message_queueing import message_queue_handler
-from modules.servers.tf2 import get_status
-from modules.tf_statistics import StatsData
+from modules.lobby_manager import lobby_manager
 
 gui_logger = get_logger("gui")
 
 
-def status_command_sender():
-    with contextlib.suppress(Exception):
-        while True:
-            get_status()
-            time.sleep(20)
-
-
 def keyboard_on_press(key):
     if key == keyboard.Key.f11:
         state_manager.switch_state()
     elif key == keyboard.Key.f10:
-        gui_logger.info(StatsData.get_data())
+        gui_logger.info(lobby_manager.get_data())
 
 
 def run_threads():
@@ -48,8 +38,6 @@ def run_threads():
 
     threading.Thread(target=parse_console_logs_and_build_conversation_history, daemon=True).start()
     threading.Thread(target=gpt3_cmd_handler, daemon=True).start()
-    if config.ENABLE_STATS:
-        threading.Thread(target=status_command_sender, daemon=True).start()
     threading.Thread(target=message_queue_handler, daemon=True).start()
     if not config.DISABLE_KEYBOARD_BINDINGS:
         keyboard.Listener(on_press=keyboard_on_press).start()

modules/bans.py

@@ -1,8 +1,10 @@
 import codecs
 import json
 from json import JSONDecodeError
+from typing import Set
 
 from modules.logs import get_logger, log_gui_general_message
+from modules.typing import Player
 
 main_logger = get_logger("main")
 gui_logger = get_logger("gui")
@@ -11,62 +13,63 @@
 
 class BansManager:
     def __init__(self, bans_file: str = None):
-        self.banned_usernames = set()
+        self.banned_players_steamid3: Set[str] = set()
         self.__bans_filename = bans_file or "bans.json"
         self.load_banned_players()
 
-    def load_banned_players(self):
+    def load_banned_players(self) -> None:
         """
         Loads the set of banned players.
         """
         try:
             with codecs.open(self.__bans_filename, "r", encoding="utf-8") as file:
-                self.banned_usernames = set(json.load(file))
-        except (EOFError, JSONDecodeError):
-            self.banned_usernames = set()
+                self.banned_players_steamid3 = set(json.load(file))
+        except (EOFError, JSONDecodeError) as e:
+            combo_logger.warning(f"Error parsing .json file [{e}].")
+            self.banned_players_steamid3 = set()
         except FileNotFoundError:
             combo_logger.warning(f"File {self.__bans_filename} could not be found.")
-            self.banned_usernames = set()
+            self.banned_players_steamid3 = set()
         except Exception as e:
             combo_logger.warning(f"Failed to load banned players. [{e}]")
-            self.banned_usernames = set()
+            self.banned_players_steamid3 = set()
 
-    def is_banned_username(self, username: str) -> bool:
+    def is_banned_player(self, player: Player) -> bool:
         """
         Checks if a given username is banned.
         """
-        main_logger.trace(f"Checking if '{username}' is banned.")
-        return username in self.banned_usernames
+        main_logger.trace(f"Checking if '{player.name}' [{player.steamid64}] is banned.")
+        return player.steamid3 in self.banned_players_steamid3
 
-    def ban_player(self, username: str) -> None:
+    def ban_player(self, player: Player) -> None:
         """
         Adds a player to the set of banned players.
         """
-        main_logger.debug(f"Trying to ban '{username}'")
-        if not self.is_banned_username(username):
-            self.banned_usernames.add(username)
-            log_gui_general_message(f"BANNED '{username}'")
-            main_logger.debug(f"Successfully banned {username}")
+        main_logger.debug(f"Trying to ban '{player.name}' [{player.steamid64}]")
+        if not self.is_banned_player(player):
+            self.banned_players_steamid3.add(player.steamid3)
+            log_gui_general_message(f"BANNED '{player.name}' [{player.steamid64}]")
+            main_logger.debug(f"Successfully banned {player.name} [{player.steamid64}]")
             with codecs.open(self.__bans_filename, "w", encoding="utf-8") as f:
-                json.dump(list(self.banned_usernames), f)
+                json.dump(list(self.banned_players_steamid3), f)
         else:
-            log_gui_general_message(f"='{username}' ALREADY BANNED")
-            main_logger.debug(f"{username} already banned.")
+            log_gui_general_message(f"='{player}' [{player.steamid64}] ALREADY BANNED")
+            main_logger.debug(f"{player.name} [{player.steamid64}] already banned.")
 
-    def unban_player(self, username: str) -> None:
+    def unban_player(self, player: Player) -> None:
         """
         Removes a player from the set of banned players.
         """
-        main_logger.debug(f"Trying to unban '{username}'")
-        if self.is_banned_username(username):
-            self.banned_usernames.remove(username)
-            log_gui_general_message(f"UNBANNED '{username}'")
-            main_logger.debug(f"Successfully unbanned {username}")
+        main_logger.debug(f"Trying to unban '{player.name}' [{player.steamid64}]")
+        if self.is_banned_player(player):
+            self.banned_players_steamid3.remove(player.steamid3)
+            log_gui_general_message(f"UNBANNED '{player.name}' [{player.steamid64}]")
+            main_logger.debug(f"Successfully unbanned {player.name} [{player.steamid64}]")
             with codecs.open(self.__bans_filename, "w", encoding="utf-8") as f:
-                json.dump(list(self.banned_usernames), f)
+                json.dump(list(self.banned_players_steamid3), f)
         else:
-            log_gui_general_message(f"USER '{username}' WAS NOT BANNED!")
-            main_logger.debug(f"{username} was not banned, cancelling.")
+            log_gui_general_message(f"USER '{player.name}' WAS NOT BANNED!")
+            main_logger.debug(f"{player.name} [{player.steamid64}] was not banned, cancelling.")
 
 
 bans_manager = BansManager()

modules/chat.py

@@ -13,6 +13,7 @@
 from modules.servers.tf2 import check_connection, set_host_username
 from modules.utils.buffered_messages import print_buffered_config_innit_messages
 from modules.utils.prompts import load_prompts
+from modules.utils.steam import set_host_steamid3
 from modules.utils.text import get_console_logline
 
 gui_logger = get_logger("gui")
@@ -37,6 +38,7 @@ def setup() -> None:
     check_for_updates()
     check_connection()
     set_host_username()
+    set_host_steamid3()
     load_prompts()
     print_buffered_config_innit_messages()
 
@@ -68,8 +70,12 @@ def parse_console_logs_and_build_conversation_history() -> None:
     controller.register_service(messaging_queue_service)
 
     for logline in get_console_logline():
+        if logline is None:
+            continue
         if not state_manager.bot_running:
             continue
-        if bans_manager.is_banned_username(logline.username):
+        if bans_manager.is_banned_player(logline.player):
+            main_logger.info(f"Player '{logline.player.name}' {logline.player.steamid3} tried to use commands, but "
+                             f"he's banned.")
             continue
         controller.process_line(logline)

modules/command_controllers.py

@@ -7,7 +7,7 @@
 from modules.conversation_history import ConversationHistory
 from modules.logs import get_logger
 from modules.set_once_dict import SetOnceDictionary
-from modules.typing import Command, LogLine
+from modules.typing import Command, LogLine, Player
 
 main_logger = get_logger("main")
 combo_logger = get_logger("combo")
@@ -17,23 +17,23 @@
 class ChatHistoryManager(BaseModel):
     GLOBAL = ConversationHistory()
 
-    def set_conversation_history_by_name(self, username: str, conv_history: ConversationHistory) -> None:
-        attr_name = self._get_conv_history_attr_name(username)
+    def set_conversation_history(self, player: Player, conv_history: ConversationHistory) -> None:
+        attr_name = self._get_conv_history_attr_name(player.steamid64)
 
         setattr(self, attr_name, conv_history)
 
-    def get_conversation_history_by_name(self, username: str) -> ConversationHistory:
-        attr_name = self._get_conv_history_attr_name(username)
+    def get_conversation_history(self, player: Player) -> ConversationHistory:
+        attr_name = self._get_conv_history_attr_name(player.steamid64)
 
         if hasattr(self, attr_name):
             return getattr(self, attr_name)
         else:
-            main_logger.info(f"Conversation history for username '{username}' doesn't exist. Creating...")
+            main_logger.info(f"Conversation history for user '{player.name}' [{player.steamid64}] doesn't exist. Creating...")
             setattr(self, attr_name, ConversationHistory())
             return getattr(self, attr_name)
 
-    def _get_conv_history_attr_name(self, username: str) -> str:
-        return f"USER_{username}_CH"
+    def _get_conv_history_attr_name(self, id64: int) -> str:
+        return f"USER_{id64}_CH"
 
     class Config(BaseConfig):
         extra = "allow"

modules/commands/clear_chat.py

@@ -1,4 +1,6 @@
 from modules.command_controllers import InitializerConfig
+from modules.lobby_manager import lobby_manager
+from modules.permissions import is_admin
 from modules.typing import LogLine
 from modules.utils.text import get_args
 from modules.logs import get_logger
@@ -9,14 +11,14 @@
 
 
 def handle_clear(logline: LogLine, shared_dict: InitializerConfig):
-    if logline.username == config.HOST_USERNAME:
+    if is_admin(logline.player):
         args = get_args(logline.prompt.removeprefix(config.CLEAR_CHAT_COMMAND).strip())
 
         if len(args) == 0:
             combo_logger.info(f"Clearing chat history for user '{logline.username}'.")
-            conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history_by_name(logline.username)
+            conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history(logline.player)
             conv_history.reset()
-            shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history_by_name(logline.username, conv_history)
+            shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history(logline.player, conv_history)
             return
 
         if r"\global" in args:
@@ -46,17 +48,25 @@ def handle_clear(logline: LogLine, shared_dict: InitializerConfig):
                 name = name.removeprefix("'")
                 name = name.removesuffix("'")
 
-                combo_logger.info(f"Clearing chat history for user '{name}'.")
-                conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history_by_name(name)
-                conv_history.reset()
-                shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history_by_name(name, conv_history)
+                player = lobby_manager.get_player_by_name(name)
+                if player is not None:
+                    combo_logger.info(f"Clearing chat history for user '{player.name}'.")
+                    conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history(player)
+                    conv_history.reset()
+                    shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history(player, conv_history)
+                else:
+                    combo_logger.info(f"Failed to find user with name: '{name}'.")
 
             except Exception as e:
                 main_logger.trace(f"Failed to parse arg in clear chat command. [{e}]")
                 continue
 
     else:
-        combo_logger.info(f"Clearing chat history for user '{logline.username}'.")
-        conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history_by_name(logline.username)
-        conv_history.reset()
-        shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history_by_name(logline.username, conv_history)
+        player = lobby_manager.get_player_by_name(logline.username)
+        if player is not None:
+            combo_logger.info(f"Clearing chat history for user '{player.name}'.")
+            conv_history = shared_dict.CHAT_CONVERSATION_HISTORY.get_conversation_history(player)
+            conv_history.reset()
+            shared_dict.CHAT_CONVERSATION_HISTORY.set_conversation_history(player, conv_history)
+        else:
+            combo_logger.info(f"Failed to find user with name: '{logline.username}'.")

modules/commands/gui/bans.py

@@ -1,21 +1,31 @@
 from modules.bans import bans_manager
+from modules.command_controllers import InitializerConfig
 from modules.logs import get_logger
+from modules.lobby_manager import lobby_manager
 
 gui_logger = get_logger("gui")
 
 
-def handle_ban(command, shared_dict):
+def handle_ban(command: str, shared_dict: InitializerConfig):
     username = command.removeprefix("ban ").strip()
-    bans_manager.ban_player(username)
+    player = lobby_manager.get_player_by_name(username)
+    if player is not None:
+        bans_manager.ban_player(player)
+        return None
+    gui_logger.info(f"Player '{username}' not found.")
 
 
-def handle_unban(command, shared_dict):
-    name = command.removeprefix("unban ").strip()
-    bans_manager.unban_player(name)
+def handle_unban(command: str, shared_dict: InitializerConfig):
+    username = command.removeprefix("unban ").strip()
+    player = lobby_manager.get_player_by_name(username)
+    if player is not None:
+        bans_manager.unban_player(player)
+        return None
+    gui_logger.info(f"Player '{username}' not found.")
 
 
-def handle_list_bans(command, shared_dict):
-    if len(bans_manager.banned_usernames) == 0:
+def handle_list_bans(command: str, shared_dict: InitializerConfig):
+    if len(bans_manager.banned_players_steamid3) == 0:
         gui_logger.info("### NO BANS ###")
     else:
-        gui_logger.info("### BANNED PLAYERS ###", *list(bans_manager.banned_usernames), sep="\n")
+        gui_logger.info("### BANNED PLAYERS ###", *list(bans_manager.banned_players_steamid3), sep="\n")