import configparser import importlib import os from collections import defaultdict from django.conf import settings from django.db.models import F, Sum from django.db.models.functions import Length from zerver.models import BotConfigData, UserProfile class ConfigError(Exception): pass def get_bot_config(bot_profile: UserProfile) -> dict[str, str]: entries = BotConfigData.objects.filter(bot_profile=bot_profile) if not entries: raise ConfigError("No config data available.") return {entry.key: entry.value for entry in entries} def get_bot_configs(bot_profile_ids: list[int]) -> dict[int, dict[str, str]]: if not bot_profile_ids: return {} entries = BotConfigData.objects.filter(bot_profile_id__in=bot_profile_ids) entries_by_uid: dict[int, dict[str, str]] = defaultdict(dict) for entry in entries: entries_by_uid[entry.bot_profile_id].update({entry.key: entry.value}) return entries_by_uid def get_bot_config_size(bot_profile: UserProfile, key: str | None = None) -> int: if key is None: return ( BotConfigData.objects.filter(bot_profile=bot_profile) .annotate(key_size=Length("key"), value_size=Length("value")) .aggregate(sum=Sum(F("key_size") + F("value_size")))["sum"] or 0 ) else: try: return len(key) + len(BotConfigData.objects.get(bot_profile=bot_profile, key=key).value) except BotConfigData.DoesNotExist: return 0 def set_bot_config(bot_profile: UserProfile, key: str, value: str) -> None: config_size_limit = settings.BOT_CONFIG_SIZE_LIMIT old_entry_size = get_bot_config_size(bot_profile, key) new_entry_size = len(key) + len(value) old_config_size = get_bot_config_size(bot_profile) new_config_size = old_config_size + (new_entry_size - old_entry_size) if new_config_size > config_size_limit: raise ConfigError( f"Cannot store configuration. Request would require {new_config_size} characters. " f"The current configuration size limit is {config_size_limit} characters." ) obj, created = BotConfigData.objects.get_or_create( bot_profile=bot_profile, key=key, defaults={"value": value} ) if not created: obj.value = value obj.save() def load_bot_config_template(bot: str) -> dict[str, str]: bot_module_name = f"zulip_bots.bots.{bot}" bot_module = importlib.import_module(bot_module_name) assert bot_module.__file__ is not None bot_module_path = os.path.dirname(bot_module.__file__) config_path = os.path.join(bot_module_path, f"{bot}.conf") if os.path.isfile(config_path): config = configparser.ConfigParser() with open(config_path) as conf: config.read_file(conf) return dict(config.items(bot)) else: return {}