class Display_Intelligence():
def __init__(self):
self.link = Database("Intelligence")
super().__init__()
def display_ioc(self, type):
result = {}
result["data"] = []
select_result = self.link.select(type,
query={},
sort="disclosure_time")
if select_result:
for alert in select_result:
alert.pop("_id")
result["data"].append(alert)
self.link.close()
result["data"] = list_dict_duplicate_removal(result["data"])
return result
def ioc_count(self):
result = {}
count = 0
for ioc_type in ioc:
result[ioc_type] = self.link.count(ioc_type)
count += count + result[ioc_type]
result["count"] = count
self.link.close()
# result["data"] = list_dict_duplicate_removal(result["data"])
return result
def deal_eve_content(filecontent):
link_mongo = Database(database_name="Semitic")
link_ioc = IoC()
eve_json = []
json_result = {}
if filecontent:
for i in filecontent:
eve_json.append(json.loads(i))
link_mongo.insert("eve", eve_json)
json_result = classify_eve(eve_json)
if json_result:
for type_json in json_result:
link_mongo.insert(type_json, json_result[type_json])
try:
if type_json == "tcp":
link_mongo.insert(
"alert_ioc", link_ioc.deal_tcp(json_result[type_json]))
elif type_json == "tls":
link_mongo.insert(
"alert_ioc", link_ioc.deal_tls(json_result[type_json]))
elif type_json == "http":
link_mongo.insert(
"alert_ioc", link_ioc.deal_http(json_result[type_json]))
elif type_json == "dns":
link_mongo.insert(
"alert_ioc", link_ioc.deal_dns(json_result[type_json]))
pass
except:
pass
link_mongo.close()
link_ioc.link_mongo.close()
class TestDatabase(unittest.TestCase):
def setup_method(self, test_method):
"""Create/Wipe a table called test_data where tests will be run"""
self.connection = Database(database=TEST_DATABASE)
self.connection.query('DROP TABLE IF EXISTS test_data')
self.connection.query('CREATE TABLE test_data (variable INTEGER)')
def teardown_method(self, test_method):
"""Drop the test_data table used to perform tests"""
self.connection.query('DROP TABLE test_data')
self.connection.close()
def test_commit(self):
"""Ensure that the database commit method retains
the data after reconnection.
"""
# insert test data and commit
self.connection.query('INSERT INTO test_data (variable) VALUES (1)')
self.connection.commit()
self.connection.close()
# reopen connection and query to ensure data was retained
self.connection = Database(database=TEST_DATABASE)
result = self.connection.query('SELECT * FROM test_data')
self.assertEqual(result, [(1, )])
def test_rollback(self):
"""Ensure that the database rollback method removes changes to the data
"""
# insert test data and rollback
self.connection.query('INSERT INTO test_data (variable) VALUES (1)')
self.connection.rollback()
self.connection.close()
# reopen connection and query to ensure data was removed
self.connection = Database(database=TEST_DATABASE)
result = self.connection.query('SELECT * FROM test_data')
self.assertEqual(result, [])
def test_close(self):
"""Ensure that the close method closes the database connection"""
with self.assertRaises(InterfaceError):
self.connection.close()
self.connection.query('SELECT * FROM test_data')
def test_with_commit(self):
"""Ensure that an error free with statement commits when exiting
Performs an error free with block and ensures the data remains
"""
with Database(database=TEST_DATABASE) as db:
db.query('INSERT INTO test_data (variable) VALUES (1)')
result = self.connection.query('SELECT * FROM test_data')
self.assertEqual(result, [(1, )])
def test_with_rollback(self):
"""Ensure that an error occurring in a database with context forces the database to rollback
Raises an error in the with context and ensures that modifications were removed
"""
try:
with Database(database=TEST_DATABASE) as db:
db.query('INSERT INTO test_data (variable) VALUES (1)')
raise Exception
except:
pass
result = self.connection.query('SELECT * FROM test_data')
self.assertEqual(result, [])
def test_limit(self):
"""Ensure that the limit keyword of the query method works as expected
Performs several queries with varying limits and compares the results to the expected results
"""
self.connection.query(
'INSERT INTO test_data (variable) VALUES (1), (2), (3), (4), (5)')
result = self.connection.query('SELECT * FROM test_data', limit=1)
self.assertEqual(result, [(1, )])
result = self.connection.query('SELECT * FROM test_data', limit=3)
self.assertEqual(result, [(1, ), (2, ), (3, )])
result = self.connection.query('SELECT * FROM test_data')
self.assertEqual(result, [(1, ), (2, ), (3, ), (4, ), (5, )])
def test_exists(self):
"""Ensure that the exists method accurately reports the existence of rows in the database"""
self.connection.query(
'INSERT INTO test_data (variable) VALUES (1), (2), (3), (4), (5)')
self.assertTrue(self.connection.exists('test_data'))
self.assertTrue(self.connection.exists('test_data', variable=3))
self.assertFalse(self.connection.exists('test_data', variable=6))
def test_bool(self):
"""Ensure that the database returns the correct boolean value based on whether it is open or closed"""
self.assertTrue(self.connection)
self.connection.close()
self.assertFalse(self.connection)
def test_del(self):
"""Ensure that the __del__ magic method works in the same way the close works"""
with self.assertRaises(InterfaceError):
self.connection.__del__()
self.connection.query('SELECT * FROM test_data')
class Display_Semitic():
def __init__(self):
self.link = Database("Semitic")
super().__init__()
def display_alert_rule(self):
result = {}
result["data"] = []
select_result = self.link.select("alert", query={})
if select_result:
for alert in select_result:
alert.pop("_id")
result["data"].append(alert)
result["data"] = list_dict_duplicate_removal(result["data"])
self.link.close()
return result
def display_alert_ioc(self):
result = {}
result["data"] = []
select_result = self.link.select("alert_ioc", query={})
if select_result:
for alert in select_result:
alert.pop("_id")
result["data"].append(alert)
self.link.close()
result["data"] = list_dict_duplicate_removal(result["data"])
return result
def display_service(self):
result = {}
result["data"] = []
select_result = self.link.select("service", query={})
if select_result:
for service in select_result:
service.pop("_id")
result["data"].append(service)
self.link.close()
result["data"] = list_dict_duplicate_removal(result["data"])
return result
def display_proto(self, type):
result = {}
result["data"] = []
select_result = self.link.select(type, query={})
for proto in select_result:
proto.pop("_id")
if type == "http":
if proto["uri"] == "/api/upload_eve":
continue
result["data"].append(proto)
self.link.close
result["data"] = list_dict_duplicate_removal(result["data"])
return result
def alert_count(self):
result = {}
result["rule_count"] = self.link.count("alert")
result["ioc_count"] = self.link.count("alert_ioc")
result["count"] = result["rule_count"] + result["ioc_count"]
self.link.close()
return result
def proto_count(self):
result = {}
count = 0
for proto_type in proto:
result[proto_type] = self.link.count(proto_type)
count += self.link.count(proto_type)
result["count"] = result["tcp"] + result['udp']
self.link.close()
return result
def eve_count(self):
result = {"count": self.link.count("eve")}
return result
class Search():
def __init__(self, data):
self.link = Database("Intelligence")
self.search_type = list(data.keys())[0]
self.search_value = data[self.search_type]
self.result = {}
self.result["data"] = []
super().__init__()
def search_local(self):
result_local = {}
result_local["from"] = "local"
result_local["type"] = self.search_type
result_local["value"] = self.search_value
try:
local = self.link.select(
self.search_type, query={self.search_type:
self.search_value})[0]
result_local["result"] = "success"
result_local["tags"] = local["tags"]
# result_local["time"] = local["disclosure_time"]
except:
result_local["result"] = "fail"
self.link.close()
return (result_local)
def search_venuseye(self):
check_url = {
"ip": "https://www.venuseye.com.cn/ve/ip/ioc",
"email": "https://www.venuseye.com.cn/ve/email/ioc",
"domain": "https://www.venuseye.com.cn/ve/domain/ioc",
"hash": {
"sha256": "https://www.venuseye.com.cn/ve/sample/sha256",
"sha1": "https://www.venuseye.com.cn/ve/sample/sha1",
"md5": "https://www.venuseye.com.cn/ve/sample/md5"
}
}
venuseye_reslut = {}
venuseye_reslut["from"] = "venuseye"
venuseye_reslut["type"] = self.search_type
venuseye_reslut["value"] = self.search_value
try:
if self.search_type == "hash":
if len(self.search_value) == 32:
venuseye = json.loads(
requests.post(check_url["hash"]["md5"], {
"target": self.search_value
}).content)
elif len(self.search_value) == 40:
venuseye = json.loads(
requests.post(check_url["hash"]["sha1"], {
"target": self.search_value
}).content)
elif len(self.search_value) == 64:
venuseye = json.loads(
requests.post(check_url["hash"]["sha256"], {
"target": self.search_value
}).content)
venuseye_reslut["tags"] = venuseye["data"]["tags"]
else:
venuseye = json.loads(
requests.post(check_url[self.search_type], {
"target": self.search_value
}).content)
venuseye_reslut["result"] = "success"
venuseye_reslut["tags"] = []
for ioc in venuseye["data"]["ioc"]:
for tag in ioc["categories"]:
venuseye_reslut["tags"].append(tag)
venuseye_reslut["tags"] = [
x for x in set(venuseye_reslut["tags"])
]
except:
venuseye_reslut["result"] = "fail"
return venuseye_reslut
def get_reslut(self):
self.result["data"].append(self.search_local())
self.result["data"].append(self.search_venuseye())
return self.result