diff --git a/neon_diana_utils/cli.py b/neon_diana_utils/cli.py index c4b6c47e..1bf31de1 100644 --- a/neon_diana_utils/cli.py +++ b/neon_diana_utils/cli.py @@ -41,6 +41,15 @@ def neon_diana_cli(version: bool = False): click.echo(f"Diana version {__version__}") +# Generic Utilities +@neon_diana_cli.command(help="Update RabbitMQ Configuration") +@click.argument("rabbitmq_json_path", default=None, required=False) +def update_rabbitmq_config(rabbitmq_json_path): + from neon_diana_utils.configuration import update_rmq_config + updated_file = update_rmq_config(rabbitmq_json_path) + click.echo(f"Updated configuration at: {updated_file}") + + # Core @neon_diana_cli.command(help="Configure Neon Core") @click.option("--username", "-u", help="RabbitMQ username for Neon AI") diff --git a/neon_diana_utils/configuration.py b/neon_diana_utils/configuration.py index 67c35b31..8d33bcf4 100644 --- a/neon_diana_utils/configuration.py +++ b/neon_diana_utils/configuration.py @@ -195,6 +195,41 @@ def make_keys_config(write_config: bool, return config +def update_rmq_config(config_file: str = None) -> str: + """ + Update an existing RabbitMQ configuration with new definitions from DIANA. + This can be used to handle added service users without changing existing + ones and to reset/update permissions and vhosts. + @param config_file: Path to file to be updated + @returns: Path to updated file + """ + if not config_file: + config_file = join(xdg_config_home(), "diana", "diana-backend", + "rabbitmq.json") + else: + config_file = expanduser(config_file) + + if not isfile(config_file): + raise FileNotFoundError(config_file) + + with open(config_file) as f: + real_config = json.load(f) + new_config = generate_rmq_config("", "") + existing_users = (user['name'] for user in real_config['users']) + for user in new_config['users']: + if user['name'] in existing_users: + continue + LOG.info(f"Adding user: {user['name']}") + real_config['users'].append(user) + real_config['vhosts'] = new_config['vhosts'] + real_config['permissions'] = new_config['permissions'] + + shutil.move(config_file, f"{config_file}.old") + with open(config_file, 'w+') as f: + json.dump(real_config, f, indent=2) + return config_file + + def generate_rmq_config(admin_username: str, admin_password: str, output_file: str = None) -> dict: """ @@ -215,10 +250,12 @@ def generate_rmq_config(admin_username: str, admin_password: str, continue user['password'] = secrets.token_urlsafe(32) - base_config['users'].append({'name': admin_username, - 'password': admin_password, - 'tags': ['administrator']}) - + if admin_username and admin_password: + base_config['users'].append({'name': admin_username, + 'password': admin_password, + 'tags': ['administrator']}) + else: + LOG.debug("Not adding unconfigured admin user") if output_file and validate_output_path(output_file): with open(output_file, 'w+') as f: json.dump(base_config, f, indent=2) @@ -261,23 +298,24 @@ def update_env_file(env_file: str): def _get_neon_mq_user_config(mq_user: Optional[str], mq_pass: Optional[str], - backend_config: str) -> dict: + rmq_config: str) -> dict: """ Get MQ config for neon core. @param mq_user: RabbitMQ Neon username @param mq_pass: RabbitMQ Neon password - @param backend_config: Path to Diana Backend configuration file to import + @param rmq_config: Path to RabbitMQ configuration file to import @returns dict user config to connect Neon Core to an MQ instance """ # Check for passed or previously configured MQ user - if not all((mq_user, mq_pass)) and isfile(backend_config): - if click.confirm(f"Import Neon MQ user from {backend_config}?"): - with open(backend_config) as f: - config = yaml.safe_load(f) - user_config = config.get('MQ', {}).get('users', - {}).get('chat_api_proxy', {}) - mq_user = user_config.get('user') - mq_pass = user_config.get('password') + if not all((mq_user, mq_pass)) and isfile(rmq_config): + if click.confirm(f"Import Neon MQ user from {rmq_config}?"): + with open(rmq_config) as f: + config = json.load(f) + for user in config['users']: + if "core" in user['tags']: + mq_user = user['name'] + mq_pass = user['password'] + break # Interactively configure MQ authentication user_config = {"user": mq_user, "password": mq_pass} @@ -399,7 +437,7 @@ def configure_neon_core(mq_user: str = None, # Validate output paths output_path = expanduser(output_path or join(xdg_config_home(), "diana")) - backend_config = join(output_path, "diana-backend", "diana.yaml") + rmq_config = join(output_path, "diana-backend", "rabbitmq.json") # Output to `core` subdirectory if not validate_output_path(join(output_path, "neon-core")): click.echo(f"Path exists: {output_path}") @@ -426,7 +464,7 @@ def configure_neon_core(mq_user: str = None, try: # Get MQ User Configuration - user_config = _get_neon_mq_user_config(mq_user, mq_pass, backend_config) + user_config = _get_neon_mq_user_config(mq_user, mq_pass, rmq_config) if not all((user_config['user'], user_config['password'])): # TODO: Prompt to configure MQ server/port? mq_config = dict() diff --git a/tests/test_diana_utils.py b/tests/test_diana_utils.py index 737672b3..d1bda8f9 100644 --- a/tests/test_diana_utils.py +++ b/tests/test_diana_utils.py @@ -76,6 +76,24 @@ def test_make_keys_config(self, confirm): os.remove(test_output_file) + def test_update_rmq_config(self): + from neon_diana_utils.configuration import update_rmq_config + test_file = join(dirname(__file__), "test_rabbitmq.json") + update_rmq_config(test_file) + self.assertTrue(isfile(test_file)) + self.assertTrue(isfile(f"{test_file}.old")) + with open(test_file) as f: + new_config = json.load(f) + with open(f"{test_file}.old") as f: + old_config = json.load(f) + self.assertIsInstance(new_config, dict) + self.assertIsInstance(old_config, dict) + for user in old_config['users']: + self.assertIn(user, new_config['users']) + + os.remove(test_file) + shutil.move(f"{test_file}.old", test_file) + def test_generate_rmq_config(self): from neon_diana_utils.configuration import generate_rmq_config test_output_file = join(dirname(__file__), "test_rmq.json") diff --git a/tests/test_rabbitmq.json b/tests/test_rabbitmq.json new file mode 100644 index 00000000..c9ba639a --- /dev/null +++ b/tests/test_rabbitmq.json @@ -0,0 +1,210 @@ +{ + "users": [ + { + "name": "neon_api_utils", + "password": "Klatchat2021", + "tags": [ + "backend", + "user" + ] + }, + { + "name": "neon_metrics", + "password": "old_metrics", + "tags": [ + "backend", + "service" + ] + }, + { + "name": "neon_coupons", + "password": "old_coupons", + "tags": [ + "backend", + "service" + ] + }, + { + "name": "neon_email", + "password": "old_email", + "tags": [ + "backend", + "service" + ] + }, + { + "name": "neon_script_parser", + "password": "old_script_parser", + "tags": [ + "backend", + "service" + ] + }, + { + "name": "neon_api", + "password": "old_neon_api", + "tags": [ + "backend", + "service" + ] + }, + { + "name": "neon_libretranslate", + "password": "old_libretranslate", + "tags": [ + "backend", + "service" + ] + } + ], + "vhosts": [ + { + "name": "/neon_emails" + }, + { + "name": "/neon_api" + }, + { + "name": "/neon_script_parser" + }, + { + "name": "/neon_metrics" + }, + { + "name": "/neon_coupons" + }, + { + "name": "/neon_testing" + }, + { + "name": "/translation" + }, + { + "name": "/llm" + }, + { + "name": "/neon_chat_api" + } + ], + "permissions": [ + { + "user": "neon_core", + "vhost": "/neon_chat_api", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_chat_api", + "configure": ".*", + "write": ".*", + "read": ".*(?!_request).*" + }, + { + "user": "neon_llm_chatgpt", + "vhost": "/llm", + "configure": ".*", + "write": ".*", + "read": "chat_gpt_input" + }, + { + "user": "neon_llm_fastchat", + "vhost": "/llm", + "configure": ".*", + "write": ".*", + "read": "fastchat_input" + }, + { + "user": "neon_api_utils", + "vhost": "/llm", + "configure": ".*", + "write": ".*", + "read": ".*(?!_input).*" + }, + { + "user": "neon_libretranslate", + "vhost": "/translation", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_coupons", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_email", + "vhost": "/neon_emails", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_emails", + "configure": ".*", + "write": ".*", + "read": "^(?!neon_emails_input).*" + }, + { + "user": "neon_api", + "vhost": "/neon_api", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_metrics", + "configure": ".*", + "write": ".*", + "read": "" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_api", + "configure": "./*", + "write": "./*", + "read": "./*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_script_parser", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_script_parser", + "vhost": "/neon_script_parser", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_metrics", + "vhost": "/neon_metrics", + "configure": ".*", + "write": ".*", + "read": ".*" + }, + { + "user": "neon_api_utils", + "vhost": "/neon_testing", + "configure": "./*", + "write": "./*", + "read": "./*" + }, + { + "user": "neon_coupons", + "vhost": "/neon_coupons", + "configure": ".*", + "write": ".*", + "read": ".*" + } + ] +} \ No newline at end of file