def query_batch( client: GraphClient, batch_size: int, ttl_cutoff_ms: int, last_uid: Optional[str] = None, ) -> List[Dict[str, Union[Dict, str]]]: after = "" if last_uid is None else f", after: {last_uid}" paging = f"first: {batch_size}{after}" query = f""" {{ q(func: le(last_index_time, {ttl_cutoff_ms}), {paging}) {{ uid, expand(_all_) {{ uid }} }} }} """ txn = client.txn() try: app.log.debug(f"retrieving batch: {query}") batch = txn.query(query) app.log.debug(f"retrieved batch: {batch.json}") return json.loads(batch.json)["q"] finally: txn.discard()
def _upsert(client: GraphClient, node_dict: Dict[str, Any]) -> str: node_dict["uid"] = "_:blank-0" node_key = node_dict["node_key"] query = f""" {{ q0(func: eq(node_key, "{node_key}"), first: 1) {{ uid, dgraph.type expand(_all_) }} }} """ txn = client.txn(read_only=False) try: res = json.loads(txn.query(query).json)["q0"] new_uid = None if res: node_dict["uid"] = res[0]["uid"] new_uid = res[0]["uid"] mutation = node_dict mut_res = txn.mutate(set_obj=mutation, commit_now=True) new_uid = node_dict.get("uid") or mut_res.uids["blank-0"] return cast(str, new_uid) finally: txn.discard()
def query_dgraph_type(client: GraphClient, type_name: str) -> List[QueryPredicateResult]: query = f""" schema(type: {type_name}) {{ type }} """ txn = client.txn(read_only=True) try: res = json.loads(txn.query(query).json) finally: txn.discard() if not res: return [] if not res.get("types"): return [] res = res["types"][0]["fields"] predicate_names = [] for pred in res: predicate_names.append(pred["name"]) predicate_metas = [] for predicate_name in predicate_names: predicate_metas.append(query_dgraph_predicate(client, predicate_name)) return predicate_metas
def _upsert(client: GraphClient, node_dict: Dict[str, Any]) -> str: node_dict["uid"] = "_:blank-0" node_key = node_dict["node_key"] query = f""" {{ q0(func: eq(node_key, "{node_key}"), first: 1) {{ uid, dgraph.type expand(_all_) }} }} """ with client.txn_context(read_only=False) as txn: res = json.loads(txn.query(query).json)["q0"] new_uid = None if res: node_dict["uid"] = res[0]["uid"] new_uid = res[0]["uid"] mutation = node_dict m_res = txn.mutate(set_obj=mutation, commit_now=True) uids = m_res.uids if new_uid is None: new_uid = uids["blank-0"] return str(new_uid)
def query(self, graph_client: GraphClient, first: int) -> List[V]: var_alloc, query = gen_query(self, "q0", first=first) variables = {v: k for k, v in var_alloc.allocated.items()} txn = graph_client.txn(read_only=True) with graph_client.txn_context(read_only=True) as txn: try: qres = json.loads(txn.query(query, variables=variables).json) except Exception as e: raise QueryFailedException(query, variables) from e d = qres.get("q0") if d: return [ self.associated_viewable().from_dict(node, graph_client) for node in d ] return []
def create_edge(client: GraphClient, from_uid: str, edge_name: str, to_uid: str) -> None: if edge_name[0] == "~": mut = {"uid": to_uid, edge_name[1:]: {"uid": from_uid}} else: mut = {"uid": from_uid, edge_name: {"uid": to_uid}} with client.txn_context(read_only=False) as txn: txn.mutate(set_obj=mut, commit_now=True)
def set_property(client: GraphClient, uid: str, prop_name: str, prop_value: Any) -> None: LOGGER.debug(f"Setting property {prop_name} as {prop_value} for {uid}") txn = client.txn(read_only=False) try: mutation = {"uid": uid, prop_name: prop_value} txn.mutate(set_obj=mutation, commit_now=True) finally: txn.discard()
def execute_file( self, name: str, file: str, graph: SubgraphView, sender: Connection, msg_id: str, chunk_size: int, ) -> None: try: pool = ThreadPool(processes=4) exec(file, globals()) client = GraphClient() analyzers = get_analyzer_objects(client) if not analyzers: self.logger.warning(f"Got no analyzers for file: {name}") self.logger.info(f"Executing analyzers: {[an for an in analyzers.keys()]}") for nodes in chunker([n for n in graph.node_iter()], chunk_size): self.logger.info(f"Querying {len(nodes)} nodes") def exec_analyzer( nodes: List[BaseView], sender: Connection ) -> List[BaseView]: try: self.exec_analyzers( client, file, msg_id, nodes, analyzers, sender ) return nodes except Exception as e: self.logger.error(traceback.format_exc()) self.logger.error( f"Execution of {name} failed with {e} {e.args}" ) sender.send(ExecutionFailed()) raise pool.apply_async(exec_analyzer, args=(nodes, sender)) pool.close() pool.join() sender.send(ExecutionComplete()) except Exception as e: self.logger.error(traceback.format_exc()) self.logger.error(f"Execution of {name} failed with {e} {e.args}") sender.send(ExecutionFailed()) raise
def delete_nodes(client: GraphClient, nodes: Iterator[str]) -> int: del_ = [{"uid": uid} for uid in nodes] txn = client.txn() try: mut = txn.create_mutation(del_obj=del_) app.log.debug(f"deleting nodes: {mut}") txn.mutate(mutation=mut, commit_now=True) app.log.debug(f"deleted nodes: {json.dumps(del_)}") return len(del_) finally: txn.discard()
def query_dgraph_predicate(client: GraphClient, predicate_name: str) -> QueryPredicateResult: query = f""" schema(pred: {predicate_name}) {{ }} """ txn = client.txn(read_only=True) try: res = json.loads(txn.query(query).json)["schema"][0] finally: txn.discard() return res
def set_score(client: GraphClient, uid: str, new_score: int, txn: Any = None) -> None: if not txn: txn = client.txn(read_only=False) try: mutation = {"uid": uid, "score": new_score} txn.mutate(set_obj=mutation, commit_now=True) finally: txn.discard()
async def lambda_handler(s3_event: SQSMessageBody, context: Any) -> None: graph_client = GraphClient() s3 = S3ResourceFactory(boto3).from_env() metrics = create_metrics_client() for event in s3_event["Records"]: with metrics.time_to_process_event(): try: _process_one_event(event, s3, graph_client, metrics) except: metrics.event_processed(status="failure") raise else: metrics.event_processed(status="success")
def create_edge(client: GraphClient, from_uid: int, edge_name: str, to_uid: int) -> None: if edge_name[0] == "~": mut = {"uid": to_uid, edge_name[1:]: {"uid": from_uid}} else: mut = {"uid": from_uid, edge_name: {"uid": to_uid}} txn = client.txn(read_only=False) try: res = txn.mutate(set_obj=mut, commit_now=True) LOGGER.debug("edge mutation result is: {}".format(res)) finally: txn.discard()
def delete_edges(client: GraphClient, edges: Iterator[Tuple[str, str, str]]) -> int: del_ = [ create_edge_obj(src_uid, predicate, dest_uid) for src_uid, predicate, dest_uid in edges ] txn = client.txn() try: mut = txn.create_mutation(del_obj=del_) app.log.debug(f"deleting edges: {mut}") txn.mutate(mutation=mut, commit_now=True) app.log.debug(f"deleted edges: {json.dumps(del_)}") return len(del_) finally: txn.discard()
def add_reverse_edge_type(client: GraphClient, uid_type: UidType, edge_name: str) -> None: LOGGER.debug( f"adding reverse edge type uid_type: {uid_type} edge_name: {edge_name}" ) self_type = uid_type._inner_type.self_type() existing_predicates = query_dgraph_type(client, self_type) predicates = "\n\t\t".join(existing_predicates) # In case we've already deployed this plugin if edge_name in predicates: return predicates += f"\n\t\t<~{edge_name}>" type_str = f""" type {self_type} {{ {predicates} }}\n """ op = pydgraph.Operation(schema=type_str) client.alter(op)
def execute_file(name: str, file: str, graph: SubgraphView, sender, msg_id): try: pool = ThreadPool(processes=4) exec(file, globals()) client = GraphClient() analyzers = get_analyzer_objects(client) if not analyzers: LOGGER.warning(f"Got no analyzers for file: {name}") LOGGER.info(f"Executing analyzers: {[an for an in analyzers.keys()]}") chunk_size = 100 if IS_RETRY == "True": chunk_size = 10 for nodes in chunker([n for n in graph.node_iter()], chunk_size): LOGGER.info(f"Querying {len(nodes)} nodes") def exec_analyzer(nodes, sender): try: exec_analyzers(client, file, msg_id, nodes, analyzers, sender) return nodes except Exception as e: LOGGER.error(traceback.format_exc()) LOGGER.error( f"Execution of {name} failed with {e} {e.args}") sender.send(ExecutionFailed()) raise exec_analyzer(nodes, sender) pool.apply_async(exec_analyzer, args=(nodes, sender)) pool.close() pool.join() sender.send(ExecutionComplete()) except Exception as e: LOGGER.error(traceback.format_exc()) LOGGER.error(f"Execution of {name} failed with {e} {e.args}") sender.send(ExecutionFailed()) raise
def get_or_create(gclient: GraphClient, lens_name: str, lens_type: str) -> LensView: with gclient.txn_context(read_only=False) as txn: query = """ # lens get_or_create query res($a: string) { res(func: eq(node_key, $a), first: 1) @cascade { uid, node_type: dgraph.type, node_key, } }""" variables = {"$a": f"lens-{lens_type}{lens_name}"} res = txn.query(query, variables=variables) res = json.loads(res.json)["res"] new_uid = None if res: new_uid = res[0]["uid"] else: m_res = txn.mutate( set_obj={ "lens_name": lens_name, "lens_type": lens_type, "node_key": "lens-" + lens_type + lens_name, "dgraph.type": "Lens", "score": 0, }, commit_now=True, ) uids = m_res.uids new_uid = new_uid or uids["blank-0"] self_lens_query = LensQuery().with_node_key(eq="lens-" + lens_type + lens_name) self_lens = self_lens_query.query_first(gclient) assert ( self_lens ), f"Lens must exist, but couldn't query: {self_lens_query.debug_query()}" return self_lens
def get_or_create(gclient: GraphClient, lens_name: str, lens_type: str) -> "LensView": eg_txn = gclient.txn(read_only=False) try: query = """ query res($a: string) { res(func: eq(node_key, $a), first: 1) @cascade { uid, node_type: dgraph.type, node_key, } }""" res = eg_txn.query( query, variables={"$a": "lens-" + lens_type + lens_name}) res = json.loads(res.json)["res"] new_uid = None if res: new_uid = res[0]["uid"] else: m_res = eg_txn.mutate( set_obj={ "lens": lens_name, "lens_type": lens_type, "node_key": "lens-" + lens_type + lens_name, "dgraph.type": "Lens", "score": 0, }, commit_now=True, ) uids = m_res.uids new_uid = new_uid or uids["blank-0"] finally: eg_txn.discard() self_lens = LensQuery().with_lens_name( eq=lens_name).query_first(gclient) assert self_lens, "Lens must exist" return self_lens
def get_uid(client: GraphClient, node_key: str) -> str: with client.txn_context(read_only=True) as txn: query = """ query res($a: string) { res(func: eq(node_key, $a), first: 1) @cascade { uid, } }""" res = txn.query(query, variables={"$a": node_key}) res = json.loads(res.json) if isinstance(res["res"], list): if res["res"]: return str(res["res"][0]["uid"]) else: raise Exception( f"get_uid failed for node_key: {node_key} {res}") else: return str(res["res"]["uid"])
def query_dgraph_type(client: GraphClient, type_name: str) -> List[str]: query = f""" schema(type: {type_name}) {{ }} """ LOGGER.debug(f"query: {query}") txn = client.txn(read_only=True) try: res = json.loads(txn.query(query).json) LOGGER.debug(f"res: {res}") finally: txn.discard() pred_names = [] if "types" in res: for field in res["types"][0]["fields"]: pred_name = (f"<{field['name']}>" if field["name"].startswith("~") else field["name"]) pred_names.append(pred_name) return pred_names
def query_first( self, graph_client: GraphClient, contains_node_key: Optional[str] = None, best_effort=False, ) -> Optional[V]: if contains_node_key: var_alloc, query = gen_query_parameterized(self, "q0", contains_node_key, 0) else: var_alloc, query = gen_query(self, "q0", first=1) variables = {v: k for k, v in var_alloc.allocated.items()} with graph_client.txn_context(read_only=True, best_effort=best_effort) as txn: try: qres = json.loads(txn.query(query, variables=variables).json) except Exception as e: raise QueryFailedException(query, variables) from e d = qres.get("q0") if d: return self.associated_viewable().from_dict(d[0], graph_client) return None
def test_single_file_contains_key( self, node_key, file_path, file_extension, file_mime_type, file_size, file_version, file_description, file_product, file_company, file_directory, file_inode, file_hard_links, signed, signed_status, md5_hash, sha1_hash, sha256_hash, ) -> None: node_key = "test_single_file_contains_key" + str(node_key) signed = "true" if signed else "false" graph_client = GraphClient() get_or_create_file_node( graph_client, node_key, file_path=file_path, file_extension=file_extension, file_mime_type=file_mime_type, file_size=file_size, file_version=file_version, file_description=file_description, file_product=file_product, file_company=file_company, file_directory=file_directory, file_inode=file_inode, file_hard_links=file_hard_links, signed=signed, signed_status=signed_status, md5_hash=md5_hash, sha1_hash=sha1_hash, sha256_hash=sha256_hash, ) queried_proc = FileQuery().query_first(graph_client, contains_node_key=node_key) assert node_key == queried_proc.node_key assert file_path == queried_proc.get_file_path() or "" assert file_extension == queried_proc.get_file_extension() or "" assert file_mime_type == queried_proc.get_file_mime_type() or "" assert file_version == queried_proc.get_file_version() or "" assert file_description == queried_proc.get_file_description() or "" assert file_product == queried_proc.get_file_product() or "" assert file_company == queried_proc.get_file_company() or "" assert file_directory == queried_proc.get_file_directory() or "" assert file_hard_links == queried_proc.get_file_hard_links() or "" assert signed == queried_proc.get_signed() or "" assert signed_status == queried_proc.get_signed_status() or "" assert md5_hash == queried_proc.get_md5_hash() or "" assert sha1_hash == queried_proc.get_sha1_hash() or "" assert sha256_hash == queried_proc.get_sha256_hash() or "" assert file_size == queried_proc.get_file_size() assert file_inode == queried_proc.get_file_inode()
async def handle_events(self, events: SQSMessageBody, context: Any) -> None: # Parse sns message self.logger.debug(f"handling events: {events} context: {context}") client = GraphClient() s3 = S3ResourceFactory(boto3).from_env() load_plugins( self.model_plugins_bucket, s3.meta.client, os.path.abspath(MODEL_PLUGINS_DIR), ) for event in events["Records"]: data = parse_s3_event(s3, event) # FIXME: this code assumes inner_message is json envelope = OldEnvelope.deserialize(data) message = json.loads(envelope.inner_message) LOGGER.info(f'Executing Analyzer: {message["key"]}') with self.metric_reporter.histogram_ctx( "analyzer-executor.download_s3_file" ): analyzer = download_s3_file( s3, self.analyzers_bucket, message["key"], ).decode("utf8") analyzer_name = message["key"].split("/")[-2] subgraph = SubgraphView.from_proto(client, bytes(message["subgraph"])) # TODO: Validate signature of S3 file LOGGER.info(f"event {event} {envelope.metadata}") rx: Connection tx: Connection rx, tx = Pipe(duplex=False) p = Process( target=self.execute_file, args=(analyzer_name, analyzer, subgraph, tx, "", self.chunk_size), ) p.start() for exec_hit in self.poll_process(rx=rx, analyzer_name=analyzer_name): with self.metric_reporter.histogram_ctx( "analyzer-executor.emit_event.ms", (TagPair("analyzer_name", exec_hit.analyzer_name),), ): emit_event( self.analyzer_matched_subgraphs_bucket, s3, exec_hit, envelope.metadata, ) self.update_msg_cache(analyzer, exec_hit.root_node_key, message["key"]) self.update_hit_cache(analyzer_name, exec_hit.root_node_key) p.join()
def set_schema(client: GraphClient, schema: str) -> None: op = pydgraph.Operation(schema=schema, run_in_background=True) LOGGER.info(f"Setting dgraph schema: {schema}") client.alter(op, timeout=SecsDuration(5)) LOGGER.info(f"Completed setting dgraph schema")
def test_single_file_view_parity_eq( self, node_key, file_path, file_extension, file_mime_type, file_size, file_version, file_description, file_product, file_company, file_directory, file_inode, file_hard_links, signed, signed_status, md5_hash, sha1_hash, sha256_hash, ): node_key = "test_single_file_view_parity_eq" + str(node_key) signed = "true" if signed else "false" graph_client = GraphClient() get_or_create_file_node( graph_client, node_key, file_path=file_path, file_extension=file_extension, file_mime_type=file_mime_type, file_size=file_size, file_version=file_version, file_description=file_description, file_product=file_product, file_company=file_company, file_directory=file_directory, file_inode=file_inode, file_hard_links=file_hard_links, signed=signed, signed_status=signed_status, md5_hash=md5_hash, sha1_hash=sha1_hash, sha256_hash=sha256_hash, ) queried_file = (FileQuery().with_node_key(eq=node_key).with_file_path( eq=file_path).with_file_extension( eq=file_extension).with_file_mime_type( eq=file_mime_type).with_file_size( eq=file_size).with_file_version( eq=file_version).with_file_description( eq=file_description).with_file_product( eq=file_product). with_file_company(eq=file_company).with_file_directory( eq=file_directory).with_file_inode( eq=file_inode).with_file_hard_links( eq=file_hard_links).with_signed(eq=signed). with_signed_status(eq=signed_status).with_md5_hash( eq=md5_hash).with_sha1_hash( eq=sha1_hash).with_sha256_hash( eq=sha256_hash).query_first(graph_client)) assert node_key == queried_file.node_key assert file_path == queried_file.get_file_path() assert file_extension == queried_file.get_file_extension() assert file_mime_type == queried_file.get_file_mime_type() assert file_size == queried_file.get_file_size() assert file_version == queried_file.get_file_version() assert file_description == queried_file.get_file_description() assert file_product == queried_file.get_file_product() assert file_company == queried_file.get_file_company() assert file_directory == queried_file.get_file_directory() assert file_inode == queried_file.get_file_inode() assert file_hard_links == queried_file.get_file_hard_links() assert signed == queried_file.get_signed() assert signed_status == queried_file.get_signed_status() assert md5_hash == queried_file.get_md5_hash() assert sha1_hash == queried_file.get_sha1_hash() assert sha256_hash == queried_file.get_sha256_hash()
def lambda_handler_fn(events: Any, context: Any) -> None: # Parse sns message LOGGER.debug(f"handling events: {events} context: {context}") client = GraphClient() s3 = get_s3_client() load_plugins(os.environ["BUCKET_PREFIX"], s3, os.path.abspath(MODEL_PLUGINS_DIR)) for event in events["Records"]: if not IS_LOCAL: event = json.loads(event["body"])["Records"][0] data = parse_s3_event(s3, event) message = json.loads(data) LOGGER.info(f'Executing Analyzer: {message["key"]}') analyzer = download_s3_file( s3, f"{os.environ['BUCKET_PREFIX']}-analyzers-bucket", message["key"]) analyzer_name = message["key"].split("/")[-2] subgraph = SubgraphView.from_proto(client, bytes(message["subgraph"])) # TODO: Validate signature of S3 file LOGGER.info(f"event {event}") rx: Connection tx: Connection rx, tx = Pipe(duplex=False) p = Process(target=execute_file, args=(analyzer_name, analyzer, subgraph, tx, "")) p.start() t = 0 while True: p_res = rx.poll(timeout=5) if not p_res: t += 1 LOGGER.info( f"Polled {analyzer_name} for {t * 5} seconds without result" ) continue result: Optional[Any] = rx.recv() if isinstance(result, ExecutionComplete): LOGGER.info("execution complete") break # emit any hits to an S3 bucket if isinstance(result, ExecutionHit): LOGGER.info( f"emitting event for {analyzer_name} {result.analyzer_name} {result.root_node_key}" ) emit_event(s3, result) update_msg_cache(analyzer, result.root_node_key, message["key"]) update_hit_cache(analyzer_name, result.root_node_key) assert not isinstance( result, ExecutionFailed), f"Analyzer {analyzer_name} failed." p.join()
def __init__(self, query: BaseQuery, graph_client: GraphClient = None) -> None: self.query = query self.graph_client = graph_client or GraphClient()
def set_schema(client: GraphClient, schema: str) -> None: op = pydgraph.Operation(schema=schema) client.alter(op)