Skip to content

Commit 5b08095

Browse files
duchatranduchatran
committed
First checkin of Python inspector function
1 parent 6b8794d commit 5b08095

File tree

1 file changed

+175
-0
lines changed

1 file changed

+175
-0
lines changed

‎inspector/python/inspector.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
2+
import json
3+
import httplib
4+
import base64,zlib
5+
import urlparse
6+
import boto3
7+
import datetime
8+
import logging
9+
10+
##################################################################
11+
# Configuration #
12+
##################################################################
13+
# Enter Sumo Http source endpoint here.
14+
sumoEndpoint = "https://endpoint1.collection.sumologic.com/receiver/v1/http/<XXXX>"
15+
# include auxiliary data (e.g for assessment template, run, or target) in the collected event or not
16+
contextLookup = True
17+
18+
19+
##################################################################
20+
# Main Code #
21+
##################################################################
22+
up = urlparse.urlparse(sumoEndpoint)
23+
options = { 'hostname': up.hostname,
24+
'path': up.path,
25+
'method': 'POST'
26+
};
27+
28+
# Internal variables used for this Lambda function
29+
resourceMap = {'finding':{},'target':{},'run':{},'template':{}, 'rulesPackage':{}}
30+
# prepare logger
31+
logger = logging.getLogger()
32+
logger.setLevel(logging.INFO)
33+
34+
# main function to send data to a Sumo HTTP source
35+
def sendSumo(msg, toCompress = False):
36+
conn = httplib.HTTPSConnection(options['hostname'])
37+
if (toCompress):
38+
headers = {"Content-Encoding": "gzip"}
39+
finalData = compress(msg)
40+
else:
41+
headers = {"Content-type": "text/html","Accept": "text/plain"}
42+
finalData =msg
43+
conn.request(options['method'], options['path'], finalData,headers)
44+
response = conn.getresponse()
45+
conn.close()
46+
return (response.status,response.reason)
47+
48+
49+
# Simple function to compress data
50+
def compress(data, compresslevel=9):
51+
compress = zlib.compressobj(compresslevel, zlib.DEFLATED, 16 + zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
52+
compressedData = compress.compress(data)
53+
compressedData += compress.flush()
54+
return compressedData
55+
56+
# This function looks up an Inspector object based on its arn and type. Returned object will be used to provide extra context for the final message to Sumo
57+
def lookup(objectId,objectType = 'run'):
58+
client = boto3.client('inspector')
59+
finalObj = None
60+
61+
objectMap = resourceMap.get(objectType)
62+
if (objectMap is None):
63+
resourceMap[objectType]= objectMap = {}
64+
try:
65+
if (objectType=='run'):
66+
run = objectMap.get(objectId)
67+
if (run is None):
68+
runs = client.describe_assessment_runs(assessmentRunArns=[objectId])
69+
if (runs is not None):
70+
run = runs['assessmentRuns'][0]
71+
# For run item, we only collect important properties
72+
objectMap[objectId] = finalObj = {'name':run['name'],'createdAt':'%s' % run['createdAt'], 'state':run['state'],'durationInSeconds':run['durationInSeconds'],'startedAt':'%s' % run['startedAt'],'assessmentTemplateArn':run['assessmentTemplateArn']}
73+
else:
74+
finalObj = run
75+
elif (objectType=='template'):
76+
template = objectMap.get(objectId)
77+
if (template is None):
78+
templates = client.describe_assessment_templates(assessmentTemplateArns=[objectId])
79+
if (templates is not None):
80+
finalObj = objectMap[objectId] = templates['assessmentTemplates'][0]
81+
else:
82+
finalObj = template
83+
elif (objectType=='rulesPackage'):
84+
rulesPackage = objectMap.get(objectId)
85+
if (rulesPackage is None):
86+
rulesPackages = client.describe_rules_packages(rulesPackageArns=[objectId])
87+
if (rulesPackages is not None):
88+
finalObj = objectMap[objectId] = rulesPackages['rulesPackages'][0]
89+
else:
90+
finalObj = rulesPackage
91+
elif (objectType=='target'):
92+
target = objectMap.get(objectId)
93+
if (target is None):
94+
targets = client.describe_assessment_targets(assessmentTargetArns=[objectId])
95+
if (targets is not None):
96+
finalObj = objectMap[objectId] = targets['assessmentTargets'][0]
97+
else:
98+
finalObj = target
99+
elif (objectType == 'finding'):
100+
finding = objectMap.get(objectId)
101+
if (finding is None):
102+
findings = client.describe_findings(findingArns=[objectId])
103+
if (findings is not None):
104+
finalObj = objectMap[objectId] = findings['findings'][0]
105+
else:
106+
finalObj = finding
107+
except Exception as e:
108+
logger.error(e)
109+
raise
110+
return finalObj
111+
112+
# simple utility function to deserialize datetime objects
113+
def json_deserializer(obj):
114+
if isinstance(obj, datetime.datetime):
115+
return obj.strftime('%Y-%m-%dT%H:%M:%SZ')
116+
elif isinstance(obj, date):
117+
return obj.strftime('%Y-%m-%d')
118+
# Let the base class default method raise the TypeError
119+
return json.JSONEncoder.default(self, obj)
120+
121+
122+
def sumo_inspector_handler(event, context):
123+
if ('Records' in event):
124+
for record in event['Records']:
125+
# get actual SNS message
126+
snsObj = record['Sns']
127+
dataObj = {'Timestamp':snsObj['Timestamp'],'Message':snsObj['Message'],'MessageId':snsObj['MessageId']}
128+
msgObj = json.loads(snsObj['Message'])
129+
if (contextLookup):
130+
# do reverse lookup of each of the following items in Message: target, run, template.
131+
if ('template' in msgObj):
132+
lookupItem = lookup(msgObj['template'],'template')
133+
if (lookupItem is not None):
134+
logger.info("Got a template item back")
135+
msgObj['templateLookup']= lookupItem
136+
else:
137+
print("Could not lookup template: %s" % msgObj['template'])
138+
if ('run' in msgObj):
139+
lookupItem = lookup(msgObj['run'],'run')
140+
if (lookupItem is not None):
141+
msgObj['runLookup']= lookupItem
142+
else:
143+
logger.info("Could not lookup run: %s" % msgObj['run'])
144+
if ('target' in msgObj):
145+
lookupItem = lookup(msgObj['target'],'target')
146+
if (lookupItem is not None):
147+
msgObj['targetLookup']= lookupItem
148+
else:
149+
logger.info("Could not lookup target: %s" % msgObj['target'])
150+
if ('finding' in msgObj):
151+
# now query findings
152+
finding = lookup(msgObj['finding'],'finding')
153+
if (finding is not None):
154+
155+
# now query rulesPackage inside the finding
156+
rulesPackage = lookup(finding['serviceAttributes']['rulesPackageArn'],'rulesPackage')
157+
if (rulesPackage is not None):
158+
finding['rulesPackageLookup'] = rulesPackage
159+
else:
160+
logger.info("Cannot lookup rulesPackageArn: %s"% finding['serviceAttributes']['rulesPackageArn'])
161+
msgObj['findingDetails'] = finding
162+
# construct final data object
163+
dataObj = {'Timestamp':snsObj['Timestamp'],'Message':msgObj,'MessageId':snsObj['MessageId']}
164+
165+
# now send this object to Sumo side
166+
rs = sendSumo(json.dumps(dataObj,default=json_deserializer),toCompress=True)
167+
168+
if (rs[0]!=200):
169+
logger.info('Error sending data to sumo with code: %d and message: %s '% (rs[0],rs[1]))
170+
logger.info(json.dumps(dataObj,default=json_deserializer))
171+
else:
172+
logger.info("Sent data to Sumo successfully")
173+
else:
174+
logger.info('Unrecoganized data')
175+

0 commit comments

Comments
 (0)