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)