Added CLI Command databricks labs ucx principal-prefix-access (#949)

## Changes
This PR merges the three cli command save-azure-storage-accounts,
save-aws-iam-profiles and save-uc-compatible-roles into one common cmd
save-storage-and-principal
 - updating labs.yml and cli file to replace existing cmd with new cmd
 - updated cli_test cases to reflect the change in code

### Closes #909 


### Functionality 

- [ ] added relevant user documentation
- [X] added new CLI command
- [X] modified existing command: `databricks labs ucx ...`
- [ ] added a new workflow
- [ ] modified existing workflow: `...`
- [ ] added a new table
- [ ] modified existing table: `...`

### Tests

- [X] manually tested
- [X] added unit tests
- [ ] added integration tests
- [ ] verified on staging environment (screenshot attached)

Author: HariGS-DB

README.md

@@ -150,18 +150,19 @@ After UCX assessment workflow is executed, the assessment dashboard will be popu
 ### Scanning for legacy credentials and mapping access
 #### AWS
 Use to identify all instance profiles in the workspace, and map their access to S3 buckets. 
+Also captures the IAM roles which has UC arn listed, and map their access to S3 buckets 
 This requires `awscli` to be installed and configured.
 
 ```commandline
-databricks labs ucx save-aws-iam-profiles
+databricks labs ucx principal-prefix-access --aws-profile test-profile
 ```
 
 #### Azure
 Use to identify all storage account used by tables, identify the relevant Azure service principals and their permissions on each storage account.
 This requires `azure-cli` to be installed and configured. 
 
 ```commandline
-databricks labs ucx save-azure-storage-accounts
+databricks labs ucx principal-prefix-access --subscription-id test-subscription-id
 ```
 
 ### Producing table mapping

labs.yml

@@ -97,31 +97,20 @@ commands:
       - name: to-schema
         description: target schema to migrate tables to
 
-  - name: save-azure-storage-accounts
-    description: Identifies all storage account used by tables, identify spn and its permission on each storage accounts
+  - name: principal-prefix-access
+    description: For azure cloud, identifies all storage account used by tables in the workspace, identify spn and its 
+      permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and 
+      its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets. 
+      The output is stored in the workspace install folder.
     flags:
       - name: subscription-id
         description: Subscription to scan storage account in
+      - name: aws-profile
+        description: AWS Profile to use for authentication
 
   - name: validate-groups-membership
     description: Validate the groups to see if the groups at account level and workspace level have different membership
     table_template: |-
       Workspace Group Name\tMembers Count\tAccount Group Name\tMembers Count
       {{range .}}{{.wf_group_name}}\t{{.wf_group_members_count}}\t{{.acc_group_name}}\t{{.acc_group_members_count}}
       {{end}}
-
-  - name: save-aws-iam-profiles
-    description: | 
-      Identifies all Instance Profiles and map their access to S3 buckets.
-      Requires a working setup of AWS CLI.
-    flags:
-      - name: aws-profile
-        description: AWS Profile to use for authentication
-
-  - name: save-uc-compatible-roles
-    description: | 
-      Scan all the AWS roles that are set for UC access and produce a mapping to the S3 resources.
-      Requires a working setup of AWS CLI.
-    flags:
-      - name: aws-profile
-        description: AWS Profile to use for authentication

src/databricks/labs/ucx/cli.py

@@ -228,76 +228,46 @@ def alias(
 
 
 @ucx.command
-def save_azure_storage_accounts(w: WorkspaceClient, subscription_id: str):
-    """identifies all azure storage account used by external tables
-    identifies all spn which has storage blob reader, blob contributor, blob owner access
-    saves the data in ucx database."""
-    if not w.config.is_azure:
-        logger.error("Workspace is not on azure, please run this command on azure databricks workspaces.")
-        return
-    if w.config.auth_type != "azure_cli":
-        logger.error("In order to obtain AAD token, Please run azure cli to authenticate.")
-        return
-    if subscription_id == "":
-        logger.error("Please enter subscription id to scan storage account in.")
-        return
-    azure_resource_permissions = AzureResourcePermissions.for_cli(w)
-    logger.info("Generating azure storage accounts and service principal permission info")
-    azure_resource_permissions.save_spn_permissions()
-
-
-@ucx.command
-def save_aws_iam_profiles(w: WorkspaceClient, aws_profile: str | None = None):
-    """identifies all Instance Profiles and map their access to S3 buckets.
-    Requires a working setup of AWS CLI.
-    https://aws.amazon.com/cli/
-    The command saves a CSV to the UCX installation folder with the mapping.
-
-    The user has to be authenticated with AWS and the have the permissions to browse the resources and iam services.
-    More information can be found here:
-    https://docs.aws.amazon.com/IAM/latest/UserGuide/access_permissions-required.html
-    """
-    if not shutil.which("aws"):
-        logger.error("Couldn't find AWS CLI in path.Please obtain and install the CLI from https://aws.amazon.com/cli/")
-        return None
-    if not aws_profile:
-        aws_profile = os.getenv("AWS_DEFAULT_PROFILE")
-    if not aws_profile:
-        logger.error(
-            "AWS Profile is not specified. Use the environment variable [AWS_DEFAULT_PROFILE] "
-            "or use the '--aws-profile=[profile-name]' parameter."
-        )
-        return None
-    aws_permissions = AWSResourcePermissions.for_cli(w, aws_profile)
-    aws_permissions.save_instance_profile_permissions()
-    return None
-
-
-@ucx.command
-def save_uc_compatible_roles(w: WorkspaceClient, *, aws_profile: str | None = None):
-    """extracts all the iam roles with trust relationships to the UC master role.
-     Map these roles to the S3 buckets they have access to.
-    Requires a working setup of AWS CLI.
-    https://aws.amazon.com/cli/
-    The command saves a CSV to the UCX installation folder with the mapping.
-
-    The user has to be authenticated with AWS and the have the permissions to browse the resources and iam services.
-    More information can be found here:
-    https://docs.aws.amazon.com/IAM/latest/UserGuide/access_permissions-required.html
-    """
-    if not shutil.which("aws"):
-        logger.error("Couldn't find AWS CLI in path.Please obtain and install the CLI from https://aws.amazon.com/cli/")
-        return None
-    if not aws_profile:
-        aws_profile = os.getenv("AWS_DEFAULT_PROFILE")
-    if not aws_profile:
-        logger.error(
-            "AWS Profile is not specified. Use the environment variable [AWS_DEFAULT_PROFILE] "
-            "or use the '--aws-profile=[profile-name]' parameter."
-        )
-        return None
-    aws_permissions = AWSResourcePermissions.for_cli(w, aws_profile)
-    aws_permissions.save_uc_compatible_roles()
+def principal_prefix_access(w: WorkspaceClient, subscription_id: str | None = None, aws_profile: str | None = None):
+    """For azure cloud, identifies all storage account used by tables in the workspace, identify spn and its
+    permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and
+    its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets.
+    The output is stored in the workspace install folder.
+    Pass suscription_id for azure and aws_profile for aws."""
+    print("testing")
+    if w.config.is_azure:
+        if w.config.auth_type != "azure-cli":
+            logger.error("In order to obtain AAD token, Please run azure cli to authenticate.")
+            return None
+        if subscription_id == "":
+            logger.error("Please enter subscription id to scan storage account in.")
+            return None
+        azure_resource_permissions = AzureResourcePermissions.for_cli(w)
+        logger.info("Generating azure storage accounts and service principal permission info")
+        path = azure_resource_permissions.save_spn_permissions()
+        if path:
+            logger.info(f"storage and spn info saved under {path}")
+    elif w.config.is_aws:
+        if not shutil.which("aws"):
+            logger.error("Couldn't find AWS CLI in path. Please install the CLI from https://aws.amazon.com/cli/")
+            return None
+        if not aws_profile:
+            aws_profile = os.getenv("AWS_DEFAULT_PROFILE")
+        if not aws_profile:
+            logger.error(
+                "AWS Profile is not specified. Use the environment variable [AWS_DEFAULT_PROFILE] "
+                "or use the '--aws-profile=[profile-name]' parameter."
+            )
+            return None
+        logger.info("Generating instance profile and bucket permission info")
+        aws_permissions = AWSResourcePermissions.for_cli(w, aws_profile)
+        instance_role_path = aws_permissions.save_instance_profile_permissions()
+        logger.info(f"Instance profile and bucket info saved {instance_role_path}")
+        logger.info("Generating UC roles and bucket permission info")
+        uc_role_path = aws_permissions.save_uc_compatible_roles()
+        logger.info(f"UC roles and bucket info saved {uc_role_path}")
+    else:
+        logger.error("This cmd is only supported for azure and aws workspaces")
     return None
 
 

tests/integration/azure/test_access.py

@@ -38,3 +38,9 @@ def test_save_spn_permissions_local(ws, sql_backend, inventory_schema, make_rand
     az_res_perm = AzureResourcePermissions(installation, ws, AzureResources(ws, include_subscriptions=""), location)
     path = az_res_perm.save_spn_permissions()
     assert ws.workspace.get_status(path)
+
+
+def test_cli(ws):
+    from databricks.labs.ucx.cli import save_storage_and_principal
+
+    save_storage_and_principal(ws)

tests/unit/test_cli.py

@@ -17,11 +17,9 @@
     manual_workspace_info,
     move,
     open_remote_config,
+    principal_prefix_access,
     repair_run,
     revert_migrated_tables,
-    save_aws_iam_profiles,
-    save_azure_storage_accounts,
-    save_uc_compatible_roles,
     skip,
     sync_workspace_info,
     validate_external_locations,
@@ -231,91 +229,79 @@ def test_alias(ws):
     ws.tables.list.assert_called_once()
 
 
-def test_save_azure_storage_accounts_not_azure(ws, caplog):
-    ws.config.is_azure = False
-
-    save_azure_storage_accounts(ws, "")
-
-    assert 'Workspace is not on azure, please run this command on azure databricks workspaces.' in caplog.messages
-
-
-def test_save_azure_storage_accounts_no_azure_cli(ws, caplog):
+def test_save_storage_and_principal_azure_no_azure_cli(ws, caplog):
     ws.config.auth_type = "azure_clis"
-
-    save_azure_storage_accounts(ws, "")
+    ws.config.is_azure = True
+    principal_prefix_access(ws, "")
 
     assert 'In order to obtain AAD token, Please run azure cli to authenticate.' in caplog.messages
 
 
-def test_save_azure_storage_accounts_no_subscription_id(ws, caplog):
-    ws.config.auth_type = "azure_cli"
+def test_save_storage_and_principal_azure_no_subscription_id(ws, caplog):
+    ws.config.auth_type = "azure-cli"
     ws.config.is_azure = True
 
-    save_azure_storage_accounts(ws, "")
+    principal_prefix_access(ws, "")
 
     assert "Please enter subscription id to scan storage account in." in caplog.messages
 
 
-def test_save_azure_storage_accounts(ws, caplog):
-    ws.config.auth_type = "azure_cli"
+def test_save_storage_and_principal_azure(ws, caplog, mocker):
+    ws.config.auth_type = "azure-cli"
     ws.config.is_azure = True
-    save_azure_storage_accounts(ws, "test")
-
-    ws.statement_execution.execute_statement.assert_called()
+    azure_resource = mocker.patch("databricks.labs.ucx.azure.access.AzureResourcePermissions.save_spn_permissions")
+    principal_prefix_access(ws, "test")
+    azure_resource.assert_called_once()
 
 
 def test_validate_groups_membership(ws):
     validate_groups_membership(ws)
     ws.groups.list.assert_called()
 
 
-def test_save_aws_iam_profiles_no_profile(ws, caplog, mocker):
+def test_save_storage_and_principal_aws_no_profile(ws, caplog, mocker):
     mocker.patch("shutil.which", return_value="/path/aws")
-    save_aws_iam_profiles(ws)
+    ws.config.is_azure = False
+    ws.config.is_aws = True
+    principal_prefix_access(ws)
     assert any({"AWS Profile is not specified." in message for message in caplog.messages})
 
 
-def test_save_aws_iam_profiles_no_connection(ws, mocker):
+def test_save_storage_and_principal_aws_no_connection(ws, mocker):
     mocker.patch("shutil.which", return_value="/path/aws")
     pop = create_autospec(subprocess.Popen)
-
+    ws.config.is_azure = False
+    ws.config.is_aws = True
     pop.communicate.return_value = (bytes("message", "utf-8"), bytes("error", "utf-8"))
     pop.returncode = 127
     mocker.patch("subprocess.Popen.__init__", return_value=None)
     mocker.patch("subprocess.Popen.__enter__", return_value=pop)
     mocker.patch("subprocess.Popen.__exit__", return_value=None)
 
     with pytest.raises(ResourceWarning, match="AWS CLI is not configured properly."):
-        save_aws_iam_profiles(ws, aws_profile="profile")
+        principal_prefix_access(ws, aws_profile="profile")
 
 
-def test_save_aws_iam_profiles_no_cli(ws, mocker, caplog):
+def test_save_storage_and_principal_aws_no_cli(ws, mocker, caplog):
     mocker.patch("shutil.which", return_value=None)
-    save_aws_iam_profiles(ws, aws_profile="profile")
+    ws.config.is_azure = False
+    ws.config.is_aws = True
+    principal_prefix_access(ws, aws_profile="profile")
     assert any({"Couldn't find AWS" in message for message in caplog.messages})
 
 
-def test_save_uc_roles_no_profile(ws, caplog, mocker):
-    mocker.patch("shutil.which", return_value="/path/aws")
-    save_uc_compatible_roles(ws)
-    assert any({"AWS Profile is not specified." in message for message in caplog.messages})
-
-
-def test_save_uc_roles_no_connection(ws, mocker):
-    mocker.patch("shutil.which", return_value="/path/aws")
-    pop = create_autospec(subprocess.Popen)
-
-    pop.communicate.return_value = (bytes("message", "utf-8"), bytes("error", "utf-8"))
-    pop.returncode = 127
-    mocker.patch("subprocess.Popen.__init__", return_value=None)
-    mocker.patch("subprocess.Popen.__enter__", return_value=pop)
-    mocker.patch("subprocess.Popen.__exit__", return_value=None)
-
-    with pytest.raises(ResourceWarning, match="AWS CLI is not configured properly."):
-        save_uc_compatible_roles(ws, aws_profile="profile")
+def test_save_storage_and_principal_aws(ws, mocker, caplog):
+    mocker.patch("shutil.which", return_value=True)
+    ws.config.is_azure = False
+    ws.config.is_aws = True
+    aws_resource = mocker.patch("databricks.labs.ucx.assessment.aws.AWSResourcePermissions.for_cli")
+    principal_prefix_access(ws, aws_profile="profile")
+    aws_resource.assert_called_once()
 
 
-def test_save_uc_roles_no_cli(ws, mocker, caplog):
-    mocker.patch("shutil.which", return_value=None)
-    save_uc_compatible_roles(ws, aws_profile="profile")
-    assert any({"Couldn't find AWS" in message for message in caplog.messages})
+def test_save_storage_and_principal_gcp(ws, caplog):
+    ws.config.is_azure = False
+    ws.config.is_aws = False
+    ws.config.is_gcp = True
+    principal_prefix_access(ws)
+    assert "This cmd is only supported for azure and aws workspaces" in caplog.messages