def expand_forward_edges_in_scope(dgraph_client: DgraphClient, node: NodeView, lens: str) -> None:
for edge_name, edge_type in node._get_forward_edge_types().items():
if isinstance(edge_type, list):
inner_edge_type = edge_type[0]
else:
inner_edge_type = edge_type
edges_in_lens = edge_in_lens(dgraph_client, node.uid, edge_name, lens)
for edge in edges_in_lens:
for neighbors in edge.values():
if not isinstance(neighbors, list):
neighbors = [neighbors]
for neighbor in neighbors:
if neighbor.get('~scope'):
neighbor.pop('~scope')
node_edge = getattr(node, edge_name)
try:
neighbor_view = inner_edge_type(dgraph_client, node_key=neighbor['node_key'], uid=neighbor['uid'])
except Exception as e:
print(f'neighbor_view failed with: {e}')
continue
print(neighbor_view, neighbor_view.uid, neighbor_view.node_key)
if isinstance(node_edge, list):
node_edge.append(neighbor_view)
else:
node_edge = neighbor_view
setattr(node, edge_name, node_edge)
def expand_reverse_edges_in_scope(dgraph_client: DgraphClient, node: NodeView, lens: str) -> None:
for edge_name, (edge_type, forward_name) in node._get_reverse_edge_types().items():
if isinstance(edge_type, list):
inner_edge_type = edge_type[0]
else:
inner_edge_type = edge_type
edges_in_lens = edge_in_lens(dgraph_client, node.uid, edge_name, lens)
for edge in edges_in_lens:
for neighbors in edge.values():
if not isinstance(neighbors, list):
neighbors = [neighbors]
for neighbor in neighbors:
if neighbor.get('~scope'):
neighbor.pop('~scope')
neighbor_view = inner_edge_type(
dgraph_client,
node_key=neighbor['node_key'],
uid=neighbor['uid'],
)
node_edge = getattr(node, forward_name)
if isinstance(node_edge, list):
node_edge.append(neighbor_view)
else:
node_edge = neighbor_view
setattr(node, forward_name, node_edge)
def expand_forward_edges_in_scope(dgraph_client: DgraphClient, node: NodeView,
lens: str) -> None:
for edge_name, edge_type in node._get_forward_edge_types().items():
if isinstance(edge_type, list):
inner_edge_type = edge_type[0]
else:
inner_edge_type = edge_type
edges_in_lens = edge_in_lens(dgraph_client, node.uid, edge_name, lens)
for edge in edges_in_lens:
for neighbors in edge.values():
if not isinstance(neighbors, list):
neighbors = [neighbors]
for neighbor in neighbors:
if neighbor.get("~scope"):
neighbor.pop("~scope")
node_edge = getattr(node, edge_name)
try:
neighbor_view = inner_edge_type(
dgraph_client,
node_key=neighbor["node_key"],
uid=neighbor["uid"],
)
except Exception as e:
LOGGER.error(f"neighbor_view failed with: {e}")
continue
LOGGER.debug(neighbor_view, neighbor_view.uid,
neighbor_view.node_key)
if isinstance(node_edge, list):
node_edge.append(neighbor_view)
else:
node_edge = neighbor_view
setattr(node, edge_name, node_edge)
def lens_to_dict(dgraph_client: DgraphClient, lens_name: str) -> List[Dict[str, Any]]:
current_graph = get_lens_scope(dgraph_client, lens_name)
print(f'Getting lens as dict {current_graph}')
if not current_graph or not current_graph.get('scope'):
return []
nodes = []
for graph in current_graph['scope']:
try:
nodes.append(NodeView.from_dict(dgraph_client, graph))
except Exception as e:
print('Failed to get NodeView from dict', e)
if current_graph.get('scope'):
current_graph.pop('scope')
concrete_nodes = [n.node for n in nodes if not isinstance(n.node, DynamicNodeView)]
dynamic_nodes = [n.node for n in nodes if isinstance(n.node, DynamicNodeView)]
expanded_dynamic_nodes = []
for dynamic_node in dynamic_nodes:
expanded = expand_dynamic_node(dynamic_node)
expanded_dynamic_nodes.append(expanded)
expand_concrete_nodes(
dgraph_client,
lens_name,
concrete_nodes
)
results = [{
"node": current_graph,
"edges": []
}]
lens_risks = get_lens_risks(dgraph_client, lens_name)
for node in lens_risks:
edges = []
risks = node.get("risks", [])
if not risks:
print(f"WARN: Node in engagement graph has no connected risks {node}")
for risk in risks:
try:
risk['node_key'] = node['node_key'] + risk['analyzer_name']
edge = {
"from": node["node_key"],
"edge_name": "risks",
"to": risk['node_key']
}
edges.append(edge)
except Exception as e:
print(f'risk edge failed: {risk} {e}')
results.append({
"node": node,
"edges": edges,
})
results.extend([n.to_dict() for n in concrete_nodes])
results.extend(expanded_dynamic_nodes)
return results
def exec_analyzers(
dg_client,
file: str,
msg_id: str,
nodes: List[NodeView],
analyzers: Dict[str, Analyzer],
sender: Any,
):
if not analyzers:
LOGGER.warning("Received empty dict of analyzers")
return
if not nodes:
LOGGER.warning("Received empty array of nodes")
result_name_to_analyzer = {}
query_str = ""
for node in nodes:
querymap = defaultdict(list)
for an_name, analyzer in analyzers.items():
if check_caches(file, msg_id, node.node_key, an_name):
continue
analyzer = analyzer # type: Analyzer
queries = analyzer.get_queries()
if isinstance(queries, list) or isinstance(queries, tuple):
querymap[an_name].extend(queries)
else:
querymap[an_name].append(queries)
for an_name, queries in querymap.items():
analyzer = analyzers[an_name]
for i, query in enumerate(queries):
analyzer_query_types = get_analyzer_view_types(query)
if node.node.get_node_type() + "View" not in [
n.__name__ for n in analyzer_query_types
]:
continue
r = str(random.randint(10, 100))
result_name = f"{an_name}u{int(node.uid, 16)}i{i}r{r}".strip().lower()
result_name_to_analyzer[result_name] = (
an_name,
analyzer,
query.view_type,
)
query_str += "\n"
query_str += generate_query(
query_name=result_name,
binding_modifier=result_name,
root=query,
contains_node_key=node.node_key,
)
if not query_str:
LOGGER.warning("No nodes to query")
return
txn = dg_client.txn(read_only=True)
try:
response = json.loads(txn.query(query_str).json)
finally:
txn.discard()
analyzer_to_results = defaultdict(list)
for result_name, results in response.items():
for result in results:
analyzer_meta = result_name_to_analyzer[
result_name
] # type: Tuple[str, Analyzer, Type[Viewable]]
an_name, analyzer, view_type = (
analyzer_meta[0],
analyzer_meta[1],
analyzer_meta[2],
)
result_graph = view_type.from_dict(dg_client, result)
response_ty = inspect.getfullargspec(analyzer.on_response).annotations.get(
"response"
)
if response_ty == NodeView:
LOGGER.warning("Analyzer on_response is expecting a NodeView")
result_graph = NodeView.from_view(result_graph)
analyzer_to_results[an_name].append(result_graph)
with ThreadPoolExecutor(max_workers=6) as executor:
for an_name, result_graphs in analyzer_to_results.items():
analyzer = analyzers[an_name]
executor.submit(handle_result_graphs, analyzer, result_graphs, sender)
executor.shutdown(wait=True)