def test_async_dict_search(self):
""" test adding threat info using dict objects instead of instances """
artifact_type = "net.ip"
artifact_value = "192.168.1.1"
context = SearchContext({
"type": artifact_type,
"value": artifact_value
})
hits = search_artifacts(artifact_type,
artifact_value,
True,
context=context)
self.assertTrue(len(hits) == 0)
context.save()
# try again, still should find nothing
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 0)
context = SearchContext.load(context.id)
self.assertTrue(context.pending_searches)
TestDictArtifactSearch.add_result("Test result posted", artifact_type,
artifact_value)
TestDictArtifactSearch.search_complete_flag().wait(2)
# now should find something, other search engine should've added something
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 1)
context = SearchContext.load(context.id)
self.assertFalse(context.pending_searches)
self.assertEqual(len(hits[0].props.all()), 1)
def test_async_search(self):
""" test search operation flagged within search context """
artifact_type = "net.ip"
artifact_value = "192.168.1.1"
context = SearchContext({
"type": artifact_type,
"value": artifact_value
})
hits = search_artifacts(artifact_type,
artifact_value,
True,
context=context)
self.assertTrue(len(hits) == 0)
context.save()
# try again, still should find nothing
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 0)
context = SearchContext.load(context.id)
self.assertTrue(context.pending_searches)
TestArtifactSearch.add_result("Test result posted", artifact_type,
artifact_value)
# now should find something, other search engine should've added something
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 1)
context = SearchContext.load(context.id)
self.assertTrue(context.pending_searches)
TestArtifactSearch.search_complete_flag().wait(2)
context = SearchContext.load(context.id)
self.assertFalse(context.pending_searches)
def test_bad_searcher(self):
""" should raise error """
artifact_type = "net.ip"
artifact_value = "192.168.1.1"
with self.assertRaises(NotImplementedError):
search_artifacts(artifact_type, artifact_value, True)
def test_filter_search(self):
"""
Tests how searches match upon the correct threats only
"""
threat1 = Threat.objects.create(name="test_threat")
Artifact.objects.create(type="net.ip",
value="192.168.1.1",
threat=threat1)
Artifact.objects.create(type="file.name",
value="badprogram.exe",
threat=threat1)
threat2 = Threat.objects.create(name="another_threat")
Artifact.objects.create(type="net.ip",
value="192.168.1.1",
threat=threat2)
Artifact.objects.create(type="file.path",
value="C:\\Program Files\\Bad",
threat=threat2)
results = search_artifacts("file.name", "badprogram")
self.assertEqual(len(results), 1)
self.assertTrue(threat1 in results)
results = search_artifacts("file.path", "C:\\")
self.assertEqual(len(results), 1)
self.assertTrue(threat2 in results)
results = search_artifacts("net.ip", "192.168.1")
self.assertEqual(len(results), 2)
self.assertTrue(threat1 in results)
self.assertTrue(threat2 in results)
def post(self, request): # pylint: disable=no-self-use
""" receive artifact information, and locate all threats that match it """
response = None
def check_required_value(
arg_name,
dictionary=request.data,
description_format='Required parameter "{arg_name}" was not provided'
):
"""
Check that a required value has been provided by the client
Arguments:
arg_name - name of the value that should be present
description_format - format specification of descriptive message
"""
if arg_name not in dictionary:
raise ValueError("{} - {}".format(
description_format.format(arg_name=arg_name),
json.dumps(dictionary),
))
try:
context_data = request.data
if request.content_type.startswith('multipart/form-data'):
check_required_value('artifact')
context_data = json.loads(request.data['artifact'])
check_required_value('type', context_data)
if context_data['type'] == "file.content":
check_required_value('file', request.FILES)
context_data['value'] = request.FILES['file']
check_required_value('value', context_data)
else:
check_required_value('type')
check_required_value('value')
# Generate a search context
context = SearchContext(context_data)
matching_threats = search_artifacts(context_data['type'],
context_data['value'],
True,
context=context)
response = process_threat_results(matching_threats, context)
context.save()
except ValueError as ve:
response = Response({'error': '{}'.format(str(ve))},
status.HTTP_400_BAD_REQUEST)
# Keep broad exception handler, as this is an entry point to the service interface
except Exception as ex: # pylint: disable=broad-except
response = Response({'error': '{}'.format(str(ex))},
status.HTTP_500_INTERNAL_SERVER_ERROR)
return response
def test_sync_search_other(self):
""" if no context is provided, then search will be synchronous """
artifact_type = "net.ip"
artifact_value = "192.168.1.1"
TestArtifactSearch.add_result("Test result posted", artifact_type,
artifact_value)
hits = search_artifacts(artifact_type, artifact_value, True)
# no context, so search should be synchronous, and we should get a hit
self.assertTrue(len(hits) == 1)
def test_simple_search(self):
"""
Test the search with property results for each kind of property
"""
threat = Threat.objects.create(name="test_threat")
self.assertEqual(str(threat), threat.name)
ip_artifact = Artifact.objects.create(type="net.ip",
value="192.168.1.1",
threat=threat)
self.assertEqual(str(ip_artifact),
"{} - 192.168.1.1 (net.ip)".format(ip_artifact.id))
filename_artifact = Artifact.objects.create(type="file.name",
value="badprogram.exe",
threat=threat)
prop1 = Property.objects.create(name="P1",
type="string",
value="test value test value",
threat=threat)
self.assertEqual(str(prop1), "P1 (string)")
prop2 = Property.objects.create(name="P2",
type="number",
value="52",
threat=threat)
prop3 = Property.objects.create(
name="P3",
type="uri",
value=
"http://www.flightview.com/TravelTools/FlightTrackerQueryResults.asp",
threat=threat)
prop4 = Property.objects.create(name="P4",
type="ip",
value="192.168.1.1",
threat=threat)
prop5 = Property.objects.create(
name="P5",
type="latlng",
value='{"lat": 42.398506, "lng": -71.139872}',
threat=threat)
results = search_artifacts("file.name", "badprogram")
self.assertEqual(len(results), 1)
self.assertEqual(results[0].id, threat.id)
artifacts = results[0].artifacts.all()
self.assertTrue(ip_artifact in artifacts)
self.assertTrue(filename_artifact in artifacts)
properties = results[0].props.all()
self.assertTrue(prop1 in properties)
self.assertTrue(prop2 in properties)
self.assertTrue(prop3 in properties)
self.assertTrue(prop4 in properties)
self.assertTrue(prop5 in properties)
def test_results_already_stored(self):
""" Test duplicate entries from same searcher """
artifact_type = "net.ip"
artifact_value = "192.168.1.1"
context = SearchContext({
"type": artifact_type,
"value": artifact_value
})
hits = search_artifacts(artifact_type,
artifact_value,
True,
context=context)
self.assertTrue(len(hits) == 0)
context.save()
# try again, still should find nothing
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 0)
context = SearchContext.load(context.id)
self.assertTrue(context.pending_searches)
TestArtifactSearch.add_result("Test result posted", artifact_type,
artifact_value)
# now should find something, other search engine should've added something
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 1)
# add the same result again from the same searcher - should still only get 1 result due to duplicate detection
TestArtifactSearch.add_result("Test result posted", artifact_type,
artifact_value)
TestArtifactSearch.search_complete_flag().wait(2)
# now should find something, other search engine should've added something
hits = search_artifacts(artifact_type, artifact_value)
self.assertTrue(len(hits) == 1)
def get(self, request, request_id):
""" Fetch results corresponding to given request id """
request_id = request_id.rstrip('/')
context = SearchContext.load(request_id)
if context is None:
# We return "no content" because we have none
# 404 will cause the custom threat service caller to stop processing
return Response({"error": "'{}' was not found".format(request_id)},
status.HTTP_204_NO_CONTENT)
try:
matching_threats = search_artifacts(context.type, context.value)
response = process_threat_results(matching_threats, context)
# Keep broad exception handler, as this is an entry point to the service interface
except Exception as ex: # pylint: disable=broad-except
response = Response({'error': '{}'.format(str(ex))},
status.HTTP_500_INTERNAL_SERVER_ERROR)
return response