import configparser import importlib import os from collections import defaultdict from typing import Dict, List, Optional from django.conf import settings from django.db.models import Sum from django.db.models.functions import Length from django.db.models.query import F 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: Optional[str] = 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( "Cannot store configuration. Request would require {} characters. " "The current configuration size limit is {} characters.".format( new_config_size, config_size_limit ) ) 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 {}