mirror of https://github.com/zulip/zulip.git
84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
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)
|
|
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 {}
|