Skip to content

Commit

Permalink
fix is_node_available
Browse files Browse the repository at this point in the history
  • Loading branch information
@rpiazza
rpiazza committed Jan 14, 2025
1 parent ca03792 commit 4f934ee
Showing 1 changed file with 36 additions and 29 deletions.
65 changes: 36 additions & 29 deletions stix2/datastore/neo4j/neo4j.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,8 @@ def _insert_sdo_sco_smo(self, obj, type_name):
self.sgraph.create(node)
# check to see if the addition of this node makes it possible to create a relationship
for rel in self.relationships_to_recheck:
self._insert_sro(obj, True)
self._insert_embedded_relationships(obj)
self._insert_sro(rel, True)
self._insert_embedded_relationships(obj, obj["id"])
self._insert_external_references(external_references, node)
self._insert_extensions(extension_relationships, node)
self._insert_sub_objects(self.sub_object_relationships, node)
Expand All @@ -190,25 +190,27 @@ def _insert_sub_object(self, sub_prop, sub_obj, parent_node):
node_contents = dict()
node_contents["type"] = sub_prop
for key, value in sub_obj.items():
if isinstance(value, list):
if isinstance(value[0], str):
node_contents[key] = ",".join(value)
elif isinstance(value[0], dict):
for v in value:
self.sub_object_relationships.append((key, v))
elif key == "hashes":
node_contents[key] = hash_dict_as_string(value)
elif not isinstance(value, dict):
node_contents[key] = value
else:
self.sub_object_relationships.append((key, value))
if not key.endswith("ref") and not key.endswith("refs"):
if isinstance(value, list):
if isinstance(value[0], str):
node_contents[key] = ",".join(value)
elif isinstance(value[0], dict):
for v in value:
self.sub_object_relationships.append((key, v))
elif key == "hashes":
node_contents[key] = hash_dict_as_string(value)
elif not isinstance(value, dict):
node_contents[key] = value
else:
self.sub_object_relationships.append((key, value))
node = Node(sub_prop,
name=sub_prop + "_" + self.next_id(),
# bundlesource=self.bundlename,
**node_contents)
self.sgraph.create(node)
relationship = Relationship(parent_node, sub_prop, node)
self.sgraph.create(relationship)
self._insert_embedded_relationships(sub_obj, parent_node["id"])

def _insert_sub_objects(self, sub_objects, parent_node):
for sub in sub_objects:
Expand Down Expand Up @@ -238,26 +240,29 @@ def _insert_extensions(self, extensions, parent_node):
type_name = ext.__class__.__name__
node_contents["type"] = type_name
for key, value in ext.items():
if isinstance(value, list):
if isinstance(value[0], str):
node_contents[key] = ",".join(value)
if not key.endswith("ref") and not key.endswith("refs"):
if isinstance(value, list):
if isinstance(value[0], str):
node_contents[key] = ",".join(value)
else:
for v in value:
self.sub_object_relationships.append((key, v))
elif key == "hashes":
node_contents[key] = hash_dict_as_string(value)
else:
for v in value:
self.sub_object_relationships.append((key, v))
elif key == "hashes":
node_contents[key] = hash_dict_as_string(value)
else:
node_contents[key] = value
node_contents[key] = value
node = Node(type_name,
name=type_name + "_" + self.next_id(),
# bundlesource=self.bundlename,
**node_contents)
relationship = Relationship(parent_node, type_name, node)
self.sgraph.create(relationship)
self._insert_embedded_relationships(ext, parent_node["id"])

def _is_node_available(self, id,):
cypher_string = f'OPTIONAL MATCH (a) WHERE a.id="{str(id)}" RETURN a'
return self.sgraph.run(cypher_string)
cypher_string = f'OPTIONAL MATCH (a) WHERE a.id="{str(id)}" UNWIND [a] AS list_rows RETURN list_rows'
cursor = self.sgraph.run(cypher_string).data()
return cursor[0]["list_rows"]

def _insert_sro(self, obj, recheck=False):
reltype = str(obj['relationship_type'])
Expand All @@ -270,13 +275,14 @@ def _insert_sro(self, obj, recheck=False):
if self._is_node_available(obj["source_ref"]) and self._is_node_available(obj["target_ref"]):
cypher_string = f'MATCH (a),(b) WHERE a.id="{str(obj["source_ref"])}" AND b.id="{str(obj["target_ref"])}" CREATE (a)-[r:{reltype}]->(b) RETURN a,b'
self.sgraph.run(cypher_string)
print(f'Created {str(obj["source_ref"])} {reltype} {obj["target_ref"]}')
if recheck:
remove_sro_from_list(obj, self.relationships_to_recheck)
else:
if not recheck:
self.relationships_to_recheck.append(obj)

def _insert_embedded_relationships(self, obj, recheck=False):
def _insert_embedded_relationships(self, obj, id, recheck=False):
for k in obj.keys():
k_tokens = k.split("_")
# find refs, but ignore external_references since they aren't objects
Expand All @@ -291,15 +297,16 @@ def _insert_embedded_relationships(self, obj, recheck=False):
for ref in ref_list:
if self._is_node_available(ref):
# The "b to a" relationship is reversed in this cypher query to ensure the correct relationship direction in the graph
cypher_string = f'MATCH (a),(b) WHERE a.id="{str(ref)}" AND b.id="{str(obj["id"])}" CREATE (b)-[r:{rel_type}]->(a) RETURN a,b'
cypher_string = f'MATCH (a),(b) WHERE a.id="{str(ref)}" AND b.id="{str(id)}" CREATE (b)-[r:{k}]->(a) RETURN a,b'
self.sgraph.run(cypher_string)
print(f'Created * {str(id)} {k} {str(ref)}')
if recheck:
remove_sro_from_list(obj, self.relationships_to_recheck)
else:
if not recheck:
embedded_relationship = {"source_ref": obj["id"],
embedded_relationship = {"source_ref": id,
"target_ref": ref,
"relationship_type": rel_type}
"relationship_type": k}
self.relationships_to_recheck.append(embedded_relationship)


0 comments on commit 4f934ee

Please sign in to comment.