feat: add cross-account runbook + actions

ml_ops/sm-datazone_import/README.md

@@ -40,7 +40,7 @@ python import-sagemaker-domain.py \
 - SageMaker execution roles need DataZone API permissions in order for the Assets UI to function. See [DataZoneUserPolicy.json](./resources/DataZoneUserPolicy.json) for an example.
 - Ensure the DataZone Domain trusts SageMaker. In the AWS DataZone console navigate to Domain details and select the "Trusted services".
 
-### Potential errors and workarounds
+### Potential Errors and Workarounds
 
 **Cannot view ML assets in SageMaker Studio, missing "Assets" tab**
 

ml_ops/sm-datazone_import/cross-account-setup.md

@@ -0,0 +1,59 @@
+### Cross Account Setup
+In this scenario, we will use two accounts, a parent account hosting the DataZone domain, and a child account that contains
+a SageMaker Domain nd UserProfile that we would like to link, and import into the DataZone domain. Here AccountA is the parent
+account and AccountB is the child account.
+
+1. **[In the Parent Account]** Create a DataZone domain (make sure you are NOT using the Unified UI). It should say “Create a DataZone Domain”.
+2. **[In the Parent Account]** Create an association to the X-Account. 
+
+ - Click in to Domain - request Association by providing the the X-Account number.
+ - This will create a AWSRAMPermissionDataZoneDefault policy to allow access from the X-account. 
+ - In the X-Account, Accept the resource Share in the DataZone UI (Unified UI)
+
+3. **[In the Parent Account]** Create A Project, select the parent domain as the “DomainUnit”.
+4. **[In the Parent Account]** Add in the X-account user that will access this project.
+
+ - In the UserManagement tab, add the child user’s role (as an IAM user or SSO user) that requires access, from the X-account. Choose the AssociatedAccount option. 
+ - In the Projects tab, add the child user as a ProjectMember. They should be available in the DropDown menu. Set the respective permissions. 
+
+> NOTE :: If this step (a.) is not done, the X-account user will see this error in their projects tab when clicking into the Assocated Domain  
+
+```
+Not a DataZone user
+You cannot view or create a project because you have not been added 
+as a Amazon DataZone user. Please contact your domain admin to add 
+your IAM role: arn:aws:iam::211125770549:role/Admin as a DataZone user.
+```
+
+5. **[In the Child Account]** Create a SageMaker Domain. Ensure that you do this from the Amazon SageMaker AI console, not the Amazon SageMaker platform console (this is the unified experience, separate from this current workflow). Add users profiles to the domain
+6. **[In the Child Account]** Setup a federation role that will have permission to federate into our parent account’s Datazone portal. See `/resources` for examples of trust and permission policies. 
+
+### Running the script
+For linking the SageMaker Domain + UserProfile using HULK BYOD Flow:
+
+Make sure that the current account you are using grants access to the X-account to sts:AssumeRole. 
+For Example, AccountA is the one that houses the SageMaker Domain and UserProfile that you would like to import into the DataZone parent account, AccountB. 
+
+* We need to be sure to add the following JSON to the TrustPolicy of the Admin (or whatever role in the parent account you’d like to assume, with DataZone permissions to call batch-put-linked-types and link the SageMaker Domain and UserProfiles). 
+* Also, make sure to add the User that the current session is using.
+
+In the parent account (Account A), under the currently assumed role - we should add the following.
+This will allow our child account to assume parent account role during our current session, and link the SageMaker Domain + UserProfile that we interact with while running the script.
+
+```
+  {
+      "Sid": "",
+      "Effect": "Allow",
+      "Principal": {
+          "AWS": "arn:aws:sts::<Account_A>:assumed-role/Admin/<user>-Isengard"
+      },
+      "Action": "sts:AssumeRole"
+  }
+```
+
+* Nothing changes with regards to the regular batch-put flow. We will assume the parent account (AccountA) credentials. 
+* From the Parent Account, the script will call batch-put-linked type using the SageMaker ARN and SageMaker UserProfile ARN from the X-account (AccountB).
+* Federation link will then work for X-account and Parent Account from DZ portal → environment view.
+
+
+

ml_ops/sm-datazone_import/import-sagemaker-domain.py

@@ -1,7 +1,8 @@
 import argparse
 import boto3
 import botocore
-import time
+import datetime
+from dateutil.tz import tzlocal
 from botocore.exceptions import ClientError
 
 """
@@ -13,6 +14,24 @@
 """
 
 
+def assumed_role_session(role_arn: str, base_session: botocore.session.Session = None):
+    base_session = base_session or boto3.session.Session()._session
+    fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
+        client_creator=base_session.create_client,
+        source_credentials=base_session.get_credentials(),
+        role_arn=role_arn,
+        extra_args={},
+    )
+    creds = botocore.credentials.DeferredRefreshableCredentials(
+        method="assume-role",
+        refresh_using=fetcher.fetch_credentials,
+        time_fetcher=lambda: datetime.datetime.now(tzlocal()),
+    )
+    botocore_session = botocore.session.Session()
+    botocore_session._credentials = creds
+    return boto3.Session(botocore_session=botocore_session)
+
+
 class SageMakerDomainImporter:
     def __init__(self, region, stage, account_id) -> None:
         self.region = region
@@ -41,6 +60,7 @@ def __init__(self, region, stage, account_id) -> None:
             "datazone-byod", region_name=region, endpoint_url=dz_endpoint_url
         )
         self.iam_client = boto3.client("iam", region_name=region)
+        self.sts_client = boto3.client("sts", region_name=region)
 
     def _choose_sm_domain(self):
         # [1] First, identify the SageMaker Domain needed to be imported by noting its DomainId.
@@ -398,9 +418,27 @@ def _associate_fed_role(self):
             else:
                 print(f"Caught error: {repr(e)}")
 
+    def _cross_account_action(self):
+        # Decision to link domain from X-account. Required to assumeRole from child-account containing SM domain.
+
+        decision = input("Need to import SM profiles from a X-account? [y/n]: ")
+        if decision == "y":
+            x_account_role = input(
+                "X-account role ARN to assume? This would be the account that originally created"
+                "the DataZone domain, aka the parent account: "
+            )
+            session = assumed_role_session(x_account_role)
+            dz_endpoint_url = "https://datazone." + region + ".api.aws"  # prod
+            self.dz_client = session.client(
+                "datazone", region_name=region, endpoint_url=dz_endpoint_url
+            )
+            self.byod_client = session.client(
+                "datazone-byod", region_name=region, endpoint_url=dz_endpoint_url
+            )
+
     def _link_domain(self):
         # attach SAGEMAKER_DOMAIN
-        linkedDomainItems = [
+        linked_domain_items = [
             {
                 "itemIdentifier": f"arn:aws:sagemaker:{self.region}:{self.account_id}:domain/{self.sm_domain_id}",
                 "itemType": "SAGEMAKER_DOMAIN",
@@ -429,7 +467,7 @@ def _link_domain(self):
             domainIdentifier=self.dz_domain_id,
             projectIdentifier=self.dz_project_id,
             environmentIdentifier=self.env_id,
-            items=linkedDomainItems,
+            items=linked_domain_items,
         )
         print(link_domain_response)
         print("Linked SageMaker Domain.")
@@ -516,6 +554,7 @@ def import_interactive(self):
         self._map_users()
         self._associate_fed_role()
         self._add_environment_action()
+        self._cross_account_action()
         self._link_domain()
         self._link_users()
         self._debug_print_results()