update log_uploader and identity_reporter scripts

deployment/identity_reporter.py

@@ -3,10 +3,15 @@
 import subprocess
 import sys
 import os
-from utility import prepare_kosli_trail_name
+from utility import dynamodb_find_session_events
+from utility import dynamodb_get_session_timestamp
 
 
 kosli_flow_name = os.environ.get("KOSLI_FLOW_NAME", "")
+aws_account_id = os.environ.get("AWS_ACCOUNT_ID", "")
+aws_region_sso = os.environ.get("AWS_REGION_SSO", "")
+dynamodb_role_arn = os.environ.get("DYNAMODB_ROLE_ARN", "")
+dynamodb_table_name = os.environ.get("DYNAMODB_TABLE_NAME", "")
 
 
 def describe_ecs_task(cluster, task_arn):
@@ -20,9 +25,16 @@ def lambda_handler(event, context):
         print(f"ECS_EXEC_SESSION_ID is {ecs_exec_session_id}", file=sys.stderr)
 
         ecs_exec_principal_id = event['detail']['userIdentity']['principalId']
-        kosli_trail_name = prepare_kosli_trail_name(ecs_exec_principal_id)
+        ecs_exec_user_name = ecs_exec_principal_id.split(":", 1)[1]
+
+        sso_session_timestamp = dynamodb_get_session_timestamp(ecs_exec_user_name,
+                                                               dynamodb_role_arn,
+                                                               dynamodb_table_name,
+                                                               aws_account_id,
+                                                               aws_region_sso)
 
         # Create Kosli trail (if it is already exists, it will just add a report event to the existing trail)
