security_content/automated_detection_testing/ci/python_ci_code/main.py at develop · RedCyberShield/security_content · 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

import uuid
import sys
import boto3
import argparse
import time


def main(args):

    parser = argparse.ArgumentParser(description="Detection Testing Execution")
    parser.add_argument("-b", "--branch", required=True, help="security content branch")
    parser.add_argument("-pr", "--pr-number", required=False, help="Pull Request Number")

    args = parser.parse_args()
    branch = args.branch
    pr_number = args.pr_number

    # vars
    max_waiting_time = 7200
    current_waiting_time = 0

    # create uuid
    uuid_test = str(uuid.uuid4())

    # start aws batch job
    client = boto3.client("batch", region_name="eu-central-1")

    if pr_number:
        response = client.submit_job(
            jobName='detection_testing_' + branch,
            jobQueue='detection_testing_execution_queue',
            jobDefinition='detection_testing_execution:2',
            containerOverrides={
                'command': ['-b', branch, '-u', uuid_test, '-pr', pr_number]
            }
        )
    else:
        response = client.submit_job(
            jobName='detection_testing_' + branch,
            jobQueue='detection_testing_execution_queue',
            jobDefinition='detection_testing_execution:2',
            containerOverrides={
                'command': ['-b', branch, '-u', uuid_test]
            }
        )


    while max_waiting_time > current_waiting_time:

        dynamodb = boto3.client('dynamodb', region_name="eu-central-1")
        response = dynamodb.query(
            TableName='dt-results',
            IndexName='uuid_test-index',
            KeyConditionExpression='uuid_test = :uuid_test',
            ExpressionAttributeValues={
                    ':uuid_test': {'S': uuid_test}
            }
        )


        test_done = True
        for item in response['Items']:
            if item['status']['S'] == 'nothing to test':
                print('No new or changed detections to test.')
                sys.exit(0)

            if item['status']['S'] == 'running':
                test_done = False

        if len(response['Items']) == 0 or (not test_done):
            time.sleep(60)
            current_waiting_time = current_waiting_time + 60
        else:
            test_passed = True
            for item in response['Items']:
                if item['result']['S'] == 'failed':
                    test_passed = False
                    print('Test failed for detection: ' + item['detection']['S'] + ', ' + item['detection_path']['S'])
                else:
                    print('Test passed for detection: ' + item['detection']['S'] + ', ' + item['detection_path']['S'])
            sys.exit(not test_passed)

    sys.exit(1)


if __name__ == "__main__":
    main(sys.argv[1:])