-
Notifications
You must be signed in to change notification settings - Fork 0
/
runCommand.py
71 lines (54 loc) · 2.23 KB
/
runCommand.py
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
import os
import boto3
import time
validAttacks = ["T15470061", "T11360015", "T10870012"]
def runAttacks(ssmClient, actionId):
runAttacksResponse = ssmClient.send_command(
Targets=[
{
'Key': 'tag:Product',
'Values': [
'WSV1-TechDay2021-2',
]
},
],
DocumentName=actionId,
DocumentVersion='$LATEST',
TimeoutSeconds=600,
Comment='Run Command - Attack',
MaxConcurrency='100%',
MaxErrors='0'
)
print(str(runAttacksResponse))
return runAttacksResponse
def main(event, context):
actionId = str(os.environ.get("ACTION_ID")) if 'ACTION_ID' in os.environ else None
regionName = str(os.environ.get("REGION_NAME")) if 'REGION_NAME' in os.environ else None
instanceId = str(os.environ.get("INSTANCE_ID")) if 'INSTANCE_ID' in os.environ else None
sleepTimer = int(os.environ.get("SLEEP_TIMER")) if 'SLEEP_TIMER' in os.environ else 5
if actionId and regionName:
ssmClient = boto3.client('ssm', region_name=regionName)
runCommandResponse = runAttacks(ssmClient, actionId)
print("Sent Command. Received response - " + str(runCommandResponse))
commandId = runCommandResponse['Command']['CommandId']
time.sleep(sleepTimer)
waiter = ssmClient.get_waiter('command_executed')
waiter.wait(
CommandId=commandId,
InstanceId=instanceId,
WaiterConfig={
'Delay': 5,
'MaxAttempts': 60
}
)
getCommandInvocationResponse = ssmClient.get_command_invocation(
CommandId=commandId,
InstanceId=instanceId
)
print(str(getCommandInvocationResponse))
if getCommandInvocationResponse["Status"] == "Success":
print("Success: " + getCommandInvocationResponse["StandardOutputContent"] + " - " + getCommandInvocationResponse["StandardOutputUrl"])
return(True)
else:
raise Exception("Error: " + getCommandInvocationResponse["StandardErrorContent"] + " - " + getCommandInvocationResponse["StandardErrorUrl"])
return(False)