+        kosli_trail_name = str(sso_session_timestamp) + '-' + ecs_exec_user_name.split("@")[0]
         print(f'Creating Kosli trail {ecs_exec_principal_id} within {kosli_flow_name} flow.')
 
         kosli_client = subprocess.check_output(['./kosli', 'begin', 'trail', kosli_trail_name,
@@ -35,7 +47,8 @@ def lambda_handler(event, context):
             json.dump({"ecs_exec_role_arn": ecs_exec_user_identity}, user_identity_file)
 
         print("Reporting ECS exec user identity to the Kosli...", file=sys.stderr)
-        subprocess.run(['./kosli', 'attest', 'generic', '/tmp/user-identity.json',
+        subprocess.run(['./kosli', 'attest', 'generic', 
+                        f'--attachments=/tmp/user-identity.json',
                         f'--flow={kosli_flow_name}', 
                         f'--trail={kosli_trail_name}',
                         '--name=user-identity',
@@ -51,7 +64,8 @@ def lambda_handler(event, context):
             json.dump({"ecs_exec_service_identity": ecs_exec_task_group}, service_identity_file)
 
         print("Reporting ECS exec service identity to the Kosli...", file=sys.stderr)
-        subprocess.run(['./kosli', 'attest', 'generic', '/tmp/service-identity.json',
+        subprocess.run(['./kosli', 'attest', 'generic',
+                        f'--attachments=/tmp/service-identity.json',
                         f'--flow={kosli_flow_name}', 
                         f'--trail={kosli_trail_name}',
                         '--name=service-identity',

deployment/log_uploader.py

@@ -5,12 +5,17 @@
 import boto3
 import time
 import sys
-from utility import prepare_kosli_trail_name
+from utility import dynamodb_find_session_events
+from utility import dynamodb_get_session_timestamp
 
 
+aws_account_id = os.environ.get("AWS_ACCOUNT_ID", "")
+aws_region_sso = os.environ.get("AWS_REGION_SSO", "")
 kosli_flow_name = os.environ.get("KOSLI_FLOW_NAME", "")
 log_bucket_name = os.environ.get("LOG_BUCKET_NAME", "")
 ecs_exec_event_wait_timeout = os.environ.get("ECS_EXEC_EVENT_WAIT_TIMEOUT", 300)
+dynamodb_role_arn = os.environ.get("DYNAMODB_ROLE_ARN", "")
+dynamodb_table_name = os.environ.get("DYNAMODB_TABLE_NAME", "")
 
 def get_ecs_exec_event_principal_id(ecs_exec_session_id):
     client = boto3.client('cloudtrail')
@@ -52,19 +57,27 @@ def lambda_handler(event, context):
         s3_client.download_file(log_bucket_name, s3_object_key, local_log_file_path)
 
         print('Getting principal id of the current ECS exec session...')
-        ecs_exec_principal_id = get_ecs_exec_event_principal_id(ecs_exec_session_id)
-        kosli_trail_name = prepare_kosli_trail_name(ecs_exec_principal_id)
+        ecs_exec_principal_id = get_ecs_exec_event_principal_id(ecs_exec_session_id) 
+        ecs_exec_user_name = ecs_exec_principal_id.split(":", 1)[1]
+
+        sso_session_timestamp = dynamodb_get_session_timestamp(ecs_exec_user_name,
+                                                               dynamodb_role_arn,
+                                                               dynamodb_table_name,
+                                                               aws_account_id,
+                                                               aws_region_sso)
 
         # Create Kosli trail (if it is already exists, it will just add a report event to the existing trail)
-        print(f'Creating Kosli trail {ecs_exec_principal_id} within {kosli_flow_name} flow.', file=sys.stderr)
+        kosli_trail_name = str(sso_session_timestamp) + '-' + ecs_exec_user_name.split("@")[0]
+        print(f'Creating Kosli trail {kosli_trail_name} within {kosli_flow_name} flow.', file=sys.stderr)
 
         kosli_client = subprocess.run(['./kosli', 'begin', 'trail', kosli_trail_name,
                                                 '--template-file=evidence-template.yml',
                                                 f'--flow={kosli_flow_name}'])
 
         # Upload the log file to the Kosli
         print('Uploading ECS exec log file to the Kosli...', file=sys.stderr)
-        subprocess.run(['./kosli', 'attest', 'generic', local_log_file_path,
+        subprocess.run(['./kosli', 'attest', 'generic', 
+                         f'--attachments={local_log_file_path}',
                          f'--flow={kosli_flow_name}', 
                          f'--trail={kosli_trail_name}',
                          '--name=command-logs'])

deployment/session_saver.py

@@ -11,8 +11,6 @@
 
 sso_elevator_audit_bucket_name = os.environ.get("SSO_ELEVATOR_AUDIT_BUCKET_NAME", "")
 dynamodb_table_name = os.environ.get("DYNAMODB_TABLE_NAME", "")
-# sso_elevator_audit_bucket_name = "sso-elevator-audit-7f6d93c4cf4a8a0fa6ffaddfd70817782bddd202"
-# dynamodb_table_name = 'sso_elevator_session_events'
 
 def lambda_handler(event, context):
     try:
@@ -37,21 +35,22 @@ def lambda_handler(event, context):
 
         print(f'Downloaded SSO event log file. The event operation type is: {session_operation_type}.')
 
+        dynamodb = boto3.resource('dynamodb')
+        dynamodb_table = dynamodb.Table(dynamodb_table_name)
+
         # Put event to the DynamoDB
         if session_operation_type == "grant":
             print('Putting the SSO event to the DynamoDB table...')
-            dynamodb = boto3.resource('dynamodb')
-            table = dynamodb.Table(dynamodb_table_name)
-            response = table.put_item(Item=json.loads(log_file_content))
+            response = dynamodb_table.put_item(Item=json.loads(log_file_content))
             print("DynamoDB PutItem succeeded")
 
         # Remove event from DynamoDB
         elif session_operation_type == "revoke":
             print('Removing the SSO event from the DynamoDB table...')
-            sso_session_events = dynamodb_find_session_events(dynamodb_table_name, 
+            sso_session_events = dynamodb_find_session_events(dynamodb_table, 
                                                               session_event_requester_email, 
                                                               session_event_account_id)
-            dynamodb_delete_session_events(dynamodb_table_name, 
+            dynamodb_delete_session_events(dynamodb_table, 
                                            sso_session_events)
 
         return {
@@ -64,7 +63,3 @@ def lambda_handler(event, context):
             'statusCode': 500,
             'body': 'Handler encountered an error'
         }
-
-# event={"version":"0","id":"122c2cd7-62ea-b734-476e-5d7f10c44239","detail-type":"Object Created","source":"aws.s3","account":"933561452000","time":"2024-02-08T07:53:43Z","region":"eu-north-1","resources":["arn:aws:s3:::sso-elevator-audit-7f6d93c4cf4a8a0fa6ffaddfd70817782bddd202"],"detail":{"version":"0","bucket":{"name":"sso-elevator-audit-7f6d93c4cf4a8a0fa6ffaddfd70817782bddd202"},"object":{"key":"logs/2024/02/08/c6f61bd5-f538-4ccf-86f2-f51d69b49f7d.json","size":447,"etag":"f5ebb1401ceb384333aef699a164454a","version-id":"5D6Oj1YX_yc2BaDzoQhH8SkAYmRIiJpk","sequencer":"0065C488878DD40529"},"request-id":"476X6HYX8Y30NQ47","requester":"933561452000","source-ip-address":"16.171.111.239","reason":"PutObject"}}
-# event={"version":"0","id":"01dfaa72-b7f2-9f22-0111-03d73691600a","detail-type":"Object Created","source":"aws.s3","account":"933561452000","time":"2024-02-08T09:23:58Z","region":"eu-north-1","resources":["arn:aws:s3:::sso-elevator-audit-7f6d93c4cf4a8a0fa6ffaddfd70817782bddd202"],"detail":{"version":"0","bucket":{"name":"sso-elevator-audit-7f6d93c4cf4a8a0fa6ffaddfd70817782bddd202"},"object":{"key":"logs/2024/02/08/07492052-af9b-42ed-a42d-14b733ac0068.json","size":425,"etag":"1851729880c42f2e7809ba9766ce9e28","version-id":"Vs13oJgC_DTHvIE5OkUsVbXsBW1YwKeD","sequencer":"0065C49DAE0659F32C"},"request-id":"VZYRK3SRJCJ1HN4T","requester":"933561452000","source-ip-address":"13.53.104.245","reason":"PutObject"}}
-# lambda_handler(event, dynamodb_table_name, "")

deployment/utility.py

@@ -1,23 +1,7 @@
 import boto3
 
 
-def prepare_kosli_trail_name(principal_id):
-    principal_id_split = principal_id.split(":")
-    prefix = principal_id_split[0]
-    username = principal_id_split[1].split("@")[0]
-    username = username.replace("@", "_")
-    kosli_trail_name = f"{prefix}_{username}"
-
-    return kosli_trail_name
-
-def dynamodb_find_session_events(table_name, requester_email, account_id):
-
-    dynamodb = boto3.resource('dynamodb')
-    table = dynamodb.Table(table_name)
-
-
-    print(table_name, requester_email, account_id)
-
+def dynamodb_find_session_events(table, requester_email, account_id):
     response = table.scan(
         FilterExpression=f'#attribute1_name = :attr1_val AND #attribute2_name = :attr2_val',
         ExpressionAttributeNames={
@@ -32,7 +16,7 @@ def dynamodb_find_session_events(table_name, requester_email, account_id):
 
     return response['Items']
 
-def dynamodb_delete_session_events(table_name, events):
+def dynamodb_delete_session_events(table, events):
     if events:
         print("Matching events found. Deleting...")
 
@@ -49,4 +33,26 @@ def dynamodb_delete_session_events(table_name, events):
     else:
         print("No events found containing both attributes:",
               'requester_email', "with value:", session_event_requester_email,
-              "and", 'account_id', "with value:", session_event_account_id)
+              "and", 'account_id', "with value:", session_event_account_id)
+
+def dynamodb_get_session_timestamp(ecs_exec_user_name, dynamodb_role_arn, dynamodb_table_name, aws_account_id, aws_region_sso):
+    # Get SSO session parameters from
+    sts_client = boto3.client('sts')
+    sts_session = sts_client.assume_role(RoleArn=dynamodb_role_arn, RoleSessionName=f'reporter-dynamodb-session-{aws_account_id}')
+    KEY_ID = sts_session['Credentials']['AccessKeyId']
+    ACCESS_KEY = sts_session['Credentials']['SecretAccessKey']
+    TOKEN = sts_session['Credentials']['SessionToken']
+    dynamodb = boto3.resource('dynamodb',
+                              region_name=aws_region_sso, 
+                              aws_access_key_id=KEY_ID, 
+                              aws_secret_access_key=ACCESS_KEY, 
+                              aws_session_token=TOKEN)
+    dynamodb_table = dynamodb.Table(dynamodb_table_name)
+
+    # There could be multiple sessions, getting the latest
+    sso_session_event = dynamodb_find_session_events(dynamodb_table, 
+                                                     ecs_exec_user_name, 
+                                                     aws_account_id).pop()
+    sso_session_timestamp = int(sso_session_event['timestamp'])
+
+    return sso_session_timestamp

modules/evidence-reporter/iam.tf

@@ -47,3 +47,13 @@ data "aws_iam_policy_document" "cloudtrail_read" {
     resources = ["*"]
   }
 }
+
+data "aws_iam_policy_document" "assume_role_sso" {
+  statement {
+    sid = "assumeRoleSSO"
+    actions = [
+      "sts:AssumeRole"
+    ]
+    resources = [var.dynamodb_role_arn]
+  }
+}

modules/evidence-reporter/identity_reporter.tf

@@ -14,7 +14,7 @@ module "identity_reporter_lambda" {
   recreate_missing_package = var.recreate_missing_package
 
   attach_policy_json = true
-  policy_json        = data.aws_iam_policy_document.ecs_read.json
+  policy_json        = data.aws_iam_policy_document.identity_reporter_combined.json
 
   maximum_retry_attempts = 0
   timeout                = 30
@@ -23,10 +23,14 @@ module "identity_reporter_lambda" {
   role_name              = var.identity_reporter_name
 
   environment_variables = {
-    KOSLI_HOST      = var.kosli_host
-    KOSLI_API_TOKEN = data.aws_ssm_parameter.kosli_api_token.value
-    KOSLI_ORG       = var.kosli_org_name
-    KOSLI_FLOW_NAME = var.kosli_flow_name
+    AWS_ACCOUNT_ID      = data.aws_caller_identity.current.account_id
+    AWS_REGION_SSO      = var.aws_region_sso
+    KOSLI_HOST          = var.kosli_host
+    KOSLI_API_TOKEN     = data.aws_ssm_parameter.kosli_api_token.value
+    KOSLI_ORG           = var.kosli_org_name
+    KOSLI_FLOW_NAME     = var.kosli_flow_name
+    DYNAMODB_ROLE_ARN   = var.dynamodb_role_arn
+    DYNAMODB_TABLE_NAME = var.dynamodb_table_name
   }
 
   allowed_triggers = {
@@ -68,3 +72,11 @@ resource "aws_cloudwatch_event_target" "ecs_exec_session_started" {
   rule      = aws_cloudwatch_event_rule.ecs_exec_session_started.name
   target_id = "${var.identity_reporter_name}-ecs-exec-session-started"
 }
+
+# IAM
+data "aws_iam_policy_document" "identity_reporter_combined" {
+  source_policy_documents = concat(
+    [data.aws_iam_policy_document.ecs_read.json],
+    [data.aws_iam_policy_document.assume_role_sso.json]
+  )
+}

modules/evidence-reporter/log_uploader.tf

@@ -23,11 +23,15 @@ module "log_uploader_lambda" {
   recreate_missing_package = var.recreate_missing_package
 
   environment_variables = {
-    KOSLI_HOST      = var.kosli_host
-    KOSLI_API_TOKEN = data.aws_ssm_parameter.kosli_api_token.value
-    KOSLI_ORG       = var.kosli_org_name
-    KOSLI_FLOW_NAME = var.kosli_flow_name
-    LOG_BUCKET_NAME = var.ecs_exec_log_bucket_name
+    AWS_ACCOUNT_ID      = data.aws_caller_identity.current.account_id
+    AWS_REGION_SSO      = var.aws_region_sso
+    KOSLI_HOST          = var.kosli_host
+    KOSLI_API_TOKEN     = data.aws_ssm_parameter.kosli_api_token.value
+    KOSLI_ORG           = var.kosli_org_name
+    KOSLI_FLOW_NAME     = var.kosli_flow_name
+    LOG_BUCKET_NAME     = var.ecs_exec_log_bucket_name
+    DYNAMODB_ROLE_ARN   = var.dynamodb_role_arn
+    DYNAMODB_TABLE_NAME = var.dynamodb_table_name
   }
 
   allowed_triggers = {
@@ -70,6 +74,7 @@ data "aws_iam_policy_document" "log_uploader_combined" {
   source_policy_documents = concat(
     [data.aws_iam_policy_document.kms_read.json],
     [data.aws_iam_policy_document.s3_read.json],
-    [data.aws_iam_policy_document.cloudtrail_read.json]
+    [data.aws_iam_policy_document.cloudtrail_read.json],
+    [data.aws_iam_policy_document.assume_role_sso.json]
   )
 }

modules/evidence-reporter/variables.tf

@@ -38,8 +38,8 @@ variable "kosli_api_token_ssm_parameter_name" {
 
 variable "kosli_cli_version" {
   type        = string
-  description = "The Kosli cli version, should be set in format 2.6.6"
-  default     = "2.6.6"
+  description = "The Kosli cli version, should be set in format 2.7.8"
+  default     = "2.7.8"
 }
 
 variable "recreate_missing_package" {
@@ -64,4 +64,21 @@ variable "reporter_releases_host" {
   type        = string
   default     = "https://reporter-releases.kosli.com"
   description = "Where to download the Reporter Lambda package."
+}
+
+variable "aws_region_sso" {
+  type        = string
+  default     = "eu-north-1"
+  description = "SSO-elevator region"
+}
+
+variable "dynamodb_role_arn" {
+  type        = string
+  description = "ARN of the role to assume in the SSO aws account"
+}
+
+variable "dynamodb_table_name" {
+  type        = string
+  default     = "sso_elevator_session_events"
+  description = "The name of the SSO events DynamoDB table"
 }

modules/session-saver/dynamodb.tf

@@ -1,13 +1,9 @@
 resource "aws_dynamodb_table" "this" {
-  name           = "sso_elevator_session_events"
-  billing_mode   = "PAY_PER_REQUEST"  # You can change this to "PROVISIONED" if you want provisioned capacity
-  hash_key       = "timestamp"
+  name         = var.dynamodb_table_name
+  billing_mode = "PAY_PER_REQUEST"
+  hash_key     = "timestamp"
   attribute {
     name = "timestamp"
     type = "N"
   }
-#   ttl {
-#     attribute_name = "ttl"
-#     enabled        = true
-#   }
 }

modules/session-saver/session_saver.tf

@@ -7,9 +7,9 @@ module "session_saver_lambda" {
   handler       = "session_saver.lambda_handler"
   runtime       = "python3.11"
 
-  create_package = true
-  publish        = true
-  source_path    = "../deployment"
+  create_package           = true
+  publish                  = true
+  source_path              = "../deployment"
   recreate_missing_package = var.recreate_missing_package
 
   attach_policy_json = true
@@ -22,8 +22,8 @@ module "session_saver_lambda" {
   role_name              = var.session_saver_name
 
   environment_variables = {
-    SSO_ELEVATOR_AUDIT_BUCKET_NAME      = var.sso_elevator_bucket_name
-    DYNAMODB_TABLE_NAME = aws_dynamodb_table.this.id
+    SSO_ELEVATOR_AUDIT_BUCKET_NAME = var.sso_elevator_bucket_name
+    DYNAMODB_TABLE_NAME            = aws_dynamodb_table.this.id
   }
 
   allowed_triggers = {