def test_get_expansion_nodes_n_level_without_cursor(mocker):
"""Test get expansion nodes at n level without cursor."""
rq_id = "7c11c7ccd384fd9f377da499fc059fa08fdc33a1bb870b5bc3812d24dd421a16"
request_data = {
"data": [
{
"attributes": {},
"id": rq_id,
"type": "file"
}
]
}
mocker.spy(test_graph, "_get_expansion_nodes")
node_a = vt_graph_api.Node(
"26c808a1eb3eaa7bb29ec2ab834559f06f2636b87d5f542223426d6f238ff906",
"file")
node_b = vt_graph_api.Node(
"7c11c7ccd384fd9f377da499fc059fa08fdc33a1bb870b5bc3812d24dd421a16",
"file")
m = mocker.Mock(status_code=200, json=mocker.Mock(return_value=request_data))
mocker.patch("requests.get", return_value=m)
expansion_nodes, _ = test_graph._get_expansion_nodes(
node_a, "similar_files", 1000)
assert test_graph._get_expansion_nodes.call_count == 1
assert node_b in expansion_nodes
mocker.resetall()
def test_search_connection_third_level(mocker):
"""Test search connection and found it in the third level."""
rq_id = "7c11c7ccd384fd9f377da499fc059fa08fdc33a1bb870b5bc3812d24dd421a16"
rq_id_2 = "660903b139d5c7ec80af124e93320c18895de32135450d4acd14096e6c0dd2ef"
request_response_first_level = [{
"data": [{
"attributes": {},
"id": rq_id,
"type": "file"
}]
}]
request_response_second_level = [
{
"data": [{
"attributes": {},
"id": "nsis.sf.net",
"type": "domain"
}]
},
]
request_response_third_level = [
{
"data": [{
"attributes": {},
"id": rq_id_2,
"type": "file"
}]
},
]
side_effects = list(request_response_first_level *
len(vt_graph_api.Node.NODE_EXPANSIONS["file"]))
side_effects += (request_response_second_level *
len(vt_graph_api.Node.NODE_EXPANSIONS["file"]))
side_effects += (request_response_third_level *
len(vt_graph_api.Node.NODE_EXPANSIONS["domain"]))
node_a = vt_graph_api.Node(
"26c808a1eb3eaa7bb29ec2ab834559f06f2636b87d5f542223426d6f238ff906",
"file")
node_b = vt_graph_api.Node(
"660903b139d5c7ec80af124e93320c18895de32135450d4acd14096e6c0dd2ef",
"file")
m = mocker.Mock(status_code=200,
json=mocker.Mock(side_effect=side_effects))
mocker.patch("requests.get", return_value=m)
mocker.spy(test_graph, "_get_expansion_nodes")
mocker.spy(test_graph, "_parallel_expansion")
assert test_graph._search_connection(node_a, [node_b], 3000, 5, 1000)
assert test_graph._get_expansion_nodes.call_count == len(side_effects)
# 3 is the number of distinct nodes that the algorithm will explore
assert test_graph._parallel_expansion.call_count == 3
mocker.resetall()
def test_expansion_existing_node(mocker):
"""Test expansion existing node in graph."""
mocker.patch.object(test_graph, "_fetch_node_information")
rq_id = "fb0b6044347e972e21b6c376e37e1115dab494a2c6b9fb28b92b1e45b45d0ebc"
first_level = {
"data": [
{
"attributes": {},
"id": rq_id,
"type": "file"
}
]
}
m = mocker.Mock(status_code=200, json=mocker.Mock(return_value=first_level))
mocker.patch("requests.get", return_value=m)
added_node_id = (
"ed01ebfbc9eb5bbea545af4d01bf5f1071661840480439c6e5babe8e080e41aa")
added_node = test_graph.add_node(added_node_id, "file",
label="Investigation node")
expansion_node = vt_graph_api.Node(
"fb0b6044347e972e21b6c376e37e1115dab494a2c6b9fb28b92b1e45b45d0ebc",
"file")
assert (
[expansion_node] ==
test_graph.expand(added_node_id, "compressed_parents", 40))
assert len(test_graph.nodes) == 2
assert test_graph.nodes[expansion_node.node_id] == expansion_node
assert test_graph.links[
(added_node.node_id, expansion_node.node_id, "compressed_parents")]
mocker.resetall()
def test_search_connection_first_level(mocker):
"""Test search connection and found it in the first level."""
rq_id = "7c11c7ccd384fd9f377da499fc059fa08fdc33a1bb870b5bc3812d24dd421a16"
request_data = {"data": [{"attributes": {}, "id": rq_id, "type": "file"}]}
mocker.spy(test_graph, "_get_expansion_nodes")
mocker.spy(test_graph, "_parallel_expansion")
node_a = vt_graph_api.Node(
"26c808a1eb3eaa7bb29ec2ab834559f06f2636b87d5f542223426d6f238ff906",
"file")
node_b = vt_graph_api.Node(
"7c11c7ccd384fd9f377da499fc059fa08fdc33a1bb870b5bc3812d24dd421a16",
"file")
m = mocker.Mock(status_code=200,
json=mocker.Mock(return_value=request_data))
mocker.patch("requests.get", return_value=m)
assert test_graph._search_connection(node_a, [node_b], 1000, 5, 100)
assert test_graph._get_expansion_nodes.call_count == len(
node_a.expansions_available)
assert test_graph._parallel_expansion.call_count == 1
mocker.resetall()
def test_search_connection_second_level_real_data(mocker):
"""Test search connection end to end.
+-----------------+SOURCE_NODE+-----------------+
| + |
| +-----+---------+ |
v v v v
bundled_files carbonblack_children ... similar_files
+ + + +
+-----------+----+ +-----+-----+ | +-------------+
| | | | | | | | |
v v v v v v v v v
INTERMEDIATE_NODE ... ... ... ... ... ... ... ...
+
+-----+---------------------+
| | |
v v v
... contacted_domains ...
+
|
v
TARGET_NODE
"""
node_a = vt_graph_api.Node(
"26c808a1eb3eaa7bb29ec2ab834559f06f2636b87d5f542223426d6f238ff906",
"file")
intermediate_node = vt_graph_api.Node(INTERMEDIATE_NODE_ID, "file")
node_b = vt_graph_api.Node("nsis.sf.net", "domain")
mocker.patch("requests.get", mock_request)
mocker.spy(test_graph, "_get_expansion_nodes")
mocker.spy(test_graph, "_parallel_expansion")
total_nodes_first_level = len(node_a.expansions_available)
assert test_graph._search_connection(node_a, [node_b], 1000, 5, 100)
# Check that _get_expansion_nodes was called with the correct arguments.
calls = [
call(node_a, "bundled_files", 40),
call(node_a, "carbonblack_children", 40),
call(node_a, "carbonblack_parents", 40),
call(node_a, "compressed_parents", 40),
call(node_a, "contacted_domains", 40),
call(node_a, "contacted_ips", 40),
call(node_a, "contacted_urls", 40),
call(node_a, "email_parents", 40),
call(node_a, "embedded_domains", 40),
call(node_a, "embedded_urls", 40),
call(node_a, "embedded_ips", 40),
call(node_a, "execution_parents", 40),
call(node_a, "itw_domains", 40),
call(node_a, "itw_urls", 40),
call(node_a, "overlay_parents", 40),
call(node_a, "pcap_parents", 40),
call(node_a, "pe_resource_parents", 40),
call(node_a, "similar_files", 40),
call(intermediate_node, "bundled_files", 40),
call(intermediate_node, "carbonblack_children", 40),
call(intermediate_node, "carbonblack_parents", 40),
call(intermediate_node, "compressed_parents", 40),
call(intermediate_node, "contacted_domains", 40),
call(intermediate_node, "contacted_ips", 40),
call(intermediate_node, "contacted_urls", 40),
call(intermediate_node, "email_parents", 40),
call(intermediate_node, "embedded_domains", 40),
call(intermediate_node, "embedded_urls", 40),
call(intermediate_node, "embedded_ips", 40),
call(intermediate_node, "execution_parents", 40),
call(intermediate_node, "itw_domains", 40),
call(intermediate_node, "itw_urls", 40),
call(intermediate_node, "overlay_parents", 40),
call(intermediate_node, "pcap_parents", 40),
call(intermediate_node, "pe_resource_parents", 40),
call(intermediate_node, "similar_files", 40),
]
test_graph._get_expansion_nodes.assert_has_calls(calls, any_order=True)
total_expansion_calls = 0
for node_type in six.itervalues(EXPANSION_NODES):
total_expansion_calls += len(
vt_graph_api.Node.NODE_EXPANSIONS[node_type])
# all assertions are less than instead of equal because of the difficult of
# stopping threads when solution is found.
assert test_graph._get_expansion_nodes.call_count <= total_expansion_calls
assert test_graph._parallel_expansion.call_count <= (
1 + total_nodes_first_level + 289 # max expansions in second level
)
mocker.resetall()