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)