Add support for Okta group->AWS SSO role rel (cartography-cncf#1307)

AWS SSO role names are weird because they look like
`AWSReservedSSO_myrolename_<somehash>`. This caused our awssaml module
to not draw links from Okta groups to these SSO roles correctly. This PR
updates the module with the correct string comparisons to do this.

Screenshot showing that this works:
![Screenshot 2024-06-03 at 1 49
03 PM](https://github.com/lyft/cartography/assets/46503781/82ef7971-36f3-4f07-ac9c-7d0f856489e2)

cartography/intel/okta/__init__.py

@@ -68,7 +68,7 @@ def start_okta_ingestion(neo4j_session: neo4j.Session, config: Config) -> None:
     applications.sync_okta_applications(neo4j_session, config.okta_org_id, config.update_tag, config.okta_api_key)
     factors.sync_users_factors(neo4j_session, config.okta_org_id, config.update_tag, config.okta_api_key, state)
     origins.sync_trusted_origins(neo4j_session, config.okta_org_id, config.update_tag, config.okta_api_key)
-    awssaml.sync_okta_aws_saml(neo4j_session, config.okta_saml_role_regex, config.update_tag)
+    awssaml.sync_okta_aws_saml(neo4j_session, config.okta_saml_role_regex, config.update_tag, config.okta_org_id)
 
     # need creds with permission
     # soft fail as some won't be able to get such high priv token

cartography/intel/okta/awssaml.py

@@ -1,33 +1,48 @@
 # Okta intel module - AWS SAML
 import logging
 import re
+from collections import namedtuple
 from typing import Dict
 from typing import List
 from typing import Optional
 
 import neo4j
 
+from cartography.client.core.tx import read_list_of_dicts_tx
+from cartography.client.core.tx import read_single_value_tx
 from cartography.util import timeit
 
 
+AccountRole = namedtuple('AccountRole', ['account_id', 'role_name'])
+OktaGroup = namedtuple('OktaGroup', ['group_id', 'group_name'])
+GroupRole = namedtuple('GroupRole', ['okta_group_id', 'aws_role_arn'])
+
 logger = logging.getLogger(__name__)
 
 
 def _parse_regex(regex_string: str) -> str:
     return regex_string.replace("{{accountid}}", "P<accountid>").replace("{{role}}", "P<role>").strip()
 
 
-@timeit
-def transform_okta_group_to_aws_role(group_id: str, group_name: str, mapping_regex: str) -> Optional[Dict]:
+def _parse_okta_group_name(okta_group_name: str, mapping_regex: str) -> AccountRole | None:
+    """
+    Extract AWS account id and AWS role name from the given Okta group name using the given mapping regex.
+    """
     regex = _parse_regex(mapping_regex)
-    matches = re.search(regex, group_name)
+    matches = re.search(regex, okta_group_name)
     if matches:
-        accountid = matches.group("accountid")
-        role = matches.group("role")
-        role_arn = f"arn:aws:iam::{accountid}:role/{role}"
+        account_id = matches.group("accountid")
+        role_name = matches.group("role")
+        return AccountRole(account_id, role_name)
+    return None
+
+
+def transform_okta_group_to_aws_role(group_id: str, group_name: str, mapping_regex: str) -> Optional[Dict]:
+    account_role = _parse_okta_group_name(group_name, mapping_regex)
+    if account_role:
+        role_arn = f"arn:aws:iam::{account_role.account_id}:role/{account_role.role_name}"
         return {"groupid": group_id, "role": role_arn}
-    else:
-        return None
+    return None
 
 
 @timeit
@@ -45,6 +60,7 @@ def query_for_okta_to_aws_role_mapping(neo4j_session: neo4j.Session, mapping_reg
 
     for res in results:
         has_results = True
+        # input: okta group id, okta group name. output: aws role arn.
         mapping = transform_okta_group_to_aws_role(res["group.id"], res["group.name"], mapping_regex)
         if mapping:
             group_to_role_mapping.append(mapping)
@@ -107,8 +123,96 @@ def _load_human_can_assume_role(neo4j_session: neo4j.Session, okta_update_tag: i
     )
 
 
+def get_awssso_okta_groups(neo4j_session: neo4j.Session, okta_org_id: str) -> list[OktaGroup]:
+    """
+    Return list of all Okta group ids in the current Okta organization tied to Okta Applications with name
+    "amazon_aws_sso".
+    """
+    query = """
+    MATCH (g:OktaGroup)-[:APPLICATION]->(a:OktaApplication{name:"amazon_aws_sso"})
+           <-[:RESOURCE]-(:OktaOrganization{id: $okta_org_id})
+    RETURN g.id as group_id, g.name as group_name
+    """
+    result = neo4j_session.read_transaction(read_list_of_dicts_tx, query, okta_org_id=okta_org_id)
+    return [OktaGroup(group_name=og['group_name'], group_id=og['group_id']) for og in result]
+
+
+def get_awssso_role_arn(account_id: str, role_hint: str, neo4j_session: neo4j.Session) -> str | None:
+    """
+    Attempt to return the AWS role ARN for the given AWS account ID and role hint string.
+    This function exists to handle that AWS SSO roles have a 'AWSReservedSSO' prefix and a hashed suffix
+    Input:
+    - account_id: AWS account ID
+    - role_hint (str): The `AccountRole.role_name` returned by _parse_okta_group_name(). This is the part of the Okta
+    group name that refers to the AWS role name.
+    Output:
+    - If we are able to find it, returns the matching AWS role ARN.
+    """
+    query = """
+    MATCH (:AWSAccount{id:$account_id})-[:RESOURCE]->(role:AWSRole{path:"/aws-reserved/sso.amazonaws.com/"})
+    WHERE SPLIT(role.name, '_')[1..-1][0] = $role_hint
+    RETURN role.arn AS role_arn
+    """
+    return neo4j_session.read_transaction(read_single_value_tx, query, account_id=account_id, role_hint=role_hint)
+
+
+def query_for_okta_to_awssso_role_mapping(
+        neo4j_session: neo4j.Session,
+        awssso_okta_groups: list[OktaGroup],
+        mapping_regex: str,
+) -> list[GroupRole]:
+    """
+    Input:
+    - neo4j session
+    - str list of Okta group names
+    - str regex that tells us how to find the AWS role name and account when given an Okta group name
+    Output:
+    - list of OktaGroup id to AWSRole arn pairs.
+    """
+    result = []
+    for group in awssso_okta_groups:
+        account_role = _parse_okta_group_name(group.group_name, mapping_regex)
+        if not account_role:
+            logger.info(f"Okta group {group.group_name} has no associated AWS SSO role")
+            continue
+
+        role_arn = get_awssso_role_arn(account_role.account_id, account_role.role_name, neo4j_session)
+        if role_arn:
+            result.append(GroupRole(group.group_id, role_arn))
+    return result
+
+
+def _load_awssso_tx(tx: neo4j.Transaction, group_to_role: list[GroupRole], okta_update_tag: int) -> None:
+    ingest_statement = """
+    UNWIND $GROUP_TO_ROLE as app_data
+        MATCH (role:AWSRole{arn: app_data.aws_role_arn})
+        MATCH (group:OktaGroup{id: app_data.okta_group_id})
+        MERGE (role)<-[r:ALLOWED_BY]-(group)
+        ON CREATE SET r.firstseen = timestamp()
+        SET r.lastupdated = $okta_update_tag
+    """
+    tx.run(
+        ingest_statement,
+        GROUP_TO_ROLE=[g._asdict() for g in group_to_role],
+        okta_update_tag=okta_update_tag,
+    )
+
+
+def _load_okta_group_to_awssso_roles(
+        neo4j_session: neo4j.Session,
+        group_to_role: list[GroupRole],
+        okta_update_tag: int,
+) -> None:
+    neo4j_session.write_transaction(_load_awssso_tx, group_to_role, okta_update_tag)
+
+
 @timeit
-def sync_okta_aws_saml(neo4j_session: neo4j.Session, mapping_regex: str, okta_update_tag: int) -> None:
+def sync_okta_aws_saml(
+        neo4j_session: neo4j.Session,
+        mapping_regex: str,
+        okta_update_tag: int,
+        okta_org_id: str,
+) -> None:
     """
     Sync okta integration with saml. This will link OktaGroups to the AWSRoles they enable.
     This is for people who use the okta saml provider for AWS
@@ -127,3 +231,7 @@ def sync_okta_aws_saml(neo4j_session: neo4j.Session, mapping_regex: str, okta_up
     group_to_role_mapping = query_for_okta_to_aws_role_mapping(neo4j_session, mapping_regex)
     _load_okta_group_to_aws_roles(neo4j_session, group_to_role_mapping, okta_update_tag)
     _load_human_can_assume_role(neo4j_session, okta_update_tag)
+
+    sso_okta_groups = get_awssso_okta_groups(neo4j_session, okta_org_id)
+    group_to_ssorole_mapping = query_for_okta_to_awssso_role_mapping(neo4j_session, sso_okta_groups, mapping_regex)
+    _load_okta_group_to_awssso_roles(neo4j_session, group_to_ssorole_mapping, okta_update_tag)

docs/root/modules/aws/schema.md

@@ -555,6 +555,12 @@ Representation of an AWS [IAM Role](https://docs.aws.amazon.com/IAM/latest/APIRe
     (AWSRole)-[TRUSTS_AWS_PRINCIPAL]->(AWSPrincipal)
     ```
 
+- Members of an Okta group can assume associated AWS roles if Okta SAML is configured with AWS.
+
+    ```
+    (AWSRole)-[ALLOWED_BY]->(OktaGroup)
+    ```
+
 - AWS Roles are defined in AWS Accounts.
 
     ```

docs/root/modules/okta/schema.md

@@ -130,6 +130,10 @@ Representation of an [Okta Group](https://developer.okta.com/docs/reference/api/
      ```
     (OktaGroup)-[MEMBER_OF_OKTA_ROLE]->(OktaAdministrationRole)
     ```
+- Members of an Okta group can assume associated AWS roles if Okta SAML is configured with AWS.
+    ```
+    (AWSRole)-[ALLOWED_BY]->(OktaGroup)
+    ```
 
 ### OktaApplication
 

tests/integration/cartography/intel/okta/test_awssaml.py

@@ -0,0 +1,123 @@
+from cartography.intel.okta.awssaml import _load_okta_group_to_awssso_roles
+from cartography.intel.okta.awssaml import get_awssso_okta_groups
+from cartography.intel.okta.awssaml import get_awssso_role_arn
+from cartography.intel.okta.awssaml import GroupRole
+from cartography.intel.okta.awssaml import OktaGroup
+from tests.integration.util import check_rels
+
+
+TEST_UPDATE_TAG = 000000
+TEST_ORG_ID = 'ORG_ID'
+DEFAULT_REGEX = r"^aws\#\S+\#(?{{role}}[\w\-]+)\#(?{{accountid}}\d+)$"
+
+
+def test_get_awssso_okta_groups(neo4j_session):
+    # Arrange
+    _ensure_okta_test_data(neo4j_session)
+
+    # Act
+    groups = get_awssso_okta_groups(neo4j_session, TEST_ORG_ID)
+
+    # Assert that the data objects are created correctly
+    assert sorted(groups) == [
+        OktaGroup(group_id='0oaxm1', group_name='AWS_1234_myrole1'),
+        OktaGroup(group_id='0oaxm2', group_name='AWS_1234_myrole2'),
+        OktaGroup(group_id='0oaxm3', group_name='AWS_1234_myrole3'),
+    ]
+
+
+def _ensure_okta_test_data(neo4j_session):
+    test_groups = [
+        ('AWS_1234_myrole1', '0oaxm1'),
+        ('AWS_1234_myrole2', '0oaxm2'),
+        ('AWS_1234_myrole3', '0oaxm3'),
+    ]
+    for group in test_groups:
+        neo4j_session.run(
+            '''
+            MERGE (o:OktaOrganization{id: $ORG_ID})
+            MERGE (o)-[:RESOURCE]-> (g:OktaGroup{name: $GROUP_NAME, id: $GROUP_ID, lastupdated: $UPDATE_TAG})
+            MERGE (o)-[:RESOURCE]->(a:OktaApplication{name:"amazon_aws_sso"})
+            MERGE (a)<-[:APPLICATION]-(g)
+            ''',
+            ORG_ID=TEST_ORG_ID,
+            GROUP_NAME=group[0],
+            GROUP_ID=group[1],
+            UPDATE_TAG=TEST_UPDATE_TAG,
+        )
+
+
+def test_get_awssso_role_arn(neo4j_session):
+    # Arrange
+    _ensure_aws_test_data(neo4j_session)
+
+    # Act and assert
+    assert get_awssso_role_arn(
+        '1234',
+        'myrole1',
+        neo4j_session,
+    ) == 'arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef'
+
+    # Act and assert that we grab the role in the other account correctly
+    assert get_awssso_role_arn(
+        '2345',
+        'myrole1',
+        neo4j_session,
+    ) == 'arn:aws:iam:2345:role/AWSReservedSSO_myrole1_abcdef'
+
+    # Act and assert the None case
+    assert get_awssso_role_arn('1234', 'myrole4', neo4j_session) is None
+
+
+def _ensure_aws_test_data(neo4j_session):
+    # Arrange
+    test_sso_roles = [
+        ('AWSReservedSSO_myrole1_abcdef', 'arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef', '1234'),
+        ('AWSReservedSSO_myrole2_bcdefa', 'arn:aws:iam:1234:role/AWSReservedSSO_myrole2_bcdefa', '1234'),
+        ('AWSReservedSSO_myrole3_cdefab', 'arn:aws:iam:1234:role/AWSReservedSSO_myrole3_cdefab', '1234'),
+        # Add one that has same role name but is in a different account. Expect this to not be returned.
+        ('AWSReservedSSO_myrole1_abcdef', 'arn:aws:iam:2345:role/AWSReservedSSO_myrole1_abcdef', '2345'),
+    ]
+    for role in test_sso_roles:
+        neo4j_session.run(
+            '''
+            MERGE (o:AWSAccount{id: $account_id})
+            MERGE (o)-[:RESOURCE]->
+                  (r1:AWSRole{name: $role_name, id: $arn, arn: $arn, path: $path, lastupdated: $update_tag})
+            ''',
+            role_name=role[0],
+            arn=role[1],
+            id=role[1],
+            account_id=role[2],
+            path='/aws-reserved/sso.amazonaws.com/',
+            update_tag=TEST_UPDATE_TAG,
+        )
+
+
+def test_load_okta_group_to_awssso_roles(neo4j_session):
+    # Arrange
+    _ensure_aws_test_data(neo4j_session)
+    _ensure_okta_test_data(neo4j_session)
+    group_roles = [
+        GroupRole(okta_group_id='0oaxm1', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef'),
+        GroupRole(okta_group_id='0oaxm2', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole2_bcdefa'),
+        GroupRole(okta_group_id='0oaxm3', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole3_cdefab'),
+    ]
+
+    # Act
+    _load_okta_group_to_awssso_roles(neo4j_session, group_roles, TEST_UPDATE_TAG)
+
+    # Assert
+    assert check_rels(
+        neo4j_session,
+        'AWSRole',
+        'id',
+        'OktaGroup',
+        'name',
+        'ALLOWED_BY',
+        False,
+    ) == {
+        ('arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef', 'AWS_1234_myrole1'),
+        ('arn:aws:iam:1234:role/AWSReservedSSO_myrole2_bcdefa', 'AWS_1234_myrole2'),
+        ('arn:aws:iam:1234:role/AWSReservedSSO_myrole3_cdefab', 'AWS_1234_myrole3'),
+    }

tests/unit/cartography/intel/okta/test_awssaml.py

@@ -1,11 +1,23 @@
+from typing import Optional
+from unittest import mock
+from unittest.mock import MagicMock
+
+import cartography.intel.okta.awssaml
+from cartography.intel.okta.awssaml import _parse_okta_group_name
+from cartography.intel.okta.awssaml import AccountRole
+from cartography.intel.okta.awssaml import GroupRole
+from cartography.intel.okta.awssaml import OktaGroup
+from cartography.intel.okta.awssaml import query_for_okta_to_awssso_role_mapping
 from cartography.intel.okta.awssaml import transform_okta_group_to_aws_role
 
 
+SAMPLE_OKTA_GROUP_IDS = ['00g9oh2', '00g9oh3', '00g9oh4']
+
+
 def test_saml_with_default_regex():
     group_name = "aws#northamerica-production#Tier1_Support#828416469395"
     group_id = "groupid"
     default_regex = r"^aws\#\S+\#(?{{role}}[\w\-]+)\#(?{{accountid}}\d+)$"
-
     result = transform_okta_group_to_aws_role(group_id, group_name, default_regex)
 
     assert result
@@ -23,3 +35,46 @@ def test_saml_with_custom_regex():
     assert result
     assert result["groupid"] == group_id
     assert result["role"] == "arn:aws:iam::123456789123:role/developer"
+
+
+def test_parse_okta_group_name() -> None:
+    group_name = 'AWS_1234_myrole1'
+    mapping_regex = r"AWS_(?{{accountid}}\d+)_(?{{role}}[a-zA-Z0-9+=,.@\-_]+)"
+
+    # Act
+    account_role: Optional[AccountRole] = _parse_okta_group_name(group_name, mapping_regex)
+
+    # Assert
+    assert account_role is not None
+    assert account_role.role_name == 'myrole1'
+    assert account_role.account_id == '1234'
+
+
+@mock.patch.object(
+    cartography.intel.okta.awssaml,
+    'get_awssso_role_arn',
+    side_effect=[
+        'arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef',
+        'arn:aws:iam:1234:role/AWSReservedSSO_myrole2_bcdefa',
+        'arn:aws:iam:1234:role/AWSReservedSSO_myrole3_cdefab',
+    ],
+)
+def test_query_for_okta_to_awssso_role_mapping(mock_get_awssso_role_arn: MagicMock) -> None:
+    # Arrange
+    neo4j_session = mock.MagicMock()
+    mapping_regex = r"AWS_(?{{accountid}}\d+)_(?{{role}}[a-zA-Z0-9+=,.@\-_]+)"
+    okta_groups = [
+        OktaGroup(group_id='0oaxm1', group_name='AWS_1234_myrole1'),
+        OktaGroup(group_id='0oaxm2', group_name='AWS_1234_myrole2'),
+        OktaGroup(group_id='0oaxm3', group_name='AWS_1234_myrole3'),
+    ]
+
+    # Act
+    result = query_for_okta_to_awssso_role_mapping(neo4j_session, okta_groups, mapping_regex)
+
+    # Assert
+    assert result == [
+        GroupRole(okta_group_id='0oaxm1', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole1_abcdef'),
+        GroupRole(okta_group_id='0oaxm2', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole2_bcdefa'),
+        GroupRole(okta_group_id='0oaxm3', aws_role_arn='arn:aws:iam:1234:role/AWSReservedSSO_myrole3_cdefab'),
+    ]