diff --git a/zulip_botserver/tests/test_server.py b/zulip_botserver/tests/test_server.py index d38cd93..653ef4b 100644 --- a/zulip_botserver/tests/test_server.py +++ b/zulip_botserver/tests/test_server.py @@ -4,6 +4,7 @@ from typing import Any, Dict import unittest from .server_test_lib import BotServerTestCase import json +from collections import OrderedDict from importlib import import_module from types import ModuleType @@ -132,6 +133,34 @@ class BotServerTests(BotServerTestCase): assert opts.hostname == '127.0.0.1' assert opts.port == 5002 + def test_read_config_from_env_vars(self) -> None: + # We use an OrderedDict so that the order of the entries in + # the stringified environment variable is standard even on + # Python 3.7 and earlier. + bots_config = OrderedDict() + bots_config['hello_world'] = { + 'email': 'helloworld-bot@zulip.com', + 'key': 'value', + 'site': 'http://localhost', + 'token': 'abcd1234', + } + bots_config['giphy'] = { + 'email': 'giphy-bot@zulip.com', + 'key': 'value2', + 'site': 'http://localhost', + 'token': 'abcd1234', + } + os.environ['ZULIP_BOTSERVER_CONFIG'] = json.dumps(bots_config) + + # No bot specified; should read all bot configs + assert server.read_config_from_env_vars() == bots_config + + # Specified bot exists; should read only that section. + assert server.read_config_from_env_vars("giphy") == {'giphy': bots_config['giphy']} + + # Specified bot doesn't exist; should read the first section of the config. + assert server.read_config_from_env_vars("redefined_bot") == {'redefined_bot': bots_config['hello_world']} + def test_read_config_file(self) -> None: with self.assertRaises(IOError): server.read_config_file("nonexistentfile.conf") diff --git a/zulip_botserver/zulip_botserver/input_parameters.py b/zulip_botserver/zulip_botserver/input_parameters.py index 2d6a2fe..1b7d585 100644 --- a/zulip_botserver/zulip_botserver/input_parameters.py +++ b/zulip_botserver/zulip_botserver/input_parameters.py @@ -7,13 +7,19 @@ def parse_args() -> argparse.Namespace: ''' parser = argparse.ArgumentParser(usage=usage) - parser.add_argument( + mutually_exclusive_args = parser.add_mutually_exclusive_group(required=True) + # config-file or use-env-vars made mutually exclusive to prevent conflicts + mutually_exclusive_args.add_argument( '--config-file', '-c', action='store', - required=True, help='Config file for the Botserver. Use your `botserverrc` for multiple bots or' '`zuliprc` for a single bot.' ) + mutually_exclusive_args.add_argument( + '--use-env-vars', '-e', + action='store_true', + help='Load configuration from JSON in ZULIP_BOTSERVER_CONFIG environment variable.' + ) parser.add_argument( '--bot-config-file', action='store', diff --git a/zulip_botserver/zulip_botserver/server.py b/zulip_botserver/zulip_botserver/server.py index 6ad6762..ed5a048 100644 --- a/zulip_botserver/zulip_botserver/server.py +++ b/zulip_botserver/zulip_botserver/server.py @@ -7,6 +7,7 @@ import os import sys import importlib.util +from collections import OrderedDict from configparser import MissingSectionHeaderError, NoOptionError from flask import Flask, request from importlib import import_module @@ -28,6 +29,32 @@ def read_config_section(parser: configparser.ConfigParser, section: str) -> Dict } return section_info +def read_config_from_env_vars(bot_name: Optional[str] = None) -> Dict[str, Dict[str, str]]: + bots_config = {} # type: Dict[str, Dict[str, str]] + json_config = os.environ.get('ZULIP_BOTSERVER_CONFIG') + + if json_config is None: + raise OSError("Could not read environment variable 'ZULIP_BOTSERVER_CONFIG': Variable not set.") + + # Load JSON-formatted environment variable; use OrderedDict to + # preserve ordering on Python 3.6 and below. + env_config = json.loads(json_config, object_pairs_hook=OrderedDict) + if bot_name is not None: + if bot_name in env_config: + bots_config[bot_name] = env_config[bot_name] + else: + # If the bot name provided via the command line does not + # exist in the configuration provided via the environment + # variable, use the first bot in the environment variable, + # with name updated to match, along with a warning. + first_bot_name = list(env_config.keys())[0] + bots_config[bot_name] = env_config[first_bot_name] + logging.warning( + "First bot name in the config list was changed from '{}' to '{}'".format(first_bot_name, bot_name) + ) + else: + bots_config = dict(env_config) + return bots_config def read_config_file(config_file_path: str, bot_name: Optional[str] = None) -> Dict[str, Dict[str, str]]: parser = parse_config_file(config_file_path) @@ -178,16 +205,20 @@ def handle_bot() -> str: def main() -> None: options = parse_args() global bots_config - try: - bots_config = read_config_file(options.config_file, options.bot_name) - except MissingSectionHeaderError: - sys.exit("Error: Your Botserver config file `{0}` contains an empty section header!\n" - "You need to write the names of the bots you want to run in the " - "section headers of `{0}`.".format(options.config_file)) - except NoOptionError as e: - sys.exit("Error: Your Botserver config file `{0}` has a missing option `{1}` in section `{2}`!\n" - "You need to add option `{1}` with appropriate value in section `{2}` of `{0}`" - .format(options.config_file, e.option, e.section)) + + if options.use_env_vars: + bots_config = read_config_from_env_vars(options.bot_name) + elif options.config_file: + try: + bots_config = read_config_file(options.config_file, options.bot_name) + except MissingSectionHeaderError: + sys.exit("Error: Your Botserver config file `{0}` contains an empty section header!\n" + "You need to write the names of the bots you want to run in the " + "section headers of `{0}`.".format(options.config_file)) + except NoOptionError as e: + sys.exit("Error: Your Botserver config file `{0}` has a missing option `{1}` in section `{2}`!\n" + "You need to add option `{1}` with appropriate value in section `{2}` of `{0}`" + .format(options.config_file, e.option, e.section)) available_bots = list(bots_config.keys()) bots_lib_modules = load_lib_modules(available_bots)