diff --git a/README.md b/README.md index c83a7cb..92daf05 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,10 @@ Use tox to test: # Changelog +### V2.2.0 +- Support for loading Graphs with special relationships (Groups, Intelligence, Livehunt, Retrohunt, Commonalities). +- New method for creating groups of nodes. + ### V2.1.0 - Support for setting Graph representation. diff --git a/tests/resources/virustotal_graph_id.json b/tests/resources/virustotal_graph_id.json index 13b1d6a..b252de9 100644 --- a/tests/resources/virustotal_graph_id.json +++ b/tests/resources/virustotal_graph_id.json @@ -1,186 +1,219 @@ { "data": { - "attributes": { - "comments_count": 0, - "creation_date": 1567094335, - "graph_data": { - "description": "First Graph API test", - "version": "api-1.0.0" + "attributes": { + "comments_count": 0, + "creation_date": 1567094335, + "graph_data": { + "description": "First Graph API test", + "version": "api-1.0.0" + }, + "last_modified_date": 1567094335, + "links": [ + { + "connection_type": "contacted_ips", + "source": "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "contacted_ips", + "source": "efa0b414a831cbf724d1c67808b7483dec22a981ae670947793d114048f88057", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "contacted_ips", + "source": "720d6a4288fa43357151bdeb8dc9cdb7c27fd7db1b5f76345f5ff094d48ae5a0", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "contacted_ips", + "source": "b20ce00a6864225f05de6407fac80ddb83cd0aec00ada438c1e354cdd0d7d5df", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "contacted_ips", + "source": "5961861d2b9f50d05055814e6bfd1c6291b30719f8a4d02d4cf80c2e87753fa1", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "contacted_ips", + "source": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", + "target": "178.62.125.244" + }, + { + "connection_type": "communicating_files", + "source": "178.62.125.244", + "target": "relationships_communicating_files_17862125244" + }, + { + "connection_type": "communicating_files", + "source": "relationships_communicating_files_17862125244", + "target": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f88e" + }, + { + "connection_type": "communicating_files", + "source": "relationships_communicating_files_17862125244", + "target": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885" + }, + { + "connection_type": "commonality", + "source": "relationships_commonality_1670398662", + "target": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885" + }, + { + "connection_type": "retrohunt", + "source": "relationships_retrohunt_user1626193170", + "target": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885" + }, + { + "connection_type": "contacted_ips", + "source": "relationships_group_123456789", + "target": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + }, + { + "connection_type": "group", + "source": "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", + "target": "relationships_group_123456789" + } + ], + "nodes": [ + { + "entity_attributes": { + "has_detections": 45, + "type_tag": "docx" }, - "last_modified_date": 1567094335, - "links": [ - { - "connection_type": "contacted_ips", - "source": - "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", - "target": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" - }, - { - "connection_type": "contacted_ips", - "source": - "efa0b414a831cbf724d1c67808b7483dec22a981ae670947793d114048f88057", - "target": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" - }, - { - "connection_type": "contacted_ips", - "source": - "720d6a4288fa43357151bdeb8dc9cdb7c27fd7db1b5f76345f5ff094d48ae5a0", - "target": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" - }, - { - "connection_type": "contacted_ips", - "source": - "b20ce00a6864225f05de6407fac80ddb83cd0aec00ada438c1e354cdd0d7d5df", - "target": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" - }, - { - "connection_type": "contacted_ips", - "source": - "5961861d2b9f50d05055814e6bfd1c6291b30719f8a4d02d4cf80c2e87753fa1", - "target": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" - }, - { - "connection_type": "contacted_ips", - "source": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", - "target": "178.62.125.244" - }, - { - "connection_type": "communicating_files", - "source": "178.62.125.244", - "target": - "relationships_communicating_files_17862125244" - }, - { - "connection_type": "communicating_files", - "source": - "relationships_communicating_files_17862125244", - "target": - "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f88e" - }, - { - "connection_type": "communicating_files", - "source": - "relationships_communicating_files_17862125244", - "target": - "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885" - } - ], - "nodes": [ - { - "entity_attributes": { - "has_detections": 45, - "type_tag": "docx" - }, - "entity_id": - "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", - "index": 0, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "country": "GB" - }, - "entity_id": "178.62.125.244", - "index": 1, - "type": "ip_address", - "x": 0, - "y": 0 - }, - { - "entity_id": - "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", - "index": 2, - "type": "relationship", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "has_detections": 51, - "type_tag": "peexe" - }, - "entity_id": - "efa0b414a831cbf724d1c67808b7483dec22a981ae670947793d114048f88057", - "index": 3, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "has_detections": 55, - "type_tag": "peexe" - }, - "entity_id": - "720d6a4288fa43357151bdeb8dc9cdb7c27fd7db1b5f76345f5ff094d48ae5a0", - "index": 4, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "has_detections": 52, - "type_tag": "peexe" - }, - "entity_id": - "b20ce00a6864225f05de6407fac80ddb83cd0aec00ada438c1e354cdd0d7d5df", - "index": 5, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "has_detections": 59, - "type_tag": "peexe" - }, - "entity_id": - "5961861d2b9f50d05055814e6bfd1c6291b30719f8a4d02d4cf80c2e87753fa1", - "index": 6, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_attributes": { - "has_detections": 57, - "type_tag": "peexe" - }, - "entity_id": - "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f88e", - "index": 7, - "type": "file", - "x": 0, - "y": 0 - }, - { - "entity_id": - "relationships_communicating_files_17862125244", - "index": 8, - "type": "relationship", - "x": 0, - "y": 0 - }, - { - "entity_id": - "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885", - "index": 9, - "type": "file", - "x": 0, - "y": 0 + "entity_id": "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", + "index": 0, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "country": "GB" + }, + "entity_id": "178.62.125.244", + "index": 1, + "type": "ip_address", + "x": 0, + "y": 0 + }, + { + "entity_id": "relationships_contacted_ips_5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91", + "index": 2, + "type": "relationship", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "has_detections": 51, + "type_tag": "peexe" + }, + "entity_id": "efa0b414a831cbf724d1c67808b7483dec22a981ae670947793d114048f88057", + "index": 3, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "has_detections": 55, + "type_tag": "peexe" + }, + "entity_id": "720d6a4288fa43357151bdeb8dc9cdb7c27fd7db1b5f76345f5ff094d48ae5a0", + "index": 4, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "has_detections": 52, + "type_tag": "peexe" + }, + "entity_id": "b20ce00a6864225f05de6407fac80ddb83cd0aec00ada438c1e354cdd0d7d5df", + "index": 5, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "has_detections": 59, + "type_tag": "peexe" + }, + "entity_id": "5961861d2b9f50d05055814e6bfd1c6291b30719f8a4d02d4cf80c2e87753fa1", + "index": 6, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_attributes": { + "has_detections": 57, + "type_tag": "peexe" + }, + "entity_id": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f88e", + "index": 7, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_id": "relationships_communicating_files_17862125244", + "index": 8, + "type": "relationship", + "x": 0, + "y": 0 + }, + { + "entity_id": "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885", + "index": 9, + "type": "file", + "x": 0, + "y": 0 + }, + { + "entity_id": "relationships_commonality_1670398662", + "index": 10, + "type": "relationship", + "entity_attributes": { + "commonalities": [ + { + "commonality": "1670398662" } - ], - "private": true, - "views_count": 6 - } + ] + }, + "x": 0, + "y": 0 + }, + { + "entity_id": "relationships_retrohunt_user1626193170", + "index": 11, + "type": "relationship", + "entity_attributes": { + "relationship_type": "retrohunt", + "retrohunt_job_id": "user-1626193170" + }, + "x": 0, + "y": 0 + }, + { + "entity_id": "relationships_group_123456789", + "index": 12, + "type": "relationship", + "entity_attributes": { + "relationship_type": "group", + "grouped_node_ids": [ + "5504e04083d6146a67cb0d671d8ad5885315062c9ee08a62e40e264c2d5eab91" + ] + }, + "x": 0, + "y": 0 + } + ], + "private": true, + "views_count": 6 + } } } diff --git a/tests/test_create_group.py b/tests/test_create_group.py new file mode 100644 index 0000000..25dfc73 --- /dev/null +++ b/tests/test_create_group.py @@ -0,0 +1,73 @@ +"""Test create a Group of nodes.""" + +import pytest +import vt_graph_api +import vt_graph_api.errors + + +def create_dummy_graph(): + return vt_graph_api.VTGraph( + "Dummy api key", verbose=False, private=False, name="Graph test", + user_editors=["dummy_user"], group_viewers=["virustotal"]) + + +def test_create_empty_group(): + """Test create a group without nodes.""" + test_graph = create_dummy_graph() + test_graph.add_node("virustotal.com", "domain") + test_graph.add_node("google.com", "domain") + + with pytest.raises(vt_graph_api.errors.CreateGroupError, + match=r"A group must contain at least one node."): + test_graph.create_group([], "Group 1") + + +def test_create_group_with_nodes_already_grouped(): + """Test create a group with nodes already grouped.""" + test_graph = create_dummy_graph() + test_graph.add_node("virustotal.com", "domain") + test_graph.add_node("google.com", "domain") + test_graph.create_group(['virustotal.com', 'google.com'], 'Group 1') + + with pytest.raises(vt_graph_api.errors.CreateGroupError, + match=r"Node .+ is already in a group."): + test_graph.create_group(['virustotal.com', 'google.com'], "Group 1") + + +def test_create_group_with_node_that_does_not_exist(): + """Test create a group with nodes that are not in the graph.""" + test_graph = create_dummy_graph() + test_graph.add_node("virustotal.com", "domain") + test_graph.add_node("google.com", "domain") + test_graph.create_group(['virustotal.com', 'google.com'], 'Group 1') + + with pytest.raises(vt_graph_api.errors.CreateGroupError, + match=r"Node hola.es is not in the Graph."): + test_graph.create_group(['hola.es'], "Group 1") + + +def test_create_group(mocker): + """Test create a group.""" + test_graph = create_dummy_graph() + + test_graph.add_node("virustotal.com", "domain") + test_graph.add_node("google.com", "domain") + test_graph.create_group(['virustotal.com', 'google.com'], 'Group 1') + + mocker.patch.object(test_graph, "_push_editors") + mocker.patch.object(test_graph, "_push_viewers") + event_mocked = mocker.patch.object(test_graph, "_push_graph_to_vt") + + test_graph.save_graph() + + # Assert group relationship node is generated + group_node = event_mocked.call_args[0][0]['data']['attributes']['nodes'][-1] + assert set(group_node['entity_attributes']['grouped_node_ids']) == { + 'virustotal.com', 'google.com'} + assert len(group_node['entity_attributes']['grouped_node_ids']) == 2 + + # Assert group relationship links are generated + links = event_mocked.call_args[0][0]['data']['attributes']['links'] + group_links = [link for link in links if link['connection_type'] == 'group'] + assert len(group_links) == 2 + mocker.resetall() diff --git a/tests/test_load_graph.py b/tests/test_load_graph.py index 5094e2b..6786147 100644 --- a/tests/test_load_graph.py +++ b/tests/test_load_graph.py @@ -1,13 +1,12 @@ """Test load graph from VT.""" - import json import os import pytest +import unittest import vt_graph_api.errors import vt_graph_api.graph - with ( open(os.path.join( os.path.dirname(os.path.abspath(__file__)), @@ -36,7 +35,6 @@ "dummy": "dummy_value" } - API_KEY = "DUMMY_API_KEY" GRAPH_ID = "DUMMY_ID" @@ -104,6 +102,34 @@ def test_load_graph_with_match(mocker): assert test_graph.links[(source, target, connection_type)] assert "virustotal" in test_graph.group_editors assert "alvarogf" in test_graph.user_viewers + + special_relationship_nodes = [ + "relationships_commonality_1670398662", + "relationships_retrohunt_user1626193170" + ] + + special_relationship_links = [ + ("relationships_commonality_1670398662", + "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885", + "commonality"), + ("relationships_retrohunt_user1626193170", + "e6ecb146f469d243945ad8a5451ba1129c5b190f7d50c64580dbad4b8246f885", + "retrohunt") + ] + + for node in special_relationship_nodes: + assert test_graph.special_relationship_nodes[node] + + for source, target, connection_type in special_relationship_links: + link = {'source': source, 'target': target, + 'connection_type': connection_type} + assert link in test_graph.special_relationship_links + + group_nodes = ["relationships_group_123456789"] + + for node in group_nodes: + assert test_graph.group_nodes[node] + mocker.resetall() @@ -190,3 +216,7 @@ def test_load_graph_wrong_json(mocker): mocker.patch("requests.get", return_value=m) vt_graph_api.graph.VTGraph.load_graph(GRAPH_ID, API_KEY) mocker.resetall() + + +def test_load_graph_with_group_nodes(mocker): + mocker.resetall() diff --git a/vt_graph_api/errors.py b/vt_graph_api/errors.py index 5a8fea7..539ac90 100644 --- a/vt_graph_api/errors.py +++ b/vt_graph_api/errors.py @@ -49,3 +49,6 @@ class DownloadScreenshotError(Exception): class CreateCollectionError(Exception): pass + +class CreateGroupError(Exception): + pass diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py index 4c1deb3..b17f97d 100644 --- a/vt_graph_api/graph.py +++ b/vt_graph_api/graph.py @@ -7,7 +7,6 @@ VT API: https://virustotal.github.io/vt-graph-api/ """ - import collections import functools import json @@ -17,6 +16,7 @@ from enum import Enum +import base64 import concurrent.futures import requests import six @@ -113,20 +113,20 @@ def load_graph(graph_id, api_key): # Get graph data. graph_data_url = ( "https://www.virustotal.com/api/v3/graphs/{graph_id}" - .format(graph_id=graph_id)) + .format(graph_id=graph_id)) graph_data_response = requests.get(graph_data_url, headers=headers) if graph_data_response.status_code != 200: raise vt_graph_api.errors.LoadError( ("Error to find graph with id: {graph_id}. Response code: " + "{status_code}.").format( - graph_id=graph_id, status_code=graph_data_response.status_code)) + graph_id=graph_id, status_code=graph_data_response.status_code)) try: data = graph_data_response.json() except json.JSONDecodeError: raise vt_graph_api.errors.LoadError( "Malformed JSON response: {json_response}" - .format(json_response=graph_data_response.text)) + .format(json_response=graph_data_response.text)) try: graph_name = data["data"]["attributes"]["graph_data"]["description"] @@ -136,7 +136,7 @@ def load_graph(graph_id, api_key): except KeyError as e: raise vt_graph_api.errors.InvalidJSONError( "Unexpected error in json structure at load_graph: {msg}." - .format(msg=str(e))) + .format(msg=str(e))) # Creates empty graph. graph = vt_graph_api.graph.VTGraph( @@ -145,15 +145,22 @@ def load_graph(graph_id, api_key): # Adds users/group viewers and editors. graph._pull_viewers() graph._pull_editors() + # Adds group nodes from json data + graph._add_group_nodes_from_json_data(nodes) + # Adds intelligence nodes from json_data + graph._add_special_relationship_nodes_from_json_data(nodes) + # Adds intelligence links from json data + graph._add_special_links_from_json_data(links) # Adds nodes to the graph. graph._add_nodes_from_json_graph_data(nodes) # Adds links to the graph. graph._add_links_from_json_graph_data(links) + return graph @staticmethod def clone_graph(graph_id, api_key, name="", private=False, user_editors=None, - user_viewers=None, group_editors=None, group_viewers=None): + user_viewers=None, group_editors=None, group_viewers=None): """Clone VirusTotal Graph and make it yours according the given parameters. Args: @@ -190,8 +197,8 @@ def clone_graph(graph_id, api_key, name="", private=False, user_editors=None, return graph def __init__(self, api_key, name="", private=False, user_editors=None, - user_viewers=None, group_editors=None, group_viewers=None, - verbose=False): + user_viewers=None, group_editors=None, group_viewers=None, + verbose=False): """Creates a VT Graph Instance. Args: @@ -228,6 +235,9 @@ def __init__(self, api_key, name="", private=False, user_editors=None, self.nodes = {} self.links = {} + self.group_nodes = {} + self.special_relationship_nodes = {} + self.special_relationship_links = {} self._id_references = {} self._api_calls_lock = threading.Lock() @@ -395,20 +405,30 @@ def _add_links_from_json_graph_data(self, json_graph_data_links): try: # It is necessary to clean the given links because they have relationship # nodes + + non_special_relationship_links = ( + link for link in json_graph_data_links if + link['source'] not in self.group_nodes and link[ + 'source'] not in self.special_relationship_nodes and link[ + 'target'] not in self.special_relationship_nodes and link[ + 'target'] not in self.group_nodes + ) + + relationship_links = [] + non_relationship_links = [] + for link in non_special_relationship_links: + if link["source"].startswith("relationship"): + relationship_links.append(link) + else: + non_relationship_links.append(link) + replace_nodes = {} - relationship_links = ( - link_ for link_ in json_graph_data_links - if link_["source"].startswith("relationship")) for link in relationship_links: if link["source"] not in replace_nodes: replace_nodes[link["source"]] = [link["target"]] else: replace_nodes[link["source"]].append(link["target"]) - non_relationship_links = ( - link for link in json_graph_data_links - if not link["source"].startswith("relationship")) - for link_data in non_relationship_links: linked_nodes = replace_nodes.get( link_data["target"], [link_data["target"]]) @@ -421,6 +441,87 @@ def _add_links_from_json_graph_data(self, json_graph_data_links): "field in the VT response." ) + def _add_group_nodes_from_json_data(self, json_graph_data_nodes): + """Add all the group nodes from the given data. + + json_graph_data_nodes are the responses from querying VT. + + Args: + json_graph_data_nodes ([dict]): list of node's data with the following + structure => { + "entity_attributes": "", + "entity_id": "", + "index": "", + "type": "", + "x": "", + "y": "" + } + """ + + self.group_nodes = {node['entity_id']: node for node in + json_graph_data_nodes + if node.get("type") == "relationship" and node.get( + "entity_attributes", {}).get("relationship_type") == "group"} + + def _is_special_relationship_node(self, node): + """Checks if a node is a special relationship node. + + Args: + node: node to be checked. + + Returns: boolean + + """ + is_relationship_node = node.get("type") == "relationship" + entity_attributes = node.get("entity_attributes", {}) + relationship_type = entity_attributes.get("relationship_type") + is_commonalities = entity_attributes.get("commonalities") is not None + + return is_relationship_node and ( + relationship_type and relationship_type != "group") or is_commonalities + + def _add_special_relationship_nodes_from_json_data(self, + json_graph_data_nodes): + """Add all the special relationship nodes from the given data. + + json_graph_data_nodes are the responses from querying VT. + + Args: + json_graph_data_nodes ([dict]): list of node's data with the following + structure => { + "entity_attributes": "", + "entity_id": "", + "index": "", + "type": "", + "x": "", + "y": "" + } + """ + self.special_relationship_nodes = dict([ + (node['entity_id'], node) for node in json_graph_data_nodes if + self._is_special_relationship_node(node) + ]) + + def _add_special_links_from_json_data(self, json_graph_data_links): + """Add all the special relationship links from the given data. + json_graph_data_nodes are the responses from querying VT. + + Args: + json_graph_data_nodes ([dict]): list of node's data with the following + structure => { + "entity_attributes": "", + "entity_id": "", + "index": "", + "type": "", + "x": "", + "y": "" + } + """ + self.special_relationship_links = [ + link for link in json_graph_data_links if + link['source'] in self.special_relationship_nodes or + link['target'] in self.special_relationship_nodes] + def _pull_viewers(self): """Pull graph's users and groups viewers from VT API. @@ -455,7 +556,7 @@ def _pull_viewers(self): except KeyError as e: raise vt_graph_api.errors.InvalidJSONError( "Unexpected error in json structure at get_graph_viewers: {msg}" - .format(msg=str(e))) + .format(msg=str(e))) self.user_viewers.extend(user_viewers) self.group_viewers.extend(group_viewers) @@ -522,7 +623,7 @@ def _pull_editors(self): except KeyError as e: raise vt_graph_api.errors.InvalidJSONError( "Unexpected error in json structure at get_graph_editors: {msg}" - .format(msg=str(e))) + .format(msg=str(e))) self.user_editors.extend(user_editors) self.group_editors.extend(group_editors) @@ -578,10 +679,10 @@ def _push_graph_to_vt(self, output): if response.status_code != 200: self._log( "Saving graph error: {status_code} status code." - .format(status_code=response.status_code)) + .format(status_code=response.status_code)) raise vt_graph_api.errors.SaveGraphError( "Saving graph error: {status_code} status code." - .format(status_code=response.status_code) + .format(status_code=response.status_code) ) data = response.json() @@ -607,7 +708,7 @@ def _fetch_node_information(self, node): if response.status_code != 200: self._log( "Request to '{url}' with '{status_code}' status code" - .format(url=url, status_code=response.status_code) + .format(url=url, status_code=response.status_code) ) return @@ -638,7 +739,7 @@ def _compute_common_relationship_ids(self): # us the common expansions shared_expansions = ( set(node.children) - .intersection(set(node_.children))) + .intersection(set(node_.children))) # Two nodes could be minimized if they have the same children in the # same expansion and they have at least one child. for expansion in shared_expansions: @@ -797,7 +898,7 @@ def _get_node_id(self, node_id, fetch_vt_enterprise=False): return valid_node_id def _query_expansion_nodes(self, node, expansion, - max_nodes_per_relationship, cursor, max_retries): + max_nodes_per_relationship, cursor, max_retries): """Get expansion nodes JSON data by querying VirusTotal API. Args: @@ -834,7 +935,7 @@ def _query_expansion_nodes(self, node, expansion, try: self._log( "Expanding node {node_id} with expansion {expansion}" - .format(node_id=node.node_id, expansion=expansion)) + .format(node_id=node.node_id, expansion=expansion)) self._increment_api_counter() response = requests.get( url, headers=self._get_headers(), timeout=self.REQUEST_TIMEOUT) @@ -848,9 +949,9 @@ def _query_expansion_nodes(self, node, expansion, return data def _get_expansion_nodes(self, node, expansion, - max_nodes_per_relationship=1000, cursor=None, - max_retries=3, expansion_nodes=None, - consumed_quotas=0): + max_nodes_per_relationship=1000, cursor=None, + max_retries=3, expansion_nodes=None, + consumed_quotas=0): """Returns the nodes to be attached to the given node with the given expansion. Args: @@ -927,7 +1028,7 @@ def _get_expansion_nodes(self, node, expansion, return expansion_nodes, consumed_quotas def _parallel_expansion(self, target_nodes, solution_paths, visited_nodes, - max_api_quotas, lock, max_depth, node, params): + max_api_quotas, lock, max_depth, node, params): """Parallelize the node expansion synchronizing the api quotas consumed. Args: @@ -1010,7 +1111,7 @@ def _parallel_expansion(self, target_nodes, solution_paths, visited_nodes, return expansion_nodes def _search_connection(self, source_node, target_nodes, - max_api_quotas, max_depth, max_qps): + max_api_quotas, max_depth, max_qps): """Search connection between the node source and all of the target_nodes. source_node @@ -1082,8 +1183,8 @@ def _search_connection(self, source_node, target_nodes, return paths def _resolve_relations(self, source_node, target_nodes, - max_api_quotas, max_depth, max_qps, - fetch_info_collected_nodes): + max_api_quotas, max_depth, max_qps, + fetch_info_collected_nodes): """Try to connect the source_node with all of the nodes in target_nodes. Args: @@ -1129,8 +1230,8 @@ def _resolve_relations(self, source_node, target_nodes, return has_link def add_node(self, node_id, node_type, fetch_information=True, - fetch_vt_enterprise=True, label="", node_attributes=None, - x=0, y=0): + fetch_vt_enterprise=True, label="", node_attributes=None, + x=0, y=0): """Adds a node with id `node_id` of `node_type` type to the graph. Args: @@ -1182,7 +1283,7 @@ def add_node(self, node_id, node_type, fetch_information=True, return node def add_nodes(self, node_list, fetch_information=True, - fetch_vt_enterprise=True): + fetch_vt_enterprise=True): """Adds the node_list to the graph concurrently. Args: @@ -1265,7 +1366,7 @@ def add_link(self, source_node, target_node, connection_type=""): if source_node == target_node: raise vt_graph_api.errors.SameNodeError( "It is no possible to add links between the same node; id: {node_id}." - .format(node_id=source_node)) + .format(node_id=source_node)) source_node = self._get_node_id(source_node) target_node = self._get_node_id(target_node) @@ -1273,11 +1374,11 @@ def add_link(self, source_node, target_node, connection_type=""): if source_node not in self.nodes: raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=source_node)) + .format(node_id=source_node)) if target_node not in self.nodes: raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=target_node)) + .format(node_id=target_node)) if connection_type not in self.nodes[source_node].expansions_available: self._log("Expansion `{expansion_type}` is not standard expansion type", logging.WARNING) @@ -1287,8 +1388,8 @@ def add_link(self, source_node, target_node, connection_type=""): self.nodes[source_node].add_child(target_node, connection_type) def add_links_if_match(self, source_node, target_node, - max_api_quotas=100000, max_depth=3, max_qps=1000, - fetch_info_collected_nodes=True): + max_api_quotas=100000, max_depth=3, max_qps=1000, + fetch_info_collected_nodes=True): """Try to find a relationship between the source_node and the target_node. Adds the needed links between the source_node and the target_node if @@ -1322,7 +1423,7 @@ def add_links_if_match(self, source_node, target_node, if source_node == target_node: raise vt_graph_api.errors.SameNodeError( "It is no possible to add links between the same node; id: {node_id}." - .format(node_id=source_node)) + .format(node_id=source_node)) quotas_before_get_id = self.get_api_calls() source_node = self._get_node_id(source_node) @@ -1349,8 +1450,8 @@ def add_links_if_match(self, source_node, target_node, max_depth, max_qps, fetch_info_collected_nodes) def connect_with_graph(self, source_node, max_api_quotas=100000, - max_depth=3, max_qps=1000, - fetch_info_collected_nodes=True): + max_depth=3, max_qps=1000, + fetch_info_collected_nodes=True): """Try to connect the source_node with the current graph nodes. Args: @@ -1422,18 +1523,18 @@ def delete_link(self, source_node, target_node, connection_type): if source_node not in self.nodes: raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=source_node)) + .format(node_id=source_node)) if target_node not in self.nodes: raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=target_node)) + .format(node_id=target_node)) if (source_node, target_node, connection_type) not in self.links: raise vt_graph_api.errors.LinkNotFoundError( ("Link between {source} and {target} with {connection_type} does " + "not exists.").format( - source=source_node, target=target_node, - connection_type=connection_type)) + source=source_node, target=target_node, + connection_type=connection_type)) del self.links[(source_node, target_node, connection_type)] self.nodes[source_node].delete_child(target_node, connection_type) @@ -1451,7 +1552,7 @@ def delete_links(self, node_id): if node_id not in self.nodes: raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=node_id)) + .format(node_id=node_id)) links_to_be_deleted = [ link for link in six.iterkeys(self.links) @@ -1487,13 +1588,13 @@ def expand(self, node_id, expansion, max_nodes_per_relationship=40): if node_id not in six.iterkeys(self.nodes): raise vt_graph_api.errors.NodeNotFoundError( "Node '{node_id}' not found in nodes." - .format(node_id=node_id)) + .format(node_id=node_id)) node = self.nodes[node_id] if expansion not in node.expansions_available: raise vt_graph_api.errors.NodeNotSupportedExpansionError( "Node {node_id} cannot be expanded with {expansion} expansion." - .format(node_id=node_id, expansion=expansion)) + .format(node_id=node_id, expansion=expansion)) expansion_nodes, _ = self._get_expansion_nodes( node, expansion, max_nodes_per_relationship) @@ -1551,7 +1652,7 @@ def expand_one_level(self, node_id, max_nodes_per_relationship=40): return expansion_nodes def expand_n_level(self, level=1, max_nodes_per_relationship=40, - max_nodes=10000): + max_nodes=10000): """Expands all the nodes in the graph `level` levels. For example: @@ -1666,6 +1767,24 @@ def save_graph(self): self._add_node_to_output(output, node_id) added.add(node_id) + special_relationship_nodes = self._get_special_relationship_nodes( + len(added)) + for node in special_relationship_nodes: + output["data"]["attributes"]["nodes"].append(node) + added.add(node['entity_id']) + + special_relationship_links = self._get_special_relationship_links(added) + output["data"]["attributes"]["links"] += special_relationship_links + + group_nodes = self._get_groups_nodes(added) + for node in group_nodes: + output["data"]["attributes"]["nodes"].append(node) + added.add(node['entity_id']) + + final_links = output["data"]["attributes"]["links"] + group_links = self._get_groups_links(final_links) + output["data"]["attributes"]["links"] += group_links + self._push_graph_to_vt(output) self._push_editors() self._push_viewers() @@ -1713,7 +1832,7 @@ def get_iframe_code(self): "{graph_id}\" width=\"800\" height=\"600\">" .format(graph_id=self.graph_id)) - def download_screenshot(self, path = "."): + def download_screenshot(self, path="."): """Downloads a screenshot of the graph. Args: @@ -1741,7 +1860,7 @@ def download_screenshot(self, path = "."): r.raw.decode_content = True filename = "{graph_id}.jpg".format(graph_id=self.graph_id) file_path = os.path.join(path, filename) - with open(file_path,'wb') as f: + with open(file_path, 'wb') as f: shutil.copyfileobj(r.raw, f) else: raise vt_graph_api.errors.DownloadScreenshotError( @@ -1756,6 +1875,98 @@ def _get_nodes_by_type(self, node_type): for node in self.nodes.values() if node.node_type == node_type ]} + def _get_groups_nodes(self, final_nodes): + """This method must be called only just before saving the graph. + It returns all the loaded and generated group nodes. It checks that + the nodes that belong to the group have been added to the graph, and deletes + the ids that have not finally been added. + + In case a Group Node doesn't have any belonging node in the Graph, it is not + returned. + + Args: + final_nodes: Set with the ids of all the added nodes to the Graph. + + Returns: Array of group nodes. + + """ + group_nodes = [] + next_index = len(final_nodes) + # Delete nodes belonging to group nodes that have not been added to the Graph. + # Inject Node index. + for idx, group_node in enumerate(self.group_nodes.values()): + belonging_nodes = group_node['entity_attributes']['grouped_node_ids'] + belonging_nodes = list( + set(node for node in belonging_nodes if node in final_nodes)) + if belonging_nodes: + group_node['entity_attributes']['grouped_node_ids'] = belonging_nodes + group_node['index'] = next_index + next_index += 1 + group_nodes.append(group_node) + return group_nodes + + def _get_groups_links(self, final_links): + """Generates all the links needed to represent groups in VTGraph. + + Args: + final_links: list with all the links added in the Graph. + + Returns: List with the links needed to represent groups in VTGraph. + + """ + map_node_group = {} + + # Generate a link for each pair node_in_group -> group_node + for group_node in self.group_nodes.values(): + for node_in_group in group_node['entity_attributes']['grouped_node_ids']: + map_node_group[node_in_group] = group_node['entity_id'] + + group_links = [{'source': node, 'target': group, 'connection_type': 'group'} + for node, group in map_node_group.items()] + + # Generate a link for each link where source or target is a node + # inside a group. + for link in final_links: + if link['source'] in map_node_group: + new_link = link.copy() + new_link['source'] = map_node_group[link['source']] + group_links.append(new_link) + elif link['target'] in map_node_group: + new_link = link.copy() + new_link['target'] = map_node_group[link['target']] + group_links.append(new_link) + return group_links + + def _get_special_relationship_nodes(self, next_index): + """ + Returns a list with the special relationship nodes, and injects nodes + index. Only used for saving the Graph. + Args: + next_index: Next index to be injected. + + Returns: + + """ + special_relationship_nodes = [] + for idx, node in self.special_relationship_nodes.values(): + node['index'] = next_index + idx + special_relationship_nodes.append(node) + return special_relationship_nodes + + def _get_special_relationship_links(self, final_nodes): + """Returns links needed to represent special relationship nodes. + It checks if the link connections exist or not. Only used for saving the + Graph. + + Args: + final_nodes: Set with all the nodes added + + Returns: list of links related to special_relationship_nodes. + + """ + return [link for link in self.special_relationship_links if + link['source'] in final_nodes and link['target'] in final_nodes] + def create_collection(self, name=None, description=None): """Creates a VT Collection taking entities from current Graph. @@ -1774,13 +1985,13 @@ def create_collection(self, name=None, description=None): data = { "type": "collection", "attributes": { - "name": name if name else "Collection created from VT Graph API", + "name": name if name else "Collection created from VT Graph API", }, "relationships": { - "files": self._get_nodes_by_type('file'), - "domains": self._get_nodes_by_type('domain'), - "urls": self._get_nodes_by_type('url'), - "ip_addresses": self._get_nodes_by_type('ip_address') + "files": self._get_nodes_by_type('file'), + "domains": self._get_nodes_by_type('domain'), + "urls": self._get_nodes_by_type('url'), + "ip_addresses": self._get_nodes_by_type('ip_address') } } @@ -1795,7 +2006,7 @@ def create_collection(self, name=None, description=None): response = requests.post( url, headers=self._get_headers(), json={"data": data}) - if(response.status_code != 200): + if (response.status_code != 200): print(response.json()) raise vt_graph_api.errors.CreateCollectionError() @@ -1805,17 +2016,64 @@ def create_collection(self, name=None, description=None): collection_id=collection_id ) - def set_representation(self, representation): - """Sets Graph representation. + def _generate_group_node_id(self, nodes_id): + return ",".join(nodes_id).replace("=", "") - Args: - representation: Graph representation. See :py:class::`RepresentationType`. + def create_group(self, node_ids, group_name): """ - self.representation = representation + Args: + node_ids: + group_name: + Returns: + Raises: CreateGroupError if: + - Group doesn't contain any node. + - Group contains a node that is already in another group. + - Group contains a node that is not in the Graph. + """ + group_node_id = self._generate_group_node_id(node_ids) + node_ids_set = set(node_ids) + + # Check if the user has provided node ids + if not node_ids: + raise vt_graph_api.errors.CreateGroupError( + "A group must contain at least one node.") + + # Check if all the nodes exists. + for node_id in node_ids_set: + if node_id not in self.nodes: + raise vt_graph_api.errors.CreateGroupError( + "Node {node_id} is not in the Graph.".format( + node_id=node_id)) + + # Check if nodes are already in a group. + for group_node in self.group_nodes.values(): + for node_id in group_node["entity_attributes"]["grouped_node_ids"]: + if node_id in node_ids_set: + raise vt_graph_api.errors.CreateGroupError( + "Node {node_id} is already in a group.".format( + node_id=node_id)) + + relationship_node = { + "entity_id": "relationships_group_{group_id}".format( + group_id=group_node_id), + "type": 'relationship', + "entity_attributes": { + "grouped_node_ids": node_ids, + "relationship_type": "group" + }, + "text": group_name + } + self.group_nodes[group_node_id] = relationship_node + def set_representation(self, representation): + """Sets Graph representation. + Args: + representation: Graph representation. See :py:class::`RepresentationType`. + """ + self.representation = representation diff --git a/vt_graph_api/node.py b/vt_graph_api/node.py index c27e846..d254b53 100644 --- a/vt_graph_api/node.py +++ b/vt_graph_api/node.py @@ -41,7 +41,7 @@ class Node(object): "whois", "ssl_cert", "collection", - "reference" + "reference", ) NODE_EXPANSIONS = { "file": [ diff --git a/vt_graph_api/version.py b/vt_graph_api/version.py index e1fe331..a8dbc9a 100644 --- a/vt_graph_api/version.py +++ b/vt_graph_api/version.py @@ -4,5 +4,5 @@ """ -__version__ = '2.1.0' +__version__ = '2.2.0' __x_tool__ = 'Graph'