feat(engine): AWS CloudTrail gzipped json reader (#63)

TracecatHQ · Apr 20, 2024 · fe30744 · fe30744
1 parent a341de7
commit fe30744
Show file tree

Hide file tree

Showing 15 changed files with 1,018 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,6 +23,7 @@ classifiers = [
 ]
 dependencies = [
     "adbc-driver-sqlite",
+    "boto3==1.34.69",
     "colorlog",
     "croniter",
     "cryptography",
@@ -31,7 +32,7 @@ dependencies = [
     "lancedb==0.6.3",
     "openai",
     "orjson",
-    "polars",
+    "polars==0.20.21",
     "psycopg[binary]",
     "psycopg2-binary",
     "pyarrow",
@@ -54,7 +55,7 @@ Repository = "https://github.com/TracecatHQ/tracecat"
 
 [project.optional-dependencies]
 runner = ["aiosmtplib", "jsonpath_ng", "python-multipart"]
-dev = ["respx", "pytest", "python-dotenv", "pytest-asyncio"]
+dev = ["respx", "pytest", "python-dotenv", "pytest-asyncio", "minio"]
 
 [tool.hatch.version]
 path = "tracecat/__init__.py"

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,20 +1,128 @@
+"""Set up shared environment variables and S3 proxy server (MinIO) for integration tests.
+"""
+
+import logging
 import os
+import subprocess
+import time
 
 import pytest
 from cryptography.fernet import Fernet
+from minio import Minio
+
+# MinIO settings
+MINIO_CONTAINER_NAME = "minio_test_server"
+MINIO_ACCESS_KEY = "admin"
+MINIO_SECRET_KEY = "password"
+MINIO_PORT = 9000
+MINIO_REGION = "us-west-2"
+
+
+# S3 log sample settings
+AWS_CLOUDTRAIL__BUCKET_NAME = "aws-cloudtrail-logs"
 
 
-@pytest.fixture(autouse=True)
+logging.basicConfig(level="INFO")
+
+
+@pytest.fixture(autouse=True, scope="session")
 def setup_shared_env():
     os.environ["TRACECAT__DB_ENCRYPTION_KEY"] = Fernet.generate_key().decode()
     os.environ["TRACECAT__API_URL"] = "http://api:8000"
     os.environ["TRACECAT__RUNNER_URL"] = "http://runner:8000"
     os.environ["TRACECAT__SERVICE_KEY"] = "test-service-key"
 
+    # AWS
+    os.environ["AWS_CLOUDTRAIL__BUCKET_NAME"] = AWS_CLOUDTRAIL__BUCKET_NAME
+
+    # MinIO Client (local)
+    os.environ["MINIO_ACCESS_KEY"] = MINIO_ACCESS_KEY
+    os.environ["MINIO_SECRET_KEY"] = MINIO_SECRET_KEY
+    os.environ["MINIO_ENDPOINT"] = f"localhost:{MINIO_PORT}"
+
     try:
         yield
     finally:
         del os.environ["TRACECAT__DB_ENCRYPTION_KEY"]
         del os.environ["TRACECAT__API_URL"]
         del os.environ["TRACECAT__RUNNER_URL"]
         del os.environ["TRACECAT__SERVICE_KEY"]
+        del os.environ["MINIO_ACCESS_KEY"]
+        del os.environ["MINIO_SECRET_KEY"]
+        del os.environ["MINIO_ENDPOINT"]
+
+
+@pytest.fixture(scope="session", autouse=True)
+def minio_container():
+    # Check if the MinIO container is already running
+    existing_containers = subprocess.run(
+        [
+            "docker",
+            "ps",
+            "--filter",
+            f"name={MINIO_CONTAINER_NAME}",
+            "--format",
+            "{{.Names}}",
+        ],
+        capture_output=True,
+        text=True,
+    )
+
+    container_exists = MINIO_CONTAINER_NAME in existing_containers.stdout.strip()
+    logging.info("🐳 MinIO container exists: %r", container_exists)
+
+    if not container_exists:
+        # Setup: Start MinIO server
+        subprocess.run(
+            [
+                "docker",
+                "run",
+                "-d",
+                "--rm",
+                "--name",
+                MINIO_CONTAINER_NAME,
+                "-p",
+                f"{MINIO_PORT}:{MINIO_PORT}",
+                "-e",
+                f"MINIO_ACCESS_KEY={MINIO_ACCESS_KEY}",
+                "-e",
+                f"MINIO_SECRET_KEY={MINIO_SECRET_KEY}",
+                "minio/minio",
+                "server",
+                "/data",
+            ],
+            check=True,
+        )
+        # Wait for the server to start
+        time.sleep(5)
+        logging.info("✅ Created minio container %r", MINIO_CONTAINER_NAME)
+    else:
+        logging.info("✅ Using existing minio container %r", MINIO_CONTAINER_NAME)
+
+    # Connect to MinIO
+    client = Minio(
+        f"localhost:{MINIO_PORT}",
+        access_key=MINIO_ACCESS_KEY,
+        secret_key=MINIO_SECRET_KEY,
+        secure=False,
+    )
+
+    # Create or connect to AWS CloudTrail bucket
+    bucket = AWS_CLOUDTRAIL__BUCKET_NAME
+    if not client.bucket_exists(bucket):
+        client.make_bucket(bucket)
+        logging.info("✅ Created minio bucket %r", bucket)
+
+    yield
+    should_cleanup = os.getenv("MINIO_CLEANUP", "1").lower() in (
+        "true",
+        "1",
+    )
+    if not container_exists and should_cleanup:
+        logging.info("🧹 Cleaning up minio container %r", MINIO_CONTAINER_NAME)
+        subprocess.run(["docker", "stop", MINIO_CONTAINER_NAME], check=True)
+    else:
+        logging.info(
+            "🧹 Skipping cleanup of minio container %r. Set `MINIO_CLEANUP=1` to cleanup.",
+            MINIO_CONTAINER_NAME,
+        )
diff --git a/tests/data/log_samples/aws_cloudtrail/cloudtrail_insights.json b/tests/data/log_samples/aws_cloudtrail/cloudtrail_insights.json
@@ -0,0 +1,55 @@
+{
+    "Records": [{
+        "eventVersion": "1.08",
+        "eventTime": "2023-01-02T02:51:00Z",
+        "awsRegion": "us-east-1",
+        "eventID": "654a30ff-b0f3-4527-81b6-EXAMPLEf2393",
+        "eventType": "AwsCloudTrailInsight",
+        "recipientAccountId": "123456789012",
+        "sharedEventID": "bcbfc274-8559-4a56-beb0-EXAMPLEa6c34",
+        "insightDetails": {
+            "state": "Start",
+            "eventSource": "ssm.amazonaws.com",
+            "eventName": "UpdateInstanceInformation",
+            "insightType": "ApiCallRateInsight",
+            "insightContext": {
+                "statistics": {
+                    "baseline": {
+                        "average": 84.410596421
+                    },
+                    "insight": {
+                        "average": 669
+                    }
+                }
+            }
+        },
+        "eventCategory": "Insight"
+    },
+    {
+        "eventVersion": "1.08",
+        "eventTime": "2023-01-02T00:22:00Z",
+        "awsRegion": "us-east-1",
+        "eventID": "258de2fb-e2a9-4fb5-aeb2-EXAMPLE449a4",
+        "eventType": "AwsCloudTrailInsight",
+        "recipientAccountId": "123456789012",
+        "sharedEventID": "8b74a7bc-d5d3-4d19-9d60-EXAMPLE08b51",
+        "insightDetails": {
+            "state": "End",
+            "eventSource": "ssm.amazonaws.com",
+            "eventName": "UpdateInstanceInformation",
+            "insightType": "ApiCallRateInsight",
+            "insightContext": {
+                "statistics": {
+                    "baseline": {
+                        "average": 74.156423842
+                    },
+                    "insight": {
+                        "average": 657
+                    },
+                    "insightDuration": 1
+                }
+            }
+        },
+        "eventCategory": "Insight"
+    }]
+}
diff --git a/tests/data/log_samples/aws_cloudtrail/cloudtrail_update_trail_error.json b/tests/data/log_samples/aws_cloudtrail/cloudtrail_update_trail_error.json
@@ -0,0 +1,43 @@
+{"Records": [{
+    "eventVersion": "1.09",
+    "userIdentity": {
+        "type": "IAMUser",
+        "principalId": "AIDA6ON6E4XEGIEXAMPLE",
+        "arn": "arn:aws:iam::111122223333:user/Terry",
+        "accountId": "111122223333",
+        "accessKeyId": "AKIAIOSFODNN7EXAMPLE",
+        "userName": "Terry",
+        "sessionContext": {
+            "attributes": {
+                "creationDate": "2023-07-19T21:11:57Z",
+                "mfaAuthenticated": "false"
+            }
+        }
+    },
+    "eventTime": "2023-07-19T21:35:03Z",
+    "eventSource": "cloudtrail.amazonaws.com",
+    "eventName": "UpdateTrail",
+    "awsRegion": "us-east-1",
+    "sourceIPAddress": "192.0.2.0",
+    "userAgent": "aws-cli/2.13.0 Python/3.11.4 Linux/4.14.255-314-253.539.amzn2.x86_64 exec-env/CloudShell exe/x86_64.amzn.2 prompt/off command/cloudtrail.update-trail",
+    "errorCode": "TrailNotFoundException",
+    "errorMessage": "Unknown trail: arn:aws:cloudtrail:us-east-1:111122223333:trail/myTrail2 for the user: 111122223333",
+    "requestParameters": {
+        "name": "myTrail2",
+        "isMultiRegionTrail": true
+    },
+    "responseElements": null,
+    "requestID": "28d2faaf-3319-4649-998d-EXAMPLE72818",
+    "eventID": "694d604a-d190-4470-8dd1-EXAMPLEe20c1",
+    "readOnly": false,
+    "eventType": "AwsApiCall",
+    "managementEvent": true,
+    "recipientAccountId": "111122223333",
+    "eventCategory": "Management",
+    "tlsDetails": {
+        "tlsVersion": "TLSv1.2",
+        "cipherSuite": "ECDHE-RSA-AES128-GCM-SHA256",
+        "clientProvidedHostHeader": "cloudtrail.us-east-1.amazonaws.com"
+    },
+    "sessionCredentialFromConsole": "true"
+}]}
diff --git a/tests/data/log_samples/aws_cloudtrail/ec2_create_key_pair.json b/tests/data/log_samples/aws_cloudtrail/ec2_create_key_pair.json
@@ -0,0 +1,50 @@
+{"Records": [{
+    "eventVersion": "1.08",
+    "userIdentity": {
+        "type": "IAMUser",
+        "principalId": "AIDA6ON6E4XEGIEXAMPLE",
+        "arn": "arn:aws:iam::444455556666:user/Arnav",
+        "accountId": "444455556666",
+        "accessKeyId": "AKIAI44QH8DHBEXAMPLE",
+        "userName": "Arnav",
+        "sessionContext": {
+            "sessionIssuer": {},
+            "webIdFederationData": {},
+            "attributes": {
+                "creationDate": "2023-07-19T21:11:57Z",
+                "mfaAuthenticated": "false"
+            }
+        }
+    },
+    "eventTime": "2023-07-19T21:19:22Z",
+    "eventSource": "ec2.amazonaws.com",
+    "eventName": "CreateKeyPair",
+    "awsRegion": "us-east-1",
+    "sourceIPAddress": "192.0.2.0",
+    "userAgent": "aws-cli/2.13.5 Python/3.11.4 Linux/4.14.255-314-253.539.amzn2.x86_64 exec-env/CloudShell exe/x86_64.amzn.2 prompt/off command/ec2.create-key-pair",
+    "requestParameters": {
+        "keyName": "my-key",
+        "keyType": "rsa",
+        "keyFormat": "pem"
+    },
+    "responseElements": {
+        "requestId": "9aa4938f-720f-4f4b-9637-EXAMPLE9a196",
+        "keyName": "my-key",
+        "keyFingerprint": "1f:51:ae:28:bf:89:e9:d8:1f:25:5d:37:2d:7d:b8:ca:9f:f5:f1:6f",
+        "keyPairId": "key-abcd12345eEXAMPLE",
+        "keyMaterial": "<sensitiveDataRemoved>"
+    },
+    "requestID": "9aa4938f-720f-4f4b-9637-EXAMPLE9a196",
+    "eventID": "2ae450ff-e72b-4de1-87b0-EXAMPLE5227cb",
+    "readOnly": false,
+    "eventType": "AwsApiCall",
+    "managementEvent": true,
+    "recipientAccountId": "444455556666",
+    "eventCategory": "Management",
+    "tlsDetails": {
+        "tlsVersion": "TLSv1.2",
+        "cipherSuite": "ECDHE-RSA-AES128-GCM-SHA256",
+        "clientProvidedHostHeader": "ec2.us-east-1.amazonaws.com"
+    },
+    "sessionCredentialFromConsole": "true"
+}]}
diff --git a/tests/data/log_samples/aws_cloudtrail/ec2_start_instances.json b/tests/data/log_samples/aws_cloudtrail/ec2_start_instances.json
@@ -0,0 +1,79 @@
+{"Records": [{
+    "eventVersion": "1.08",
+    "userIdentity": {
+        "type": "IAMUser",
+        "principalId": "EXAMPLE6E4XEGITWATV6R",
+        "arn": "arn:aws:iam::123456789012:user/Mateo",
+        "accountId": "123456789012",
+        "accessKeyId": "AKIAIOSFODNN7EXAMPLE",
+        "userName": "Mateo",
+        "sessionContext": {
+            "sessionIssuer": {},
+            "webIdFederationData": {},
+            "attributes": {
+                "creationDate": "2023-07-19T21:11:57Z",
+                "mfaAuthenticated": "false"
+            }
+        }
+    },
+    "eventTime": "2023-07-19T21:17:28Z",
+    "eventSource": "ec2.amazonaws.com",
+    "eventName": "StartInstances",
+    "awsRegion": "us-east-1",
+    "sourceIPAddress": "192.0.2.0",
+    "userAgent": "aws-cli/2.13.5 Python/3.11.4 Linux/4.14.255-314-253.539.amzn2.x86_64 exec-env/CloudShell exe/x86_64.amzn.2 prompt/off command/ec2.start-instances",
+    "requestParameters": {
+        "instancesSet": {
+            "items": [
+                {
+                    "instanceId": "i-EXAMPLE56126103cb"
+                },
+                {
+                    "instanceId": "i-EXAMPLEaff4840c22"
+                }
+            ]
+        }
+    },
+    "responseElements": {
+        "requestId": "e4336db0-149f-4a6b-844d-EXAMPLEb9d16",
+        "instancesSet": {
+            "items": [
+                {
+                    "instanceId": "i-EXAMPLEaff4840c22",
+                    "currentState": {
+                        "code": 0,
+                        "name": "pending"
+                    },
+                    "previousState": {
+                        "code": 80,
+                        "name": "stopped"
+                    }
+                },
+                {
+                    "instanceId": "i-EXAMPLE56126103cb",
+                    "currentState": {
+                        "code": 0,
+                        "name": "pending"
+                    },
+                    "previousState": {
+                        "code": 80,
+                        "name": "stopped"
+                    }
+                }
+            ]
+        }
+    },
+    "requestID": "e4336db0-149f-4a6b-844d-EXAMPLEb9d16",
+    "eventID": "e755e09c-42f9-4c5c-9064-EXAMPLE228c7",
+    "readOnly": false,
+    "eventType": "AwsApiCall",
+    "managementEvent": true,
+    "recipientAccountId": "123456789012",
+    "eventCategory": "Management",
+     "tlsDetails": {
+        "tlsVersion": "TLSv1.2",
+        "cipherSuite": "ECDHE-RSA-AES128-GCM-SHA256",
+        "clientProvidedHostHeader": "ec2.us-east-1.amazonaws.com"
+    },
+    "sessionCredentialFromConsole": "true"
+}]}