def __init__(self, id: str):
self.engagement_key = id
self.eg_client: DgraphClient = DgraphClient(
DgraphClientStub("alpha0.engagementgraphcluster.grapl:9080"))
self.mg_client: DgraphClient = DgraphClient(
DgraphClientStub("alpha0.mastergraphcluster.grapl:9080"))
self.node_copier: NodeCopier = NodeCopier(id, self.mg_client,
self.eg_client)
def test_process_with_created_files(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"763ddbda-8812-4a07-acfe-83402b92379d",
parent_process,
)
created_file = {
"file_path": "/folder/file.txt",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
created_file_view = upsert(
local_client,
"File",
FileView,
"575f103e-1a11-4650-9f1b-5b72e44dfec3",
created_file,
)
create_edge(
local_client,
parent_process_view.uid,
"created_files",
created_file_view.uid,
)
queried_process = (
ProcessQuery()
.with_node_key(eq="763ddbda-8812-4a07-acfe-83402b92379d")
.with_process_id(eq=100)
.with_process_name(contains="word")
.with_created_timestamp(eq=created_timestamp)
.with_created_files(
FileQuery()
.with_node_key(eq="575f103e-1a11-4650-9f1b-5b72e44dfec3")
.with_file_path(eq="/folder/file.txt")
)
.query_first(local_client)
)
assert queried_process
assert queried_process.process_id == 100
assert len(queried_process.created_files) == 1
created_file = queried_process.created_files[0]
assert created_file.file_path == "/folder/file.txt"
def test_process_query_view_parity(self, process_props: ProcessProps):
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_proc = get_or_create_process(self, local_client, process_props,)
queried_proc = (
ProcessQuery()
.with_node_key(eq=created_proc.node_key)
.query_first(local_client)
)
assert queried_proc
assert process_props["node_key"] == queried_proc.node_key
assert "Process" == queried_proc.get_node_type()
assert process_props["process_id"] == queried_proc.get_process_id()
assert process_props["arguments"] == escape_dgraph_str(
queried_proc.get_arguments()
)
assert (
process_props["created_timestamp"] == queried_proc.get_created_timestamp()
)
assert None == queried_proc.get_asset()
assert process_props["terminate_time"] == queried_proc.get_terminate_time()
assert process_props["image_name"] == escape_dgraph_str(
queried_proc.get_image_name()
)
assert process_props["process_name"] == escape_dgraph_str(
queried_proc.get_process_name()
)
def test_process_query_view_miss(self, process_props: ProcessProps) -> None:
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_proc = get_or_create_process(self, local_client, process_props)
assert (
created_proc.process_id is not None
and created_proc.arguments is not None
and created_proc.created_timestamp is not None
and created_proc.terminate_time is not None
and created_proc.image_name is not None
and created_proc.process_name is not None
)
queried_proc = (
ProcessQuery()
.with_node_key(eq=created_proc.node_key)
.with_process_id(eq=Not(created_proc.process_id))
.with_arguments(eq=Not(created_proc.arguments))
.with_created_timestamp(eq=Not(created_proc.created_timestamp))
.with_terminate_time(eq=Not(created_proc.terminate_time))
.with_image_name(eq=Not(created_proc.image_name))
.with_process_name(eq=Not(created_proc.process_name))
.query_first(local_client)
)
assert not queried_proc
def test_with_wrote_files(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"test_with_wrote_files-8f0761fb-2ffe-4d4b-ab38-68e5489f56dc",
parent_process,
)
wrote_file = {
"file_path": "/folder/file.txt",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
wrote_file_view = upsert(
local_client,
"File",
FileView,
"test_with_wrote_files-2325c49a-95b4-423f-96d0-99539fe03833",
wrote_file,
)
create_edge(
local_client,
parent_process_view.uid,
"wrote_files",
wrote_file_view.uid,
)
queried_process = (ProcessQuery().with_node_key(
eq="test_with_wrote_files-8f0761fb-2ffe-4d4b-ab38-68e5489f56dc"
).with_process_id(eq=100).with_process_name(
contains="word"
).with_created_timestamp(eq=created_timestamp).with_wrote_files(
FileQuery().with_node_key(
eq="test_with_wrote_files-2325c49a-95b4-423f-96d0-99539fe03833"
).with_file_path(eq="/folder/file.txt")).query_first(local_client))
assert queried_process
assert (queried_process.node_key ==
"test_with_wrote_files-8f0761fb-2ffe-4d4b-ab38-68e5489f56dc")
assert queried_process.process_id == 100
assert queried_process.process_name == "word.exe"
assert len(queried_process.wrote_files) == 1
assert (queried_process.wrote_files[0].node_key ==
"test_with_wrote_files-2325c49a-95b4-423f-96d0-99539fe03833")
assert queried_process.wrote_files[0].file_path == "/folder/file.txt"
def analyzer(graph: Subgraph, sender: Connection):
try:
print('analyzing')
client = DgraphClient(DgraphClientStub('db.mastergraph:9080'))
_analyzer(client, graph, sender)
except Exception as e:
print('analyzer failed: {}'.format(e))
sender.send(None)
def lambda_handler(events: Any, context: Any) -> None:
# Parse sns message
print("handling")
print(events)
print(context)
alpha_names = os.environ["MG_ALPHAS"].split(",")
client_stubs = [DgraphClientStub("{}:9080".format(name)) for name in alpha_names]
client = DgraphClient(*client_stubs)
s3 = get_s3_client()
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)
print(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
print(f'event {event}')
rx, tx = Pipe(duplex=False) # type: Tuple[Connection, Connection]
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
print(f"Polled {analyzer_name} for {t * 5} seconds without result")
continue
result = rx.recv() # type: Optional[Any]
if isinstance(result, ExecutionComplete):
print("execution complete")
break
# emit any hits to an S3 bucket
if isinstance(result, ExecutionHit):
print(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 test_with_bin_file(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"635952af-87f3-4a2a-a65d-3f1859db9525",
parent_process,
)
bin_file = {
"file_path": "/folder/file.txt",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
bin_file_view = upsert(
local_client,
"File",
FileView,
"9f16e0c9-33c0-4d18-9878-ef686373570b",
bin_file,
)
create_edge(
local_client,
parent_process_view.uid,
"bin_file",
bin_file_view.uid,
)
queried_process = (ProcessQuery().with_node_key(
eq="635952af-87f3-4a2a-a65d-3f1859db9525"
).with_process_id(eq=100).with_process_name(
contains="word").with_created_timestamp(
eq=created_timestamp).with_bin_file(FileQuery().with_node_key(
eq="9f16e0c9-33c0-4d18-9878-ef686373570b").with_file_path(
eq="/folder/file.txt")).query_first(local_client))
assert queried_process
assert "635952af-87f3-4a2a-a65d-3f1859db9525"
assert queried_process.process_id == 100
assert queried_process.process_name == "word.exe"
assert queried_process.created_timestamp == created_timestamp
bin_file = queried_process.bin_file
assert bin_file.node_key == "9f16e0c9-33c0-4d18-9878-ef686373570b"
assert bin_file.file_path == "/folder/file.txt"
def test_with_read_files(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"test_with_read_files-669a3693-d960-401c-8d29-5d669ffcd660",
parent_process,
)
read_file = {
"file_path": "/folder/file.txt",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
read_file_view = upsert(
local_client,
"File",
FileView,
"test_with_read_files-aa9248ec-36ee-4177-ba1a-999de735e682",
read_file,
)
create_edge(
local_client,
parent_process_view.uid,
"read_files",
read_file_view.uid,
)
queried_process = (ProcessQuery().with_process_id(
eq=100).with_process_name(contains="word").with_created_timestamp(
eq=created_timestamp).with_read_files(
FileQuery().with_file_path(
eq="/folder/file.txt")).query_first(local_client))
assert queried_process
assert (queried_process.node_key ==
"test_with_read_files-669a3693-d960-401c-8d29-5d669ffcd660")
assert queried_process.process_id == 100
assert queried_process.process_name == "word.exe"
assert len(queried_process.read_files) == 1
assert (queried_process.read_files[0].node_key ==
"test_with_read_files-aa9248ec-36ee-4177-ba1a-999de735e682")
assert queried_process.read_files[0].file_path == "/folder/file.txt"
def main() -> None:
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
test_ipc(local_client)
parent = {
"process_id": 100,
"process_name": "word.exe",
} # type: Dict[str, Property]
child = {
"process_id": 1234,
"process_name": "cmd.exe"
} # type: Dict[str, Property]
external_ip = {
"external_ip": "56.123.14.24",
} # type: Dict[str, Property]
parent_view = upsert(local_client, ProcessView,
"ea75f056-61a1-479d-9ca2-f632d2c67205", parent)
child_view = upsert(local_client, ProcessView,
"10f585c2-cf31-41e2-8ca5-d477e78be3ac", child)
external_ip_view = upsert(
local_client,
ExternalIpView,
"8bc20354-e8c5-49fc-a984-2927b24c1a29",
external_ip,
)
create_edge(local_client, parent_view.uid, "children", child_view.uid)
create_edge(local_client, child_view.uid, "created_connections",
external_ip_view.uid)
queried_child_0 = ProcessQuery().with_process_id(
eq=1234).query_first(local_client)
assert queried_child_0
assert queried_child_0.node_key == child_view.node_key
queried_child_1 = (ProcessQuery().with_process_id(eq=1234).query_first(
local_client,
contains_node_key="10f585c2-cf31-41e2-8ca5-d477e78be3ac"))
assert queried_child_1
assert queried_child_1.node_key == child_view.node_key
assert queried_child_1.node_key == queried_child_0.node_key
p = (ProcessQuery().with_process_name(
eq="cmd.exe").with_parent().with_created_connection().query_first(
local_client)) # type: Optional[ProcessView]
assert p
def execute_file(file: str, graph: SubgraphView, sender):
print("executing analyzer")
alpha_names = os.environ["MG_ALPHAS"].split(",")
client_stubs = [DgraphClientStub("{}:9080".format(name)) for name in alpha_names]
client = DgraphClient(*client_stubs)
exec(file, globals())
for node in graph.node_iter():
# TODO: Check node + analyzer file hash in redis cache, avoid reprocess of hits
print("File executed: {}".format(analyzer(client, node, sender))) # type: ignore
def main() -> None:
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
parent = {
"process_id": 100,
"process_name": "word.exe",
} # type: Dict[str, Property]
child = {"process_id": 1234, "process_name": "cmd.exe"} # type: Dict[str, Property]
parent_view = upsert(
local_client,
"Process",
ProcessView,
"ea75f056-61a1-479d-9ca2-f632d2c67205",
parent,
)
child_view = upsert(
local_client,
"Process",
ProcessView,
"10f585c2-cf31-41e2-8ca5-d477e78be3ac",
child,
)
create_edge(local_client, parent_view.uid, "children", child_view.uid)
queried_child_0 = ProcessQuery().with_process_id(eq=1234).query_first(local_client)
assert queried_child_0
assert queried_child_0.node_key == child_view.node_key
queried_child_1 = (
ProcessQuery()
.with_process_id(eq=1234)
.query_first(
local_client, contains_node_key="10f585c2-cf31-41e2-8ca5-d477e78be3ac"
)
)
assert queried_child_1
assert queried_child_1.node_key == child_view.node_key
assert queried_child_1.node_key == queried_child_0.node_key
p = (
ProcessQuery().with_parent().query_first(local_client)
) # type: Optional[ProcessView]
assert p
p.get_parent()
def execute_file(name: str, file: str, graph: SubgraphView, sender, msg_id):
alpha_names = os.environ["MG_ALPHAS"].split(",")
try:
pool = ThreadPool(processes=4)
exec(file, globals())
client_stubs = [
DgraphClientStub(f"{a_name}:9080") for a_name in alpha_names
]
client = DgraphClient(*client_stubs)
analyzers = get_analyzer_objects(client)
if not analyzers:
print(f'Got no analyzers for file: {name}')
print(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):
print(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:
print(traceback.format_exc())
print(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:
print(traceback.format_exc())
print(f'Execution of {name} failed with {e} {e.args}')
sender.send(ExecutionFailed())
raise
def main() -> None:
local_client = DgraphClient(DgraphClientStub('localhost:9080'))
test_ipc(local_client)
parent = {
'process_id': 100,
'process_name': 'word.exe'
} # type: Dict[str, Property]
child = {
'process_id': 1234,
'process_name': 'cmd.exe'
} # type: Dict[str, Property]
external_ip = {
'external_ip': '56.123.14.24',
} # type: Dict[str, Property]
parent_view = upsert(local_client, ProcessView,
'ea75f056-61a1-479d-9ca2-f632d2c67205', parent)
child_view = upsert(local_client, ProcessView,
'10f585c2-cf31-41e2-8ca5-d477e78be3ac', child)
external_ip_view = upsert(local_client, ExternalIpView,
'8bc20354-e8c5-49fc-a984-2927b24c1a29',
external_ip)
create_edge(local_client, parent_view.uid, 'children', child_view.uid)
create_edge(local_client, child_view.uid, 'created_connections',
external_ip_view.uid)
queried_child_0 = ProcessQuery().with_process_id(
eq=1234).query_first(local_client)
assert queried_child_0
assert queried_child_0.node_key == child_view.node_key
queried_child_1 = (ProcessQuery().with_process_id(eq=1234).query_first(
local_client,
contains_node_key='10f585c2-cf31-41e2-8ca5-d477e78be3ac'))
assert queried_child_1
assert queried_child_1.node_key == child_view.node_key
assert queried_child_1.node_key == queried_child_0.node_key
p = (ProcessQuery().with_process_name(
eq="cmd.exe").with_parent().with_created_connection().query_first(
local_client)) # type: Optional[ProcessView]
assert p
def lambda_handler(events: Any, context: Any) -> None:
# Parse sns message
print("handling")
print(events)
print(context)
alpha_names = os.environ["MG_ALPHAS"].split(",")
client_stubs = [DgraphClientStub("{}:9080".format(name)) for name in alpha_names]
client = DgraphClient(*client_stubs)
for event in events["Records"]:
data = parse_s3_event(event)
message = json.loads(data)
# TODO: Use env variable for s3 bucket
analyzer = download_s3_file("grapl-analyzers-bucket", message["key"])
subgraph = SubgraphView.from_proto(client, bytes(message["subgraph"]))
# TODO: Validate signature of S3 file
print("creating queue")
rx, tx = Pipe(duplex=False) # type: Tuple[Connection, Connection]
print("creating process")
p = Process(target=execute_file, args=(analyzer, subgraph, tx))
print("running process")
p.start()
while True:
print("waiting for results")
p_res = rx.poll(timeout=5)
if not p_res:
print("Polled for 5 seconds without result")
continue
result = rx.recv() # type: Optional[Any]
if isinstance(result, ExecutionComplete):
print("execution complete")
# TODO: ACK Result
break
# emit any hits to an S3 bucket
if isinstance(result, ExecutionHit):
print("emitting event for result: {}".format(result))
emit_event(result)
assert not isinstance(
result, ExecutionFailed
), "Result was none. Analyzer failed."
p.join()
def test_with_deleted_files(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"test_with_deleted_files-47527d73-22c4-4e0f-bf7d-184bf1f206e2",
parent_process,
)
deleted_file = {
"file_path": "/folder/file.txt",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
deleted_file_view = upsert(
local_client,
"File",
FileView,
"test_with_deleted_files8b8364ea-9b47-476b-8cf0-0f724adff10f",
deleted_file,
)
create_edge(
local_client,
parent_process_view.uid,
"deleted_files",
deleted_file_view.uid,
)
queried_process = (
ProcessQuery()
.with_process_id(eq=100)
.with_process_name(contains="word")
.with_created_timestamp(eq=created_timestamp)
.with_deleted_files(FileQuery().with_file_path(eq="/folder/file.txt"))
.query_first(local_client)
)
assert queried_process
assert queried_process.process_id == 100
def test_process_query_view_parity_contains(
self,
node_key,
process_id,
created_timestamp,
asset_id,
terminate_time,
image_name,
process_name,
arguments,
):
node_key = "test_process_query_view_parity_contains" + str(node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
get_or_create_process_node(
local_client,
node_key,
process_id,
arguments,
created_timestamp,
asset_id,
terminate_time,
image_name,
process_name,
)
query = ProcessQuery().with_node_key(eq=node_key)
# Don't f**k with newlines due to a dgraph bug
# https://github.com/dgraph-io/dgraph/issues/4694
if len(arguments) > 3 and "\n" not in arguments:
query.with_arguments(contains=arguments[:len(arguments) - 1])
if len(asset_id) > 3 and "\n" not in asset_id:
query.with_asset_id(contains=asset_id[:len(asset_id) - 1])
if len(image_name) > 3 and "\n" not in image_name:
query.with_image_name(contains=image_name[:len(image_name) - 1])
if len(process_name) > 3 and "\n" not in process_name:
query.with_process_name(contains=process_name[:len(process_name) -
1])
queried_proc = query.query_first(local_client)
assert queried_proc
assert "Process" == queried_proc.get_node_type()
assert process_id == queried_proc.get_process_id()
assert node_key == queried_proc.node_key
assert arguments == queried_proc.get_arguments()
assert created_timestamp == queried_proc.get_created_timestamp()
assert asset_id == queried_proc.get_asset_id()
assert terminate_time == queried_proc.get_terminate_time()
assert image_name == queried_proc.get_image_name()
assert process_name == queried_proc.get_process_name()
def execute_file(name: str, file: str, graph: SubgraphView, sender, msg_id):
try:
pool = ThreadPool(processes=4)
exec(file, globals())
client_stubs = (DgraphClientStub(f"{host}:{port}")
for host, port in mg_alphas())
client = DgraphClient(*client_stubs)
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 test_process_query_view_parity_contains(
self,
node_key,
process_id,
created_timestamp,
terminate_time,
image_name,
process_name,
arguments,
):
node_key = "test_process_query_view_parity_contains" + str(node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
get_or_create_process_node_deprecated(
local_client,
node_key,
process_id,
arguments,
created_timestamp,
terminate_time,
image_name,
process_name,
)
query = ProcessQuery().with_node_key(eq=node_key)
# Don't f**k with newlines due to a dgraph bug
# https://github.com/dgraph-io/dgraph/issues/4694
for prop in [arguments, image_name, process_name]:
hypothesis.assume(len(prop) > 3)
hypothesis.assume("\n" not in prop)
hypothesis.assume("\\" not in prop)
# These fail because dgraph doesn't like the query
# (regexp(process_name, /00\\//))
query.with_arguments(contains=arguments[: len(arguments) - 1])
query.with_image_name(contains=image_name[: len(image_name) - 1])
query.with_process_name(contains=process_name[: len(process_name) - 1])
queried_proc = query.query_first(local_client)
assert queried_proc
assert "Process" == queried_proc.get_node_type()
assert process_id == queried_proc.get_process_id()
assert node_key == queried_proc.node_key
assert arguments == queried_proc.get_arguments()
assert created_timestamp == queried_proc.get_created_timestamp()
assert terminate_time == queried_proc.get_terminate_time()
assert image_name == queried_proc.get_image_name()
assert process_name == queried_proc.get_process_name()
def exec_analyzer(analyzer, node, sender):
try:
client_stubs = [
DgraphClientStub(f"{a_name}:9080")
for a_name in alpha_names
]
client = DgraphClient(*client_stubs)
analyzer(client, node, sender)
return node
except Exception as e:
print(traceback.format_exc())
print(f'Execution of {name} failed with {e} {e.args}')
sender.send(ExecutionFailed())
raise
def lambda_handler(events: Any, context: Any) -> None:
mg_client = DgraphClient(DgraphClientStub('db.mastergraph:9080'))
eg_client = DgraphClient(DgraphClientStub('db.engagementgraph:9080'))
create_process_schema(eg_client)
uid_map = {}
for event in events['Records']:
print('Copying engagement')
data = parse_s3_event(event)
incident_graph = json.loads(data)
label = incident_graph['label'].encode('utf-8')
node_refs = incident_graph['node_refs']
edges = incident_graph['edges']
engagement_key = get_engagement_key(label, [n['uid'] for n in node_refs])
if should_throttle(engagement_key, eg_client):
print('Throttling: {}'.format(engagement_key))
continue
print('Creating engagement: {}'.format(engagement_key))
copier = NodeCopier(engagement_key, mg_client, eg_client)
print('node_refs: {}'.format(node_refs))
print('edges: {}'.format(edges))
for node_ref in node_refs:
new_uid = copier.copy_node(node_ref['uid'])
uid_map[node_ref['uid']] = new_uid
print('new_uid: {}'.format(new_uid))
for edge in edges:
copier.copy_edge(uid_map[edge[0]], edge[1], uid_map[edge[2]])
print('Copied engagement successfully')
print('Engagement creation was successful')
def test_process_query_view_parity_eq(
self,
node_key,
process_id,
created_timestamp,
terminate_time,
image_name,
process_name,
arguments,
):
node_key = "test_process_query_view_parity_eq" + str(node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
get_or_create_process_node_deprecated(
local_client,
node_key,
process_id,
arguments,
created_timestamp,
terminate_time,
image_name,
process_name,
)
queried_proc = (
ProcessQuery()
.with_node_key(eq=node_key)
.with_process_id(eq=process_id)
.with_arguments(eq=arguments)
.with_created_timestamp(eq=created_timestamp)
.with_terminate_time(eq=terminate_time)
.with_image_name(eq=image_name)
.with_process_name(eq=process_name)
.query_first(local_client)
)
# assert process_view.process_id == queried_proc.get_process_id()
assert node_key == queried_proc.node_key
assert "Process" == queried_proc.get_node_type()
assert process_id == queried_proc.get_process_id()
assert arguments == queried_proc.get_arguments()
assert created_timestamp == queried_proc.get_created_timestamp()
assert terminate_time == queried_proc.get_terminate_time()
assert image_name == queried_proc.get_image_name()
assert process_name == queried_proc.get_process_name()
def main() -> None:
local_client = DgraphClient(DgraphClientStub('localhost:9080'))
parent = {
'process_id': 100,
'process_name': 'word.exe'
} # type: Dict[str, Property]
child = {
'process_id': 1234,
'process_name': 'cmd.exe'
} # type: Dict[str, Property]
parent_view = upsert(local_client, 'Process', ProcessView,
'ea75f056-61a1-479d-9ca2-f632d2c67205', parent)
child_view = upsert(local_client, 'Process', ProcessView,
'10f585c2-cf31-41e2-8ca5-d477e78be3ac', child)
create_edge(local_client, parent_view.uid, 'children', child_view.uid)
queried_child_0 = ProcessQuery().with_process_id(
eq=1234).query_first(local_client)
assert queried_child_0
assert queried_child_0.node_key == child_view.node_key
queried_child_1 = (ProcessQuery().with_process_id(eq=1234).query_first(
local_client,
contains_node_key='10f585c2-cf31-41e2-8ca5-d477e78be3ac'))
assert queried_child_1
assert queried_child_1.node_key == child_view.node_key
assert queried_child_1.node_key == queried_child_0.node_key
p = (ProcessQuery().with_parent().query_first(local_client)
) # type: Optional[ProcessView]
assert p
p.get_parent()
def test_single_process_contains_key(
self,
node_key,
process_id,
created_timestamp,
asset_id,
terminate_time,
image_name,
process_name,
arguments,
):
node_key = "test_single_process_contains_key" + str(node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
get_or_create_process_node(
local_client,
node_key,
process_id,
arguments,
created_timestamp,
asset_id,
terminate_time,
image_name,
process_name,
)
queried_proc = ProcessQuery().query_first(local_client,
contains_node_key=node_key)
# assert process_view.process_id == queried_proc.get_process_id()
assert node_key == queried_proc.node_key
assert "Process" == queried_proc.get_node_type()
assert process_id == queried_proc.get_process_id()
assert arguments == queried_proc.get_arguments()
assert created_timestamp == queried_proc.get_created_timestamp()
assert asset_id == queried_proc.get_asset_id()
assert terminate_time == queried_proc.get_terminate_time()
assert image_name == queried_proc.get_image_name()
assert process_name == queried_proc.get_process_name()
def test_single_process_contains_key(self, process_props: ProcessProps) -> None:
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_proc = get_or_create_process(self, local_client, process_props)
# Setup complete, do some queries
queried_proc = ProcessQuery().query_first(
local_client, contains_node_key=created_proc.node_key
)
assert queried_proc
assert created_proc.get_process_id() == queried_proc.get_process_id()
assert created_proc.node_key == queried_proc.node_key
assert "Process" == queried_proc.get_node_type()
assert created_proc.get_arguments() == queried_proc.get_arguments()
assert (
created_proc.get_created_timestamp() == queried_proc.get_created_timestamp()
)
assert created_proc.get_terminate_time() == queried_proc.get_terminate_time()
assert created_proc.get_image_name() == queried_proc.get_image_name()
assert created_proc.get_process_name() == queried_proc.get_process_name()
assert not queried_proc.get_asset()
def test_process_query_view_miss(
self,
node_key,
process_id,
created_timestamp,
asset_id,
terminate_time,
image_name,
process_name,
arguments,
):
node_key = "test_process_query_view_miss" + str(node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
process = {
"process_id": process_id,
"arguments": arguments,
"created_timestamp": created_timestamp,
"asset_id": asset_id,
"terminate_time": terminate_time,
"image_name": image_name,
"process_name": process_name,
} # type: Dict[str, Property]
process_view = cast(ProcessView,
upsert(local_client, "Process", ProcessView,
node_key, process)) # type: ProcessView
queried_proc = (ProcessQuery().with_node_key(
eq=node_key).with_process_id(eq=Not(process_id)).with_arguments(
eq=Not(arguments)).with_created_timestamp(
eq=Not(created_timestamp)).with_asset_id(
eq=Not(asset_id)).with_terminate_time(
eq=Not(terminate_time)).with_image_name(
eq=Not(image_name)).with_process_name(eq=Not(
process_name)).query_first(local_client))
assert not queried_proc
def test__single_ip_addr_node__query_by_node_key(
self,
node_key,
first_seen_timestamp,
last_seen_timestamp,
ip_address,
):
# current function's name, but don't need to copy-paste replace
node_key = node_key_for_test(self, node_key)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created = get_or_create_ip_address_node(
local_client=local_client,
node_key=node_key,
first_seen_timestamp=as_millis(first_seen_timestamp),
last_seen_timestamp=as_millis(last_seen_timestamp),
ip_address=str(ip_address),
)
queried_ip_address_node = (IpAddressQuery().with_ip_address(
).with_first_seen_timestamp().with_last_seen_timestamp().query_first(
local_client, contains_node_key=node_key))
assert_views_equal(expected=created, actual=queried_ip_address_node)
def test_single_process_connected_to_asset_node(
self, asset_props: AssetProps, process_props: ProcessProps,
):
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_asset = get_or_create_asset(self, local_client, asset_props)
created_proc = get_or_create_process(self, local_client, process_props)
create_edge(
local_client, created_asset.uid, "asset_processes", created_proc.uid,
)
# Setup complete, do some queries
queried_proc = (
ProcessQuery()
.with_asset(AssetQuery().with_hostname(created_asset.get_hostname()))
.query_first(local_client, contains_node_key=created_proc.node_key)
)
assert queried_proc
fetch_all_properties(queried_proc)
assert_equal_props(created_proc, queried_proc)
queried_asset = queried_proc.get_asset()
assert_equal_identity(created_asset, queried_asset)
def test_parent_children_edge(self) -> None:
# Given: a process with a pid 100 & process_name word.exe,
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
created_timestamp = int(time.time())
parent_process = {
"process_id": 100,
"process_name": "word.exe",
"created_timestamp": created_timestamp,
} # type: Dict[str, Property]
parent_process_view = upsert(
local_client,
"Process",
ProcessView,
"0e84f2ce-f711-46ce-bc9e-1b13c9ba6d6c",
parent_process,
)
child_process = {
"process_id": 110,
"process_name": "malware.exe",
"created_timestamp": created_timestamp + 1000,
} # type: Dict[str, Property]
child_process_view = upsert(
local_client,
"Process",
ProcessView,
"46d2862f-cb58-4062-b35e-bb310b8d5b0d",
child_process,
)
create_edge(
local_client, parent_process_view.uid, "children", child_process_view.uid,
)
queried_process = (
ProcessQuery()
.with_node_key(eq="0e84f2ce-f711-46ce-bc9e-1b13c9ba6d6c")
.with_process_id(eq=100)
.with_process_name(contains="word")
.with_created_timestamp(eq=created_timestamp)
.with_children(
ProcessQuery()
.with_node_key(eq="46d2862f-cb58-4062-b35e-bb310b8d5b0d")
.with_process_id(eq=110)
.with_process_name(eq="malware.exe")
.with_created_timestamp(eq=created_timestamp + 1000)
)
.query_first(local_client)
)
assert queried_process
assert queried_process.node_key == "0e84f2ce-f711-46ce-bc9e-1b13c9ba6d6c"
assert queried_process.process_id == 100
assert queried_process.process_name == "word.exe"
assert queried_process.created_timestamp == created_timestamp
assert len(queried_process.children) == 1
child = queried_process.children[0]
assert child.node_key == "46d2862f-cb58-4062-b35e-bb310b8d5b0d"
assert child.process_id == 110
assert child.process_name == "malware.exe"
assert child.created_timestamp == created_timestamp + 1000
def test_single_file_view_parity_eq(
self,
node_key,
file_path,
asset_id,
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)
local_client = DgraphClient(DgraphClientStub("localhost:9080"))
get_or_create_file_node(
local_client,
node_key,
file_path,
asset_id,
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,
)
queried_file = (FileQuery().with_node_key(eq=node_key).with_file_path(
eq=file_path).with_asset_id(eq=asset_id).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(local_client))
assert node_key == queried_file.node_key
assert file_path == queried_file.get_file_path()
assert asset_id == queried_file.get_asset_id()
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()
