SentinelAI/test_lambda.py at main · playroom-security/SentinelAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#!/usr/bin/env python3
"""
Test script for SentinelAI Lambda function routing.
Tests the main lambda_function.py handler with different event types.
"""

import sys
import json
sys.path.insert(0, '/Users/dcyberguy/Projects/AI/SentinelAI')
sys.path.insert(0, '/Users/dcyberguy/Projects/AI/SentinelAI/sentinelai_lambda')

from lambda_function import lambda_handler


def test_iam_analysis():
    """Test IAM policy analysis routing."""
    print("🧪 Testing IAM Policy Analysis...")

    event = {
        "policy": {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": "*",
                    "Resource": "*"
                }
            ]
        }
    }

    result = lambda_handler(event, None)
    print(f"Status Code: {result['statusCode']}")
    if result['statusCode'] == 200:
        body = json.loads(result['body'])
        print(f"Issues Found: {body.get('issues_found', 'N/A')}")
        print("✅ IAM Analysis: PASSED")
    else:
        print("❌ IAM Analysis: FAILED")
    print()


def test_cloudtrail_analysis():
    """Test CloudTrail log analysis routing."""
    print("🧪 Testing CloudTrail Analysis...")

    event = {
        "logs": [
            {
                "eventName": "CreateAccessKey",
                "sourceIPAddress": "203.0.113.42",
                "userIdentity": {"type": "IAMUser", "userName": "test-user"}
            }
        ]
    }

    result = lambda_handler(event, None)
    print(f"Status Code: {result['statusCode']}")
    if result['statusCode'] == 200:
        body = json.loads(result['body'])
        print(f"Events Analyzed: {body.get('events_analyzed', 'N/A')}")
        print("✅ CloudTrail Analysis: PASSED")
    else:
        print("❌ CloudTrail Analysis: FAILED")
    print()


def test_invalid_event():
    """Test invalid event handling."""
    print("🧪 Testing Invalid Event Handling...")

    event = {"invalid": "data"}

    result = lambda_handler(event, None)
    print(f"Status Code: {result['statusCode']}")
    if result['statusCode'] == 400:
        body = json.loads(result['body'])
        print(f"Error: {body.get('error', 'N/A')}")
        print("✅ Invalid Event Handling: PASSED")
    else:
        print("❌ Invalid Event Handling: FAILED")
    print()


def test_analysis_type_routing():
    """Test analysis_type routing."""
    print("🧪 Testing Analysis Type Routing...")

    # Test IAM via analysis_type
    event = {
        "analysis_type": "iam",
        "policy": {
            "Version": "2012-10-17",
            "Statement": [{"Effect": "Allow", "Action": "s3:GetObject", "Resource": "*"}]
        }
    }

    result = lambda_handler(event, None)
    print(f"IAM via analysis_type - Status Code: {result['statusCode']}")

    # Test CloudTrail via analysis_type
    event = {
        "analysis_type": "cloudtrail",
        "logs": [{"eventName": "ListUsers", "sourceIPAddress": "192.168.1.1"}]
    }

    result = lambda_handler(event, None)
    print(f"CloudTrail via analysis_type - Status Code: {result['statusCode']}")

    print("✅ Analysis Type Routing: PASSED")
    print()


if __name__ == "__main__":
    print("🚀 SentinelAI Lambda Function Tests")
    print("=" * 50)

    test_iam_analysis()
    test_cloudtrail_analysis()
    test_invalid_event()
    test_analysis_type_routing()

    print("🎉 All tests completed!")