def get_node_by_id(g: Graph, id: str) -> Optional[str]:
"""Queries Database for any type--ID
:param g: OGM Graph ref
:param id: STIX Node id_
"""
ret = ''
try:
id_type = id.split('--')[0]
_uuid = id.split('--')[1]
# if id_type.has
class_name = id_type.replace('-', '')
if class_name in g.registry:
klass = g.registry[class_name]
queried_proper = g.query(klass).filter(
klass.id_.endswith(_uuid) & (klass.id_.startswith(id_type))).one()
ret = queried_proper
except Exception as e:
print('**** query error:')
print(e)
traceback.print_exc()
print('\n')
if len(ret.id_):
return ret
else:
return None
def insert_bundle(g: Graph, stix_objs: List[Dict]):
"""Insert bundle into OrientDB graph. Tested on v3.0.31
:param g: OGM Graph ref
:param id: STIX Node id_
"""
objects = format_stix_to_db(stix_objs)
o: Dict
for o in tqdm(objects, 'Objects'):
print("\nObject:")
print(o)
o_type = o['type']
if '-' in o_type:
o_type = o_type.replace('-', '')
o_id = o['id_']
id_type = o_id.strip().split('--')[0]
_uuid = o_id.strip().split('--')[1]
class_name = id_type.replace('-', '')
if class_name in g.registry:
klass = g.registry[class_name]
#print(klass)
#TODO: Perhaps hit fixed get_node_by_id to to check for object existing.
try:
if len(
g.query(klass).filter(
klass.id_.endswith(_uuid)
& (klass.id_.startswith(id_type)))):
# if len(g.gremlin(f"g.V().has('id_', '{o_id}')")) or len(g.gremlin(f"g.E().has('id_', '{o_id}')")):
print(f"Object with id: {o_id} found, moving to next object")
continue
else:
if o_type == 'relationship':
r_type = o['relationship_type'].replace('-', '_')
src_id_type = o['source_ref'].strip().split('--')[0]
s_id = o['source_ref'].strip().split('--')[1]
tar_id_type = o['target_ref'].strip().split('--')[0]
t_id = o['target_ref'].strip().split('--')[1]
if r_type in g.registry:
temp = g.registry[r_type]
#Checks all props of a relationship (src/target_refs and id_ to determine of we already have this relationship in the database)
#TODO: what if we want to update?? Flag??? If keep and update flag is true, we replace anything currently in the database with the new file?
if (len(
g.query(temp).filter(
klass.id_.endswith(_uuid)
& (klass.id_.startswith(id_type))))):
if (len(
g.query(temp).filter(
klass.source_ref.startswith(src_id_type)
& (klass.source_ref.endswith(s_id))))):
if (len(
g.query(temp).filter(
klass.target_ref.startswith(
tar_id_type)
& (klass.target_ref.endswith(t_id))))):
print(
f"Relationship with id: {o_id} found, moving on to next object"
)
continue
except Exception as e:
print('**** gremlin error:')
print(e)
traceback.print_exc()
print('\n')
if o['type'] != 'relationship':
if o_type in g.registry:
klass = g.registry[o_type]
else:
print('!' * 40)
print(f'Stix Object of type {o["type"]} not in g.registry.')
try:
klass = class_from_json(o.copy(), g)
print(f'Created DB class {klass}!')
except Exception as e:
print('Exception creating DB class')
print(e)
traceback.print_exc()
print('\n')
print('!' * 40, '\n')
try:
klass = g.registry[o_type]
o = check_mandatory_props(klass, o)
res = klass.objects.create(**o)
except Exception as e:
print('****error inserting:')
print(o)
print(e)
print('\n')
else:
src = get_node_by_id(g, o['source_ref'])
target = get_node_by_id(g, o['target_ref'])
if not src or not target:
print(
'\n\n\n!!!!!!!!!!!!!! ERROR: Missing node when creating edge'
)
print(f'\tsrc: {src}')
print(f'\target: {target}')
print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n\n\n')
continue
r_type = o['relationship_type'].replace('-', '_')
if r_type not in g.registry:
print('!' * 40)
print(f'Stix Relationship of type {r_type} not in g.registry.')
try:
klass = class_from_json(o.copy(), g)
print(f'Created DB relationship class {klass}!')
print('!' * 40)
except Exception as e:
print(f'Exception creating DB relationship class for {o}')
print(e)
traceback.print_exc()
print('!' * 40, '\n')
klass = g.registry[r_type]
res = g.create_edge(klass, src, target, **o)
