ail-framework/bin/lib/timeline_engine.py

212 lines
7.8 KiB
Python
Executable file

#!/usr/bin/env python3
# -*-coding:UTF-8 -*
import os
import sys
from uuid import uuid4
sys.path.append(os.environ['AIL_BIN'])
##################################
# Import Project packages
##################################
from lib.ConfigLoader import ConfigLoader
config_loader = ConfigLoader()
r_meta = config_loader.get_db_conn("Kvrocks_Timeline")
config_loader = None
# CORRELATION_TYPES_BY_OBJ = {
# "chat": ["item", "username"], # item ???
# "cookie-name": ["domain"],
# "cryptocurrency": ["domain", "item"],
# "cve": ["domain", "item"],
# "decoded": ["domain", "item"],
# "domain": ["cve", "cookie-name", "cryptocurrency", "decoded", "etag", "favicon", "hhhash", "item", "pgp", "title", "screenshot", "username"],
# "etag": ["domain"],
# "favicon": ["domain", "item"],
# "hhhash": ["domain"],
# "item": ["chat", "cve", "cryptocurrency", "decoded", "domain", "favicon", "pgp", "screenshot", "title", "username"],
# "pgp": ["domain", "item"],
# "screenshot": ["domain", "item"],
# "title": ["domain", "item"],
# "username": ["chat", "domain", "item"],
# }
#
# def get_obj_correl_types(obj_type):
# return CORRELATION_TYPES_BY_OBJ.get(obj_type)
# def sanityze_obj_correl_types(obj_type, correl_types):
# obj_correl_types = get_obj_correl_types(obj_type)
# if correl_types:
# correl_types = set(correl_types).intersection(obj_correl_types)
# if not correl_types:
# correl_types = obj_correl_types
# if not correl_types:
# return []
# return correl_types
class Timeline:
def __init__(self, global_id, name):
self.id = global_id
self.name = name
def _get_block_obj_global_id(self, block):
return r_meta.hget(f'block:{self.id}:{self.name}', block)
def _set_block_obj_global_id(self, block, global_id):
return r_meta.hset(f'block:{self.id}:{self.name}', block, global_id)
def _get_block_timestamp(self, block, position):
return r_meta.zscore(f'line:{self.id}:{self.name}', f'{position}:{block}')
def _get_nearest_bloc_inf(self, timestamp):
inf = r_meta.zrevrangebyscore(f'line:{self.id}:{self.name}', float(timestamp), 0, start=0, num=1, withscores=True)
if inf:
inf, score = inf[0]
if inf.startswith('end'):
inf_key = f'start:{inf[4:]}'
inf_score = r_meta.zscore(f'line:{self.id}:{self.name}', inf_key)
if inf_score == score:
inf = inf_key
return inf
else:
return None
def _get_nearest_bloc_sup(self, timestamp):
sup = r_meta.zrangebyscore(f'line:{self.id}:{self.name}', float(timestamp), '+inf', start=0, num=1, withscores=True)
if sup:
sup, score = sup[0]
if sup.startswith('start'):
sup_key = f'end:{sup[6:]}'
sup_score = r_meta.zscore(f'line:{self.id}:{self.name}', sup_key)
if score == sup_score:
sup = sup_key
return sup
else:
return None
def get_first_obj_id(self):
first = r_meta.zrange(f'line:{self.id}:{self.name}', 0, 0)
if first: # start:block
first = first[0]
if first.startswith('start:'):
first = first[6:]
else:
first = first[4:]
return self._get_block_obj_global_id(first)
def get_last_obj_id(self):
last = r_meta.zrevrange(f'line:{self.id}:{self.name}', 0, 0)
if last: # end:block
last = last[0]
if last.startswith('end:'):
last = last[4:]
else:
last = last[6:]
return self._get_block_obj_global_id(last)
def get_objs_ids(self):
objs = set()
for block in r_meta.zrange(f'line:{self.id}:{self.name}', 0, -1):
if block:
if block.startswith('start:'):
objs.add(self._get_block_obj_global_id(block[6:]))
return objs
# def get_objs_ids(self):
# objs = {}
# last_obj_id = None
# for block, timestamp in r_meta.zrange(f'line:{self.id}:{self.name}', 0, -1, withscores=True):
# if block:
# if block.startswith('start:'):
# last_obj_id = self._get_block_obj_global_id(block[6:])
# objs[last_obj_id] = {'first_seen': timestamp}
# else:
# objs[last_obj_id]['last_seen'] = timestamp
# return objs
def _update_bloc(self, block, position, timestamp):
r_meta.zadd(f'line:{self.id}:{self.name}', {f'{position}:{block}': timestamp})
def _add_bloc(self, obj_global_id, timestamp, end=None):
if end:
timestamp_end = end
else:
timestamp_end = timestamp
new_bloc = str(uuid4())
r_meta.zadd(f'line:{self.id}:{self.name}', {f'start:{new_bloc}': timestamp, f'end:{new_bloc}': timestamp_end})
self._set_block_obj_global_id(new_bloc, obj_global_id)
return new_bloc
def add_timestamp(self, timestamp, obj_global_id):
inf = self._get_nearest_bloc_inf(timestamp)
sup = self._get_nearest_bloc_sup(timestamp)
if not inf and not sup:
# create new bloc
new_bloc = self._add_bloc(obj_global_id, timestamp)
return new_bloc
# timestamp < first_seen
elif not inf:
sup_pos, sup_id = sup.split(':')
sup_obj = self._get_block_obj_global_id(sup_id)
if sup_obj == obj_global_id:
self._update_bloc(sup_id, 'start', timestamp)
# create new bloc
else:
new_bloc = self._add_bloc(obj_global_id, timestamp)
return new_bloc
# timestamp > first_seen
elif not sup:
inf_pos, inf_id = inf.split(':')
inf_obj = self._get_block_obj_global_id(inf_id)
if inf_obj == obj_global_id:
self._update_bloc(inf_id, 'end', timestamp)
# create new bloc
else:
new_bloc = self._add_bloc(obj_global_id, timestamp)
return new_bloc
else:
inf_pos, inf_id = inf.split(':')
sup_pos, sup_id = sup.split(':')
inf_obj = self._get_block_obj_global_id(inf_id)
if inf_id == sup_id:
# reduce bloc + create two new bloc
if obj_global_id != inf_obj:
# get end timestamp
sup_timestamp = self._get_block_timestamp(sup_id, 'end')
# reduce original bloc
self._update_bloc(inf_id, 'end', timestamp - 1)
# Insert new bloc
new_bloc = self._add_bloc(obj_global_id, timestamp)
# Recreate end of the first bloc by a new bloc
self._add_bloc(inf_obj, timestamp + 1, end=sup_timestamp)
return new_bloc
# timestamp in existing bloc
else:
return inf_id
# different blocs: expend sup/inf bloc or create a new bloc if
elif inf_pos == 'end' and sup_pos == 'start':
# Extend inf bloc
if obj_global_id == inf_obj:
self._update_bloc(inf_id, 'end', timestamp)
return inf_id
sup_obj = self._get_block_obj_global_id(sup_id)
# Extend sup bloc
if obj_global_id == sup_obj:
self._update_bloc(sup_id, 'start', timestamp)
return sup_id
# create new bloc
new_bloc = self._add_bloc(obj_global_id, timestamp)
return new_bloc
# inf_pos == 'start' and sup_pos == 'end'
# else raise error ???