From 2af835df30c910afc425caa3998152d77d99fc0c Mon Sep 17 00:00:00 2001 From: Jose Luis Franco Arza Date: Tue, 19 Nov 2024 17:02:21 +0100 Subject: [PATCH] Add support for RBAC in weaviate-cli. This commit adds the option to add, remove, get roles. Assigning roles to users, as well as permissions to roles, revoking them (roles as well as permissions). It also adds the option to specify a user from your config as the user to pass in the client connection. This allows having multiple api-keys assigned to different users without having the need to keep different configs. --- README.md | 21 ++ cli.py | 15 +- requirements-dev.txt | 3 +- setup.cfg | 3 +- .../test_managers/test_config_manager.py | 99 +++++++++ test/unittests/test_utils.py | 193 +++++++++++++++++- weaviate_cli/commands/add.py | 87 ++++++++ weaviate_cli/commands/create.py | 56 ++++- weaviate_cli/commands/delete.py | 36 +++- weaviate_cli/commands/get.py | 116 ++++++++++- weaviate_cli/commands/revoke.py | 86 ++++++++ weaviate_cli/defaults.py | 26 ++- weaviate_cli/managers/collection_manager.py | 1 + weaviate_cli/managers/config_manager.py | 19 +- weaviate_cli/managers/role_manager.py | 132 ++++++++++++ weaviate_cli/managers/user_manager.py | 48 +++++ weaviate_cli/utils.py | 152 ++++++++++++++ 17 files changed, 1068 insertions(+), 25 deletions(-) create mode 100644 weaviate_cli/commands/add.py create mode 100644 weaviate_cli/commands/revoke.py create mode 100644 weaviate_cli/managers/role_manager.py create mode 100644 weaviate_cli/managers/user_manager.py diff --git a/README.md b/README.md index 37b641a..d05678c 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,27 @@ Here you can see an example on how the configuration file should look like if yo } ``` +If you want to allow using different users for different actions in your cluster, you can specify the different users in the configuration file and use the `--user` option to specify which user to use for a specific action. +An example of how the configuration file should look like is the following: + +```json +{ + "host": "your-weaviate-host", + "auth": { + "type": "user", + "user1": "your-api-key-for-user1", + "user2": "your-api-key-for-user2" + } +} +``` +It's important to note that the "type" key must be set to "user" and the users must be specified in the auth section. +When using the `weaviate-cli` command, you can specify the user to use for an action by using the `--user` option. For example: + +```bash +weaviate-cli --user user1 create collection --collection movies --vectorizer transformers +weaviate-cli --user user2 get collection --collection movies +``` + ## Requirements - Python 3.9+ diff --git a/cli.py b/cli.py index 19bec4d..ac6d0c5 100644 --- a/cli.py +++ b/cli.py @@ -1,3 +1,4 @@ +from typing import Optional import click import sys from weaviate_cli.managers.config_manager import ConfigManager @@ -8,6 +9,8 @@ from weaviate_cli.commands.query import query from weaviate_cli.commands.restore import restore from weaviate_cli.commands.cancel import cancel +from weaviate_cli.commands.add import add +from weaviate_cli.commands.revoke import revoke from weaviate_cli import __version__ @@ -26,6 +29,12 @@ def print_version(ctx, param, value): is_flag=False, help="If specified cli uses the config specified with this path.", ) +@click.option( + "--user", + required=False, + type=str, + help="If specified cli uses the user specified in the config file.", +) @click.option( "--version", is_flag=True, @@ -35,10 +44,10 @@ def print_version(ctx, param, value): help="Prints the version of the CLI.", ) @click.pass_context -def main(ctx: click.Context, config_file): +def main(ctx: click.Context, config_file: Optional[str], user: Optional[str]): """Weaviate CLI tool""" try: - ctx.obj = {"config": ConfigManager(config_file)} + ctx.obj = {"config": ConfigManager(config_file, user)} except Exception as e: click.echo(f"Fatal Error: {e}") sys.exit(1) @@ -51,6 +60,8 @@ def main(ctx: click.Context, config_file): main.add_command(restore) main.add_command(query) main.add_command(cancel) +main.add_command(add) +main.add_command(revoke) if __name__ == "__main__": main() diff --git a/requirements-dev.txt b/requirements-dev.txt index 561b4a7..52cac41 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,4 +1,5 @@ -weaviate-client>=4.0.0 +#weaviate-client>=4.0.0 +weaviate-client @ git+https://github.com/weaviate/weaviate-python-client.git@1.28/add-support-for-rbac click==8.1.7 twine pytest diff --git a/setup.cfg b/setup.cfg index 6ae917b..43b22a7 100644 --- a/setup.cfg +++ b/setup.cfg @@ -37,7 +37,8 @@ classifiers = include_package_data = True python_requires = >=3.9 install_requires = - weaviate-client>=4.5.0 + weaviate-client @ git+https://github.com/weaviate/weaviate-python-client.git@1.28/add-support-for-rbac + #weaviate-client>=4.5.0 click==8.1.7 semver>=3.0.2 numpy>=1.24.0 diff --git a/test/unittests/test_managers/test_config_manager.py b/test/unittests/test_managers/test_config_manager.py index 2249153..4952eee 100644 --- a/test/unittests/test_managers/test_config_manager.py +++ b/test/unittests/test_managers/test_config_manager.py @@ -1,5 +1,6 @@ import socket import pytest +import weaviate from unittest.mock import MagicMock, mock_open, patch from weaviate_cli.managers.config_manager import ConfigManager import json @@ -40,3 +41,101 @@ def test_check_host_docker_internal_failed(mock_socket): # Test failed connection mock_socket.side_effect = socket.error() assert config._ConfigManager__check_host_docker_internal() == False + + +def test_init_with_user(): + config_data = { + "host": "localhost", + "http_port": "8080", + "grpc_port": "50051", + "auth": {"type": "user", "admin": "admin-key", "jose": "jose-key"}, + } + + with patch("builtins.open", mock_open(read_data=json.dumps(config_data))): + with patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = True + config = ConfigManager(config_file="test_config.json", user="jose") + assert config.user == "jose" + assert config.config == config_data + + +def test_get_client_with_user_auth(): + config_data = { + "host": "localhost", + "http_port": "8080", + "grpc_port": "50051", + "auth": {"type": "user", "admin": "admin-key", "jose": "jose-key"}, + } + + with patch("builtins.open", mock_open(read_data=json.dumps(config_data))): + with patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = True + with patch("weaviate.connect_to_local") as mock_connect: + config = ConfigManager(config_file="test_config.json", user="jose") + config.get_client() + + mock_connect.assert_called_once() + call_kwargs = mock_connect.call_args.kwargs + assert isinstance( + call_kwargs["auth_credentials"], weaviate.auth.AuthApiKey + ) + assert call_kwargs["auth_credentials"].api_key == "jose-key" + + +def test_get_client_with_invalid_user(): + config_data = { + "host": "localhost", + "http_port": "8080", + "grpc_port": "50051", + "auth": {"type": "user", "admin": "admin-key"}, + } + + with patch("builtins.open", mock_open(read_data=json.dumps(config_data))): + with patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = True + config = ConfigManager(config_file="test_config.json", user="jose") + with pytest.raises(Exception) as exc_info: + config.get_client() + assert str(exc_info.value) == "User 'jose' not found in config file" + + +def test_get_client_missing_user(): + config_data = { + "host": "localhost", + "http_port": "8080", + "grpc_port": "50051", + "auth": {"type": "user", "admin": "admin-key"}, + } + + with patch("builtins.open", mock_open(read_data=json.dumps(config_data))): + with patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = True + config = ConfigManager(config_file="test_config.json") + with pytest.raises(Exception) as exc_info: + config.get_client() + assert ( + str(exc_info.value) == "User must be specified when auth type is 'user'" + ) + + +def test_get_client_with_api_key_auth(): + config_data = { + "host": "localhost", + "http_port": "8080", + "grpc_port": "50051", + "auth": {"type": "api_key", "api_key": "test-key"}, + } + + with patch("builtins.open", mock_open(read_data=json.dumps(config_data))): + with patch("os.path.isfile") as mock_isfile: + mock_isfile.return_value = True + with patch("weaviate.connect_to_local") as mock_connect: + config = ConfigManager(config_file="test_config.json") + config.get_client() + + mock_connect.assert_called_once() + call_kwargs = mock_connect.call_args.kwargs + assert isinstance( + call_kwargs["auth_credentials"], weaviate.auth.AuthApiKey + ) + assert call_kwargs["auth_credentials"].api_key == "test-key" diff --git a/test/unittests/test_utils.py b/test/unittests/test_utils.py index eda2e29..78caabe 100644 --- a/test/unittests/test_utils.py +++ b/test/unittests/test_utils.py @@ -1,9 +1,15 @@ import pytest from unittest.mock import MagicMock -from weaviate_cli.utils import get_client_from_context, get_random_string, pp_objects +from weaviate_cli.utils import ( + get_client_from_context, + get_random_string, + pp_objects, + parse_permission, +) from weaviate.collections import Collection from io import StringIO import sys +from weaviate.rbac.models import RBAC def test_get_client_from_context(mock_click_context, mock_client): @@ -95,3 +101,188 @@ def test_pp_objects_missing_properties(): assert "test-uuid-5678" in output assert "test_name" in output assert "None" in output + + +def test_parse_permission_collections(): + # Test basic schema permissions + assert parse_permission("read_collections:Movies") == [ + RBAC.permissions.collections.read(collection="Movies") + ] + + assert parse_permission("create_collections:*") == [ + RBAC.permissions.collections.create(collection="*") + ] + + # Test without collection specified (should default to "*") + assert parse_permission("update_collections") == [ + RBAC.permissions.collections.update(collection="*") + ] + + assert parse_permission("delete_collections:Movies") == [ + RBAC.permissions.collections.delete(collection="Movies") + ] + + # Test schema with * wildcard + assert parse_permission("create_collections:*data*") == [ + RBAC.permissions.collections.create(collection="*data*") + ] + + assert parse_permission("update_collections:*") == [ + RBAC.permissions.collections.update(collection="*") + ] + + assert parse_permission("manage_collections") == [ + RBAC.permissions.collections.manage(collection="*") + ] + + # Test schema with crud shorthand + assert parse_permission("crud_collections:MyCollection*") == [ + RBAC.permissions.collections.create(collection="MyCollection*"), + RBAC.permissions.collections.read(collection="MyCollection*"), + RBAC.permissions.collections.update(collection="MyCollection*"), + RBAC.permissions.collections.delete(collection="MyCollection*"), + ] + + +def test_parse_permission_roles(): + # Test roles permissions + assert parse_permission("manage_roles") == [RBAC.permissions.roles.manage()] + + assert parse_permission("read_roles") == [RBAC.permissions.roles.read()] + + assert parse_permission("manage_roles:TenantReader") == [ + RBAC.permissions.roles.manage(role="TenantReader") + ] + + assert parse_permission("read_roles:TenantReader") == [ + RBAC.permissions.roles.read(role="TenantReader") + ] + + +def test_parse_permission_standalone(): + # Test standalone permissions + assert parse_permission("manage_users") == [RBAC.permissions.users.manage()] + assert parse_permission("read_cluster") == [RBAC.permissions.cluster.read()] + + +def test_parse_permission_backup(): + # Test manage_backups permission + assert parse_permission("manage_backups") == [ + RBAC.permissions.backups.manage(collection="*") + ] + + assert parse_permission("manage_backups:Movies") == [ + RBAC.permissions.backups.manage(collection="Movies") + ] + + assert parse_permission("manage_backups:User.*") == [ + RBAC.permissions.backups.manage(collection="User.*") + ] + + +def test_parse_permission_nodes(): + assert parse_permission("read_nodes:verbose") == [ + RBAC.permissions.nodes.read(verbosity="verbose") + ] + + assert parse_permission("read_nodes:minimal") == [ + RBAC.permissions.nodes.read(verbosity="minimal") + ] + + assert parse_permission("read_nodes:verbose:Movies") == [ + RBAC.permissions.nodes.read(verbosity="verbose", collection="Movies") + ] + + +def test_parse_permission_crud(): + # Test crud shorthand for schema + crud_perms = parse_permission("crud_collections") + expected_crud = [ + RBAC.permissions.collections.create(collection="*"), + RBAC.permissions.collections.read(collection="*"), + RBAC.permissions.collections.update(collection="*"), + RBAC.permissions.collections.delete(collection="*"), + ] + + assert crud_perms == expected_crud + + # Test crud shorthand for specific collection + crud_tenant_perms = parse_permission("crud_collections:Movies") + expected_tenant_crud = [ + RBAC.permissions.collections.create(collection="Movies"), + RBAC.permissions.collections.read(collection="Movies"), + RBAC.permissions.collections.update(collection="Movies"), + RBAC.permissions.collections.delete(collection="Movies"), + ] + assert crud_tenant_perms == expected_tenant_crud + + +def test_parse_permission_partial_crud(): + # Test cr (create, read) shorthand + cr_perms = parse_permission("cr_collections:Movies") + expected_cr = [ + RBAC.permissions.collections.create(collection="Movies"), + RBAC.permissions.collections.read(collection="Movies"), + ] + assert cr_perms == expected_cr + + # Test ud (update, delete) shorthand + ud_perms = parse_permission("ud_collections:Movies") + expected_ud = [ + RBAC.permissions.collections.update(collection="Movies"), + RBAC.permissions.collections.delete(collection="Movies"), + ] + assert ud_perms == expected_ud + + # Test ru (read, update) shorthand + ru_perms = parse_permission("ru_collections:Movies") + expected_ru = [ + RBAC.permissions.collections.read(collection="Movies"), + RBAC.permissions.collections.update(collection="Movies"), + ] + assert ru_perms == expected_ru + + +def test_parse_permission_data(): + # Test crud shorthand for data + crud_perms = parse_permission("crud_data") + expected_crud = [ + RBAC.permissions.data.create(collection="*"), + RBAC.permissions.data.read(collection="*"), + RBAC.permissions.data.update(collection="*"), + RBAC.permissions.data.delete(collection="*"), + ] + assert crud_perms == expected_crud + + # Test individual permissions + create_perm = parse_permission("create_data:Movies") + assert create_perm == [RBAC.permissions.data.create(collection="Movies")] + + read_perm = parse_permission("read_data:Movies") + assert read_perm == [RBAC.permissions.data.read(collection="Movies")] + + update_perm = parse_permission("update_data:Movies") + assert update_perm == [RBAC.permissions.data.update(collection="Movies")] + + delete_perm = parse_permission("delete_data:Movies") + assert delete_perm == [RBAC.permissions.data.delete(collection="Movies")] + + manage_perm = parse_permission("manage_data:Movies") + assert manage_perm == [RBAC.permissions.data.manage(collection="Movies")] + + +def test_parse_permission_invalid(): + # Test invalid action + with pytest.raises(ValueError, match="Invalid resource type: action"): + parse_permission("invalid_action:Movies") + + # Test invalid crud combination + with pytest.raises(ValueError, match="Invalid crud combination: xyz"): + parse_permission("xyz_collections:Movies") + + # Test invalid format + with pytest.raises(ValueError, match="Invalid permission format"): + parse_permission("read_collections:Movies:extra") + + with pytest.raises(ValueError, match="Invalid permission format"): + parse_permission("read_nodes:verbose:Movies:extra") diff --git a/weaviate_cli/commands/add.py b/weaviate_cli/commands/add.py new file mode 100644 index 0000000..e9ff055 --- /dev/null +++ b/weaviate_cli/commands/add.py @@ -0,0 +1,87 @@ +import click +import sys +from weaviate_cli.managers.user_manager import UserManager +from weaviate_cli.managers.role_manager import RoleManager +from weaviate_cli.utils import get_client_from_context + + +@click.group() +def add() -> None: + """Add resources to other resources in Weaviate.""" + pass + + +@add.command("role") +@click.option( + "--role", + multiple=True, + required=True, + help="The name of the role to add. Can be specified multiple times. Example: --role MoviesAdmin --role TenantAdmin", +) +@click.option( + "--to-user", + required=True, + help="The user to add the role to.", +) +@click.pass_context +def add_role_cli(ctx: click.Context, role: tuple[str], to_user: str) -> None: + """Add a role to a user.""" + client = None + try: + client = get_client_from_context(ctx) + user_manager = UserManager(client) + user_manager.add_role(role=role, to_user=to_user) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() + + +@add.command("permission") +@click.option( + "-p", + "--permission", + multiple=True, + required=True, + help="""Permission in format action:collection. Can be specified multiple times. + + Allowed actions: + - User management: manage_users + - Role management: manage_roles, read_roles + - Cluster statistics read: read_cluster + - Backup management: manage_backups + - Collections permissions: create_collections, read_collections, update_collections, delete_collections, manage_collections + - Data permissions: create_data, read_data, update_data, delete_data + - CRUD shorthands: crud_collections, crud_data + - Nodes read: read_nodes + + Example: --permission crud_collections:* --permission read_data:Movies --permission manage_backups:Movies --permission read_cluster --permission read_nodes:verbose:Movies""", +) +@click.option( + "--to-role", + required=True, + help="The name of the role to add the permission to.", +) +@click.pass_context +def add_permission_cli( + ctx: click.Context, permission: tuple[str], to_role: str +) -> None: + """Add a permission to a role.""" + + client = None + try: + client = get_client_from_context(ctx) + role_manager = RoleManager(client) + role_manager.add_permission(permission=permission, to_role=to_role) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() diff --git a/weaviate_cli/commands/create.py b/weaviate_cli/commands/create.py index bc2ae77..9cea6ab 100644 --- a/weaviate_cli/commands/create.py +++ b/weaviate_cli/commands/create.py @@ -6,11 +6,15 @@ from weaviate_cli.managers.collection_manager import CollectionManager from weaviate_cli.managers.tenant_manager import TenantManager from weaviate_cli.managers.data_manager import DataManager +from weaviate_cli.managers.role_manager import RoleManager from weaviate.exceptions import WeaviateConnectionError -from weaviate_cli.defaults import CreateBackupDefaults -from weaviate_cli.defaults import CreateCollectionDefaults -from weaviate_cli.defaults import CreateTenantsDefaults -from weaviate_cli.defaults import CreateDataDefaults +from weaviate_cli.defaults import ( + CreateBackupDefaults, + CreateCollectionDefaults, + CreateTenantsDefaults, + CreateDataDefaults, + CreateRoleDefaults, +) # Create Group @@ -322,3 +326,47 @@ def create_data_cli( finally: if client: client.close() + + +@create.command("role") +@click.option( + "--name", + default=CreateRoleDefaults.name, + help="The name of the role to create.", +) +@click.option( + "-p", + "--permission", + multiple=True, + required=True, + help="""Permission in format action:collection. Can be specified multiple times. + + Allowed actions: + - User management: manage_users + - Role management: manage_roles, read_roles + - Cluster statistics read: read_cluster + - Backup management: manage_backups + - Collections permissions: create_collections, read_collections, update_collections, delete_collections, manage_collections + - Data permissions: create_data, read_data, update_data, delete_data + - Nodes read: read_nodes + - CRUD shorthands: crud_collections, crud_data + + Example: --permission crud_collections:* --permission read_data:Movies --permission manage_backups:Movies --permission read_cluster --permission read_nodes:verbose:Movies""", +) +@click.pass_context +def create_role_cli(ctx: click.Context, name: str, permission: tuple[str]) -> None: + """Create a role in Weaviate.""" + client = None + try: + client = get_client_from_context(ctx) + role_man = RoleManager(client) + role_man.create_role(name=name, permissions=permission) + click.echo(f"Role '{name}' created successfully in Weaviate.") + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() diff --git a/weaviate_cli/commands/delete.py b/weaviate_cli/commands/delete.py index e2e44b6..4e6dfb9 100644 --- a/weaviate_cli/commands/delete.py +++ b/weaviate_cli/commands/delete.py @@ -4,10 +4,13 @@ from weaviate_cli.utils import get_client_from_context from weaviate_cli.managers.collection_manager import CollectionManager from weaviate_cli.managers.data_manager import DataManager -from weaviate.exceptions import WeaviateConnectionError -from weaviate_cli.defaults import DeleteCollectionDefaults -from weaviate_cli.defaults import DeleteTenantsDefaults -from weaviate_cli.defaults import DeleteDataDefaults +from weaviate_cli.managers.role_manager import RoleManager +from weaviate_cli.defaults import ( + DeleteCollectionDefaults, + DeleteTenantsDefaults, + DeleteDataDefaults, + DeleteRoleDefaults, +) # Delete Group @@ -130,3 +133,28 @@ def delete_data_cli(ctx, collection, limit, consistency_level, uuid): finally: if client: client.close() + + +@delete.command("role") +@click.option( + "--name", + default=DeleteRoleDefaults.name, + help="The name of the role to delete.", +) +@click.pass_context +def delete_role_cli(ctx, name): + """Delete a role in Weaviate.""" + + client = None + try: + client = get_client_from_context(ctx) + role_manager = RoleManager(client) + role_manager.delete_role(name=name) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() diff --git a/weaviate_cli/commands/get.py b/weaviate_cli/commands/get.py index 4bfa478..0816f17 100644 --- a/weaviate_cli/commands/get.py +++ b/weaviate_cli/commands/get.py @@ -1,15 +1,23 @@ import sys +from typing import List, Optional import click +from weaviate_cli.managers.role_manager import RoleManager from weaviate_cli.managers.tenant_manager import TenantManager +from weaviate_cli.managers.user_manager import UserManager from weaviate_cli.utils import get_client_from_context from weaviate_cli.managers.collection_manager import CollectionManager from weaviate_cli.managers.backup_manager import BackupManager from weaviate_cli.managers.shard_manager import ShardManager from weaviate.exceptions import WeaviateConnectionError -from weaviate_cli.defaults import GetBackupDefaults -from weaviate_cli.defaults import GetTenantsDefaults -from weaviate_cli.defaults import GetShardsDefaults -from weaviate_cli.defaults import GetCollectionDefaults +from weaviate.rbac.models import Role +from weaviate_cli.defaults import ( + GetBackupDefaults, + GetTenantsDefaults, + GetShardsDefaults, + GetCollectionDefaults, + GetRoleDefaults, + GetUserDefaults, +) # Get Group @@ -136,3 +144,103 @@ def get_backup_cli(ctx, backend, backup_id, restore): finally: if client: client.close() + + +@get.command("role") +@click.option( + "--name", + default=GetRoleDefaults.name, + help="The name of the role to get.", +) +@click.option( + "--for-user", + default=GetRoleDefaults.for_user, + help="The user to get the roles of.", +) +@click.option( + "--all", + is_flag=True, + help="Get all roles in Weaviate.", +) +@click.pass_context +def get_role_cli(ctx, name, for_user, all): + """Get a specific role or all roles in Weaviate. If --role is provided, get the specific role. If --for-user is provided, get the roles of the specific user.""" + + client = None + try: + if all and name: + raise Exception("Either --all or --name is required. Cannot provide both.") + if all and for_user: + raise Exception( + "Either --all or --for-user is required. Cannot provide both." + ) + if not all and not name and not for_user: + raise Exception("Either --all or --name or --for-user is required.") + if name and for_user: + raise Exception( + "Either --name or --for-user is required. Cannot provide both." + ) + client = get_client_from_context(ctx) + role_man = RoleManager(client) + roles: List[Role] = [] + if all: + roles = role_man.get_all_roles() + else: + if name: + roles.append(role_man.get_role(name=name)) + elif for_user: + roles = role_man.get_roles_from_user(for_user=for_user).values() + else: + raise Exception("Either --name or --for-user is required.") + for role in roles: + role_man.print_role(role) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() + + +@get.command("user") +# @click.option( +# "--user", +# default=GetUserDefaults.user_id, +# help="The ID of the user to get.", +# ) +@click.option( + "--of-role", + default=GetUserDefaults.of_role, + help="The name of the role to get the users of.", +) +# @click.option( +# "--all", +# is_flag=True, +# help="Get all users in Weaviate.", +# ) +@click.pass_context +def get_user_cli(ctx, of_role: Optional[str]): + """Get users of a specific role.""" + + client = None + try: + if not of_role: + raise Exception("Currently only --of-role is supported.") + client = get_client_from_context(ctx) + user_man = UserManager(client) + users = user_man.get_user_from_role(of_role=of_role).values() + print(f"Users of role '{of_role}':") + separator = "-" * 50 + print(f"\n{separator}") + for user in users: + user_man.print_user(user) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() diff --git a/weaviate_cli/commands/revoke.py b/weaviate_cli/commands/revoke.py new file mode 100644 index 0000000..790a3e5 --- /dev/null +++ b/weaviate_cli/commands/revoke.py @@ -0,0 +1,86 @@ +import click +import sys +from weaviate_cli.managers.user_manager import UserManager +from weaviate_cli.managers.role_manager import RoleManager +from weaviate_cli.utils import get_client_from_context + + +@click.group() +def revoke() -> None: + """Revoke resources from other resources in Weaviate.""" + pass + + +@revoke.command("role") +@click.option( + "--role", + multiple=True, + required=True, + help="The name of the role to revoke. Can be specified multiple times. Example: --role MoviesAdmin --role TenantAdmin", +) +@click.option( + "--from-user", + required=True, + help="The user to revoke the role from.", +) +@click.pass_context +def revoke_role_cli(ctx: click.Context, role: tuple[str], from_user: str) -> None: + """Revoke a role from a user.""" + client = None + try: + client = get_client_from_context(ctx) + user_manager = UserManager(client) + user_manager.revoke_role(role=role, from_user=from_user) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() + + +@revoke.command("permission") +@click.option( + "-p", + "--permission", + multiple=True, + required=True, + help="""Permission in format action:collection. Can be specified multiple times. + + Allowed actions: + - User management: manage_users + - Role management: manage_roles, read_roles + - Cluster statistics read: read_cluster + - Backup management: manage_backups + - Collections permissions: create_collections, read_collections, update_collections, delete_collections, manage_collections + - Data permissions: create_data, read_data, update_data, delete_data + - CRUD shorthands: crud_collections, crud_data + - Nodes read: read_nodes + + Example: --permission crud_collections:* --permission read_data:Movies --permission manage_backups:Movies --permission read_cluster --permission read_nodes:verbose:Movies""", +) +@click.option( + "--from-role", + required=True, + help="The role to revoke the permission from.", +) +@click.pass_context +def revoke_permission_cli( + ctx: click.Context, permission: tuple[str], from_role: str +) -> None: + """Revoke a permission from a role.""" + client = None + try: + client = get_client_from_context(ctx) + role_manager = RoleManager(client) + role_manager.revoke_permission(permission=permission, from_role=from_role) + except Exception as e: + click.echo(f"Error: {e}") + if client: + client.close() + sys.exit(1) + finally: + if client: + client.close() diff --git a/weaviate_cli/defaults.py b/weaviate_cli/defaults.py index 8d0d49f..fcabe2e 100644 --- a/weaviate_cli/defaults.py +++ b/weaviate_cli/defaults.py @@ -1,5 +1,5 @@ -from dataclasses import dataclass -from typing import Optional +from dataclasses import dataclass, field +from typing import Optional, List, Dict @dataclass @@ -47,6 +47,12 @@ class CreateDataDefaults: vector_dimensions: int = 1536 +@dataclass +class CreateRoleDefaults: + name: str = "NewRole" + permission: tuple = () + + @dataclass class CancelBackupDefaults: backend: str = "s3" @@ -74,6 +80,11 @@ class DeleteDataDefaults: uuid: Optional[str] = None +@dataclass +class DeleteRoleDefaults: + name: str = "NewRole" + + @dataclass class GetCollectionDefaults: collection: Optional[str] = None @@ -90,6 +101,12 @@ class GetShardsDefaults: collection: Optional[str] = None +@dataclass +class GetRoleDefaults: + name: Optional[str] = None + for_user: Optional[str] = None + + @dataclass class GetBackupDefaults: backend: str = "s3" @@ -97,6 +114,11 @@ class GetBackupDefaults: restore: bool = False +@dataclass +class GetUserDefaults: + of_role: Optional[str] = None + + @dataclass class QueryDataDefaults: collection: str = "Movies" diff --git a/weaviate_cli/managers/collection_manager.py b/weaviate_cli/managers/collection_manager.py index cfe03cb..1a8585b 100644 --- a/weaviate_cli/managers/collection_manager.py +++ b/weaviate_cli/managers/collection_manager.py @@ -272,6 +272,7 @@ def update_collection( if auto_tenant_activation is not None else col_obj.config.get().multi_tenancy_config.auto_tenant_activation ) + col_obj.config.update( description=description, vectorizer_config=( diff --git a/weaviate_cli/managers/config_manager.py b/weaviate_cli/managers/config_manager.py index 2736458..1e6a0e0 100644 --- a/weaviate_cli/managers/config_manager.py +++ b/weaviate_cli/managers/config_manager.py @@ -19,8 +19,11 @@ class ConfigManager: default_port: int = 8080 default_grpc_port: int = 50051 - def __init__(self, config_file: Optional[str] = None) -> None: + def __init__( + self, config_file: Optional[str] = None, user: Optional[str] = None + ) -> None: """Initialize config manager with optional config file path""" + self.user = user if config_file: assert os.path.isfile( config_file @@ -74,14 +77,18 @@ def __get_host(self, port: int = 8080) -> str: def get_client(self) -> weaviate.WeaviateClient: """Get weaviate client from config""" - auth_config: Optional[weaviate.auth.AuthCredentials] = None if "auth" in self.config: - if ( - "type" in self.config["auth"] - and self.config["auth"]["type"] == "api_key" - ): + if self.config["auth"].get("type") == "user": + if not self.user: + raise Exception("User must be specified when auth type is 'user'") + if self.user not in self.config["auth"]: + raise Exception(f"User '{self.user}' not found in config file") + auth_config = weaviate.auth.AuthApiKey( + api_key=self.config["auth"][self.user] + ) + elif self.config["auth"].get("type") == "api_key": auth_config = weaviate.auth.AuthApiKey( api_key=self.config["auth"]["api_key"] ) diff --git a/weaviate_cli/managers/role_manager.py b/weaviate_cli/managers/role_manager.py new file mode 100644 index 0000000..3487a5e --- /dev/null +++ b/weaviate_cli/managers/role_manager.py @@ -0,0 +1,132 @@ +from typing import Optional, List, Dict +import json +from weaviate_cli.utils import parse_permission +from weaviate import WeaviateClient +from weaviate.rbac.models import Role, RBAC +from weaviate_cli.defaults import ( + CreateRoleDefaults, + DeleteRoleDefaults, + GetRoleDefaults, +) + + +class RoleManager: + def __init__(self, client: WeaviateClient): + self.client = client + + def create_role( + self, + name: str = CreateRoleDefaults.name, + permissions: tuple = CreateRoleDefaults.permission, + ) -> None: + try: + rbac_permissions = [] + for perm in permissions: + parsed = parse_permission(perm) + if isinstance(parsed, list): + rbac_permissions.extend(parsed) + else: + rbac_permissions.append(parsed) + + self.client.roles.create(name=name, permissions=rbac_permissions) + + except Exception as e: + raise Exception(f"Error creating role '{name}': {e}") + + def get_role(self, name: str) -> Role: + try: + return self.client.roles.by_name(name) + except Exception as e: + raise Exception(f"Error getting role '{name}': {e}") + + def get_roles_from_user(self, for_user: str) -> Dict[str, Role]: + try: + return self.client.roles.by_user(user=for_user) + except Exception as e: + raise Exception(f"Error getting roles from user '{for_user}': {e}") + + def delete_role(self, name: str = DeleteRoleDefaults.name) -> None: + try: + if not self.client.roles.exists(name): + raise Exception(f"Role '{name}' does not exist.") + self.client.roles.delete(name) + assert not self.client.roles.exists(name) + except Exception as e: + raise Exception(f"Error deleting role '{name}': {e}") + + def get_all_roles(self) -> List[Role]: + try: + return self.client.roles.list_all() + except Exception as e: + raise Exception(f"Error getting all roles: {e}") + + def add_permission(self, permission: str, to_role: str) -> None: + try: + rbac_permissions = [] + for perm in permission: + parsed = parse_permission(perm) + rbac_permissions = rbac_permissions + parsed + self.client.roles.add_permissions( + permissions=rbac_permissions, role=to_role + ) + except Exception as e: + raise Exception( + f"Error adding permission '{permission}' to role '{to_role}': {e}" + ) + + def revoke_permission(self, permission: str, from_role: str) -> None: + try: + rbac_permissions = [] + for perm in permission: + parsed = parse_permission(perm) + rbac_permissions = rbac_permissions + parsed + self.client.roles.remove_permissions( + permissions=rbac_permissions, role=from_role + ) + except Exception as e: + raise Exception( + f"Error revoking permission '{permission}' from role '{from_role}': {e}" + ) + + def print_role(self, role: Optional[Role] = None) -> None: + """Print a role and its permissions in a human readable format.""" + if not role: + raise ValueError("Role is required. Please provide a valid role.") + + separator = "-" * 50 + print(f"\n{separator}") + print(f"Role: {role.name}") + + if role.cluster_actions: + print("\nCluster Actions:") + for action in role.cluster_actions: + print(f" - {action.value}") + + if role.nodes_permissions: + print("\nNodes Permissions:") + for perm in role.nodes_permissions: + print( + f" - Verbosity: {perm.verbosity}, Collection: {perm.collection if perm.collection else '*'}, Action: {perm.action.value}" + ) + + if role.backups_permissions: + print("\nBackups Permissions:") + for perm in role.backups_permissions: + print(f" - Collection: {perm.collection}, Action: {perm.action.value}") + + if role.roles_permissions: + print("\nRoles Permissions:") + for perm in role.roles_permissions: + print(f" - Role: {perm.role}, Action: {perm.action.value}") + + if role.config_permissions: + print("\nCollections (schema) Permissions:") + for perm in role.config_permissions: + print( + f" - Collection: {perm.collection}, Tenant: {perm.tenant}, Action: {perm.action.value}" + ) + + if role.data_permissions: + print("\nData Permissions:") + for perm in role.data_permissions: + print(f" - Collection: {perm.collection}, Action: {perm.action.value}") diff --git a/weaviate_cli/managers/user_manager.py b/weaviate_cli/managers/user_manager.py new file mode 100644 index 0000000..7771be5 --- /dev/null +++ b/weaviate_cli/managers/user_manager.py @@ -0,0 +1,48 @@ +from typing import Optional, Dict +from weaviate import WeaviateClient +from weaviate.rbac.models import User +from weaviate_cli.defaults import ( + GetUserDefaults, +) + + +class UserManager: + def __init__(self, client: WeaviateClient): + self.client = client + + def get_user_from_role( + self, of_role: str = GetUserDefaults.of_role + ) -> Dict[str, User]: + """Get all roles assigned to a user.""" + try: + return self.client.roles.users(of_role) + except Exception as e: + raise Exception(f"Error getting users for role '{of_role}': {e}") + + def add_role( + self, + role: tuple[str], + to_user: str, + ) -> None: + """Assign a role to a user.""" + try: + self.client.roles.assign(roles=list(role), user=to_user) + except Exception as e: + raise Exception(f"Error assigning role '{role}' to user '{to_user}': {e}") + + def revoke_role( + self, + role: tuple[str], + from_user: str, + ) -> None: + """Revoke a role from a user.""" + try: + self.client.roles.revoke(roles=list(role), user=from_user) + except Exception as e: + raise Exception( + f"Error revoking role '{role}' from user '{from_user}': {e}" + ) + + def print_user(self, user: User) -> None: + """Print user roles in a human readable format.""" + print(f"User: {user.name}") diff --git a/weaviate_cli/utils.py b/weaviate_cli/utils.py index e6aaa93..1bfe98b 100644 --- a/weaviate_cli/utils.py +++ b/weaviate_cli/utils.py @@ -5,6 +5,8 @@ import string import random import weaviate +from weaviate.rbac.models import RBAC +from typing import Union, List def get_client_from_context(ctx) -> weaviate.Client: @@ -59,3 +61,153 @@ def pp_objects(response, main_properties): footer = f"{'':<37}" * (len(main_properties) + 1) + f"{'':<11}{'':<11}{'':<11}" print(footer) print(f"Total: {len(objects)} objects") + + +def parse_permission(perm: str) -> Union[RBAC.permissions, List[RBAC.permissions]]: + """ + Convert a permission string to RBAC permission object(s). + Format: action:collection[:tenant] + + Supports: + - Basic permissions: read_schema:Movies + - CRUD shorthand: crud_schema:Movies + - Partial CRUD: cr_schema:Movies (create+read) + - User management: manage_users + - Role permissions: manage_roles, read_roles + - Cluster permissions: read_cluster + - Backup permissions: manage_backups + - Collections permissions: create_collections, read_collections, update_collections, delete_collections, manage_collections + - Data permissions: create_data, read_data, update_data, delete_data + - Nodes permissions: read_nodes + Args: + perm (str): Permission string + + Returns: + Union[RBAC.permissions, List[RBAC.permissions]]: Single permission or list for crud/partial crud + """ + valid_resources = [ + "collections", + "data", + "roles", + "users", + "cluster", + "backups", + "nodes", + ] + parts = perm.split(":") + if len(parts) > 3 or (parts[0] != "read_nodes" and len(parts) > 2): + raise ValueError( + f"Invalid permission format: {perm}. Expected format: action:collection/role/verbosity. Example: manage_roles:custom, crud_collections:Movies, read_nodes:verbose:Movies" + ) + action = parts[0] + role = parts[1] if len(parts) > 1 and "roles" in action else "*" + verbosity = parts[1] if len(parts) > 1 and action == "read_nodes" else None + if action == "read_nodes": + # For read_nodes the first part belongs to the verbosity + # and the second (optional) belongs to the collection + # Example: read_nodes:verbose:Movies + collection = parts[2] if len(parts) > 2 else "*" + else: + collection = ( + parts[1] + if len(parts) > 1 + and action + not in ["manage_roles", "manage_users", "read_roles", "read_cluster"] + else "*" + ) + + # Handle standalone permissions first + if action in ["manage_users", "read_cluster", "manage_backups", "read_nodes"]: + return [ + _create_permission( + action=action, collection=collection, verbosity=verbosity + ) + ] + + out_permissions = [] + # Handle crud and partial crud cases + if "_" in action: + parts = action.split("_", 2) + prefix = parts[0] + resource = parts[1] if len(parts) > 1 else None + if resource not in valid_resources: + # Find closest matching resource type + closest = min( + valid_resources, + key=lambda x: ( + sum(c1 != c2 for c1, c2 in zip(x, resource)) + if len(x) == len(resource) + else abs(len(x) - len(resource)) + ), + ) + suggestion = f"\nDid you mean '{closest}'?" if closest else "" + raise ValueError(f"Invalid resource type: {resource}. {suggestion}") + + action_map = {"c": "create", "r": "read", "u": "update", "d": "delete"} + + if prefix in ["create", "read", "update", "delete", "manage"]: + actions = [prefix] + else: + # Handle partial crud (curd, cr, ru, ud, etc) + if not all(c in action_map for c in prefix): + raise ValueError(f"Invalid crud combination: {prefix}") + actions = [action_map[c] for c in prefix] + + for act in actions: + out_permissions.append( + _create_permission( + action=f"{act}_{resource}", + role=role, + collection=collection, + ) + ) + + return out_permissions + + raise ValueError(f"Invalid permission action: {action}") + + +def _create_permission( + action: str, role: str = "*", collection: str = "*", verbosity: str = "minimal" +) -> RBAC.permissions: + """Helper function to create individual RBAC permission objects.""" + + # Handle standalone permissions + if action == "manage_users": + return RBAC.permissions.users.manage() + elif action == "read_cluster": + return RBAC.permissions.cluster.read() + elif action == "manage_backups": + return RBAC.permissions.backups.manage(collection=collection) + elif action == "read_nodes": + return RBAC.permissions.nodes.read(verbosity=verbosity, collection=collection) + # Handle roles permissions + elif action in ["manage_roles", "read_roles"]: + action_prefix = action.split("_")[0] # will be either "manage" or "read" + return getattr(RBAC.permissions.roles, action_prefix)(role=role) + + # Handle schema permissions + elif action in [ + "create_collections", + "read_collections", + "update_collections", + "delete_collections", + "manage_collections", + ]: + action_prefix = action.split("_")[0] + return getattr(RBAC.permissions.collections, action_prefix)( + collection=collection + ) + + # Handle data permissions + elif action in [ + "create_data", + "read_data", + "update_data", + "delete_data", + "manage_data", + ]: + action_prefix = action.split("_")[0] + return getattr(RBAC.permissions.data, action_prefix)(collection=collection) + + raise ValueError(f"Invalid permission action: {action}.")