From b6613baa80246af85872a9259a6a4d6adb232c3f Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Tue, 26 Jul 2022 16:37:02 +0200
Subject: [PATCH 01/10] Support group nodes

---
 vt_graph_api/graph.py | 182 +++++++++++++++++++++++++++++-------------
 1 file changed, 127 insertions(+), 55 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index 4c1deb3..d3cbb4c 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
@@ -113,20 +112,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 +135,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 +144,20 @@ 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 group links from json data
+    graph._add_group_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 +194,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 +232,8 @@ def __init__(self, api_key, name="", private=False, user_editors=None,
 
     self.nodes = {}
     self.links = {}
+    self.group_nodes = {}
+    self.group_links = {}
 
     self._id_references = {}
     self._api_calls_lock = threading.Lock()
@@ -395,9 +401,14 @@ 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_group_links = [link for link in json_graph_data_links if
+                         link['source'] not in self.group_nodes and link[
+                           'target'] not in self.group_nodes]
+
       replace_nodes = {}
       relationship_links = (
-          link_ for link_ in json_graph_data_links
+          link_ for link_ in non_group_links
           if link_["source"].startswith("relationship"))
       for link in relationship_links:
         if link["source"] not in replace_nodes:
@@ -406,7 +417,7 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
           replace_nodes[link["source"]].append(link["target"])
 
       non_relationship_links = (
-          link for link in json_graph_data_links
+          link for link in non_group_links
           if not link["source"].startswith("relationship"))
 
       for link_data in non_relationship_links:
@@ -421,6 +432,37 @@ 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.
+
+    Raises:
+      InvalidJSONError: whether the API response does not have the correct
+        structure.
+
+    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 = dict([
+        (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 _add_group_links_from_json_data(self, json_graph_data_links):
+    self.group_links = [link for link in json_graph_data_links if
+                        link['source'] in self.group_nodes or link[
+                          'target'] in self.group_nodes]
+
   def _pull_viewers(self):
     """Pull graph's users and groups viewers from VT API.
 
@@ -455,7 +497,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 +564,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 +620,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 +649,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 +680,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 +839,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 +876,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 +890,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 +969,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 +1052,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 +1124,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 +1171,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 +1224,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 +1307,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 +1315,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 +1329,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 +1364,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 +1391,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 +1464,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 +1493,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 +1529,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 +1593,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 +1708,16 @@ def save_graph(self):
       self._add_node_to_output(output, node_id)
       added.add(node_id)
 
+    group_nodes = self._get_groups_nodes(added)
+    output_nodes = output["data"]["attributes"]["nodes"]
+    output["data"]["attributes"]["nodes"] = [*output_nodes, *group_nodes]
+    for node in group_nodes:
+      added.add(node['entity_id'])
+
+    group_links = self._get_groups_links(added)
+    output_links = output["data"]["attributes"]["links"]
+    output["data"]["attributes"]["links"] = [*output_links, *group_links]
+
     self._push_graph_to_vt(output)
     self._push_editors()
     self._push_viewers()
@@ -1713,7 +1765,7 @@ def get_iframe_code(self):
         "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
         .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 +1793,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 +1808,23 @@ 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):
+    group_nodes = []
+    current_index = len(final_nodes)
+    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'] = current_index + idx
+        group_nodes.append(group_node)
+    return group_nodes
+
+  def _get_groups_links(self, final_nodes):
+    return [link for link in self.group_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 +1843,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 +1864,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()
 
@@ -1814,8 +1883,11 @@ def set_representation(self, representation):
     self.representation = representation
 
 
+if __name__ == "__main__":
+  import os
 
+  graph = VTGraph.load_graph(
+      "gb49cd88b5e6c46198623e84ed6e39c0378c661fb599d443b8170e2abb789ab65",
+      api_key=os.environ["VT_API_KEY"])
 
-
-
-
+  graph.save_graph()

From e3dc0063d84963ea02e538138abad09551bf3979 Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Tue, 26 Jul 2022 17:02:33 +0200
Subject: [PATCH 02/10] Supported intelligence nodes

---
 vt_graph_api/graph.py | 58 ++++++++++++++++++++++++++++++++-----------
 vt_graph_api/node.py  |  2 +-
 2 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index d3cbb4c..49ecb41 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -148,6 +148,10 @@ def load_graph(graph_id, api_key):
     graph._add_group_nodes_from_json_data(nodes)
     # Adds group links from json data
     graph._add_group_links_from_json_data(links)
+    # Adds intelligence nodes from json_data
+    graph._add_intelligence_nodes_from_json_data(nodes)
+    # Adds intelligence links from json data
+    graph._add_intelligence_links_from_json_data(links)
     # Adds nodes to the graph.
     graph._add_nodes_from_json_graph_data(nodes)
     # Adds links to the graph.
@@ -234,6 +238,8 @@ def __init__(self, api_key, name="", private=False, user_editors=None,
     self.links = {}
     self.group_nodes = {}
     self.group_links = {}
+    self.intelligence_nodes = {}
+    self.intelligence_links = {}
 
     self._id_references = {}
     self._api_calls_lock = threading.Lock()
@@ -402,13 +408,17 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
       # It is necessary to clean the given links because they have relationship
       # nodes
 
-      non_group_links = [link for link in json_graph_data_links if
-                         link['source'] not in self.group_nodes and link[
-                           'target'] not in self.group_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.intelligence_nodes and link[
+            'target'] not in self.intelligence_nodes and link[
+            'target'] not in self.group_nodes
+      ]
 
       replace_nodes = {}
       relationship_links = (
-          link_ for link_ in non_group_links
+          link_ for link_ in non_special_relationship_links
           if link_["source"].startswith("relationship"))
       for link in relationship_links:
         if link["source"] not in replace_nodes:
@@ -417,7 +427,7 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
           replace_nodes[link["source"]].append(link["target"])
 
       non_relationship_links = (
-          link for link in non_group_links
+          link for link in non_special_relationship_links
           if not link["source"].startswith("relationship"))
 
       for link_data in non_relationship_links:
@@ -463,6 +473,18 @@ def _add_group_links_from_json_data(self, json_graph_data_links):
                         link['source'] in self.group_nodes or link[
                           'target'] in self.group_nodes]
 
+  def _add_intelligence_nodes_from_json_data(self, json_graph_data_nodes):
+    self.intelligence_nodes = dict([
+        (node['entity_id'], node) for node in json_graph_data_nodes
+        if node.get("type") == "relationship" and node.get(
+            "entity_attributes", {}).get("relationship_type") == "intelligence"])
+
+  def _add_intelligence_links_from_json_data(self, json_graph_data_links):
+    self.intelligence_links = [link for link in json_graph_data_links if
+                               link['source'] in self.intelligence_nodes or
+                               link[
+                                 'target'] in self.intelligence_nodes]
+
   def _pull_viewers(self):
     """Pull graph's users and groups viewers from VT API.
 
@@ -1708,6 +1730,16 @@ def save_graph(self):
       self._add_node_to_output(output, node_id)
       added.add(node_id)
 
+    intelligence_nodes = self._get_intelligence_nodes()
+    output_nodes = output["data"]["attributes"]["nodes"]
+    output["data"]["attributes"]["nodes"] = [*output_nodes, *intelligence_nodes]
+    for node in intelligence_nodes:
+      added.add(node['entity_id'])
+
+    intelligence_links = self._get_intelligence_links(added)
+    output_links = output["data"]["attributes"]["links"]
+    output["data"]["attributes"]["links"] = [*output_links, *intelligence_links]
+
     group_nodes = self._get_groups_nodes(added)
     output_nodes = output["data"]["attributes"]["nodes"]
     output["data"]["attributes"]["nodes"] = [*output_nodes, *group_nodes]
@@ -1825,6 +1857,13 @@ def _get_groups_links(self, final_nodes):
     return [link for link in self.group_links if
             link['source'] in final_nodes and link['target'] in final_nodes]
 
+  def _get_intelligence_nodes(self):
+    return list(self.intelligence_nodes.values())
+
+  def _get_intelligence_links(self, final_nodes):
+    return [link for link in self.intelligence_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.
 
@@ -1882,12 +1921,3 @@ def set_representation(self, representation):
     """
     self.representation = representation
 
-
-if __name__ == "__main__":
-  import os
-
-  graph = VTGraph.load_graph(
-      "gb49cd88b5e6c46198623e84ed6e39c0378c661fb599d443b8170e2abb789ab65",
-      api_key=os.environ["VT_API_KEY"])
-
-  graph.save_graph()
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": [

From 335203d70c78789f84a854ea337608dcfcc57a30 Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 11:19:47 +0200
Subject: [PATCH 03/10] Add tests

---
 README.md                                |   4 +
 tests/resources/virustotal_graph_id.json | 393 ++++++++++++-----------
 tests/test_create_group.py               |  72 +++++
 tests/test_load_graph.py                 |  36 ++-
 vt_graph_api/errors.py                   |   3 +
 vt_graph_api/graph.py                    | 217 ++++++++++---
 vt_graph_api/version.py                  |   2 +-
 7 files changed, 500 insertions(+), 227 deletions(-)
 create mode 100644 tests/test_create_group.py

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..861ea92
--- /dev/null
+++ b/tests/test_create_group.py
@@ -0,0 +1,72 @@
+"""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"Nodes .+ are already in groups."):
+    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(['virustotal.com', '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..bc1b018 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
\ No newline at end of file
diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index 49ecb41..ec35099 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -16,6 +16,7 @@
 
 from enum import Enum
 
+import base64
 import concurrent.futures
 import requests
 import six
@@ -146,12 +147,10 @@ def load_graph(graph_id, api_key):
     graph._pull_editors()
     # Adds group nodes from json data
     graph._add_group_nodes_from_json_data(nodes)
-    # Adds group links from json data
-    graph._add_group_links_from_json_data(links)
     # Adds intelligence nodes from json_data
-    graph._add_intelligence_nodes_from_json_data(nodes)
+    graph._add_special_relationship_nodes_from_json_data(nodes)
     # Adds intelligence links from json data
-    graph._add_intelligence_links_from_json_data(links)
+    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.
@@ -237,9 +236,8 @@ def __init__(self, api_key, name="", private=False, user_editors=None,
     self.nodes = {}
     self.links = {}
     self.group_nodes = {}
-    self.group_links = {}
-    self.intelligence_nodes = {}
-    self.intelligence_links = {}
+    self.special_relationship_nodes = {}
+    self.special_relationship_links = {}
 
     self._id_references = {}
     self._api_calls_lock = threading.Lock()
@@ -411,8 +409,8 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
       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.intelligence_nodes and link[
-            'target'] not in self.intelligence_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
       ]
 
@@ -447,10 +445,6 @@ def _add_group_nodes_from_json_data(self, json_graph_data_nodes):
 
     json_graph_data_nodes are the responses from querying VT.
 
-    Raises:
-      InvalidJSONError: whether the API response does not have the correct
-        structure.
-
     Args:
       json_graph_data_nodes ([dict]): list of node's data with the following
       structure => {
@@ -468,22 +462,57 @@ def _add_group_nodes_from_json_data(self, json_graph_data_nodes):
         if node.get("type") == "relationship" and node.get(
             "entity_attributes", {}).get("relationship_type") == "group"])
 
-  def _add_group_links_from_json_data(self, json_graph_data_links):
-    self.group_links = [link for link in json_graph_data_links if
-                        link['source'] in self.group_nodes or link[
-                          'target'] in self.group_nodes]
+  def _add_special_relationship_nodes_from_json_data(self,
+      json_graph_data_nodes):
+    """Add all the special relationship nodes from the given data.
 
-  def _add_intelligence_nodes_from_json_data(self, json_graph_data_nodes):
-    self.intelligence_nodes = dict([
-        (node['entity_id'], node) for node in json_graph_data_nodes
-        if node.get("type") == "relationship" and node.get(
-            "entity_attributes", {}).get("relationship_type") == "intelligence"])
+    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": ""
+      }
+    """
+    special_relationship_nodes = []
+
+    for node in json_graph_data_nodes:
+      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
+
+      if (is_relationship_node and
+          (relationship_type and relationship_type != "group") or
+          is_commonalities):
+        special_relationship_nodes.append((node['entity_id'], node))
+
+    self.special_relationship_nodes = dict(special_relationship_nodes)
+
+  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.
 
-  def _add_intelligence_links_from_json_data(self, json_graph_data_links):
-    self.intelligence_links = [link for link in json_graph_data_links if
-                               link['source'] in self.intelligence_nodes or
-                               link[
-                                 'target'] in self.intelligence_nodes]
+    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.
@@ -1730,15 +1759,17 @@ def save_graph(self):
       self._add_node_to_output(output, node_id)
       added.add(node_id)
 
-    intelligence_nodes = self._get_intelligence_nodes()
+    special_relationship_nodes = self._get_special_relationship_nodes()
     output_nodes = output["data"]["attributes"]["nodes"]
-    output["data"]["attributes"]["nodes"] = [*output_nodes, *intelligence_nodes]
-    for node in intelligence_nodes:
+    output["data"]["attributes"]["nodes"] = [*output_nodes,
+                                             *special_relationship_nodes]
+    for node in special_relationship_nodes:
       added.add(node['entity_id'])
 
-    intelligence_links = self._get_intelligence_links(added)
+    special_relationship_links = self._get_special_relationship_links(added)
     output_links = output["data"]["attributes"]["links"]
-    output["data"]["attributes"]["links"] = [*output_links, *intelligence_links]
+    output["data"]["attributes"]["links"] = [*output_links,
+                                             *special_relationship_links]
 
     group_nodes = self._get_groups_nodes(added)
     output_nodes = output["data"]["attributes"]["nodes"]
@@ -1746,9 +1777,9 @@ def save_graph(self):
     for node in group_nodes:
       added.add(node['entity_id'])
 
-    group_links = self._get_groups_links(added)
-    output_links = output["data"]["attributes"]["links"]
-    output["data"]["attributes"]["links"] = [*output_links, *group_links]
+    final_links = output["data"]["attributes"]["links"]
+    group_links = self._get_groups_links(final_links)
+    output["data"]["attributes"]["links"] = [*final_links, *group_links]
 
     self._push_graph_to_vt(output)
     self._push_editors()
@@ -1841,6 +1872,14 @@ def _get_nodes_by_type(self, node_type):
     ]}
 
   def _get_groups_nodes(self, final_nodes):
+    """Returns all the loaded and generated group nodes.
+
+    Args:
+      final_nodes: Set with the ids of all the added nodes.
+
+    Returns: Array of group nodes.
+
+    """
     group_nodes = []
     current_index = len(final_nodes)
     for idx, group_node in enumerate(self.group_nodes.values()):
@@ -1853,15 +1892,53 @@ def _get_groups_nodes(self, final_nodes):
         group_nodes.append(group_node)
     return group_nodes
 
-  def _get_groups_links(self, final_nodes):
-    return [link for link in self.group_links if
-            link['source'] in final_nodes and link['target'] in final_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):
+    """Returns a list with the special relationship nodes."""
+    return list(self.special_relationship_nodes.values())
+
+  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.
+
+    Args:
+      final_nodes: Set with all the nodes added
 
-  def _get_intelligence_nodes(self):
-    return list(self.intelligence_nodes.values())
+    Returns: list of links related to special_relationship_nodes.
 
-  def _get_intelligence_links(self, final_nodes):
-    return [link for link in self.intelligence_links if
+    """
+    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):
@@ -1913,6 +1990,61 @@ def create_collection(self, name=None, description=None):
         collection_id=collection_id
     )
 
+  def _generate_group_node_id(self, nodes_id):
+    return ",".join(nodes_id).replace("=", "")
+
+  def create_group(self, node_ids, group_name):
+    """
+
+    Args:
+      node_ids:
+      group_name:
+
+    Returns:
+
+    Raises: CreateGroupError if the
+
+    """
+    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.
+    nodes_already_in_a_group = set()
+    for group_node in self.group_nodes.values():
+      for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
+        nodes_already_in_a_group.add(node_id)
+
+    intersection = nodes_already_in_a_group.intersection(node_ids)
+    if (len(intersection) > 0):
+      raise vt_graph_api.errors.CreateGroupError(
+          "Nodes {intersection} are already in groups.".format(
+              intersection=",".join(list(intersection))))
+
+    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.
 
@@ -1920,4 +2052,3 @@ def set_representation(self, representation):
       representation: Graph representation. See :py:class::`RepresentationType`.
     """
     self.representation = representation
-
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'

From c2e142d0bba2d32a014e513b85b7ad3fbbee9f6b Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 11:21:07 +0200
Subject: [PATCH 04/10] Add new line

---
 vt_graph_api/errors.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vt_graph_api/errors.py b/vt_graph_api/errors.py
index bc1b018..539ac90 100644
--- a/vt_graph_api/errors.py
+++ b/vt_graph_api/errors.py
@@ -51,4 +51,4 @@ class CreateCollectionError(Exception):
   pass
 
 class CreateGroupError(Exception):
-  pass
\ No newline at end of file
+  pass

From 7ce2c728faed731f034400948313f9bfe85b0253 Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 11:25:57 +0200
Subject: [PATCH 05/10] Concat lists

---
 vt_graph_api/graph.py | 485 ++++++++++++++++++++++--------------------
 1 file changed, 249 insertions(+), 236 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index ec35099..558660a 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -1761,294 +1761,307 @@ def save_graph(self):
 
     special_relationship_nodes = self._get_special_relationship_nodes()
     output_nodes = output["data"]["attributes"]["nodes"]
-    output["data"]["attributes"]["nodes"] = [*output_nodes,
-                                             *special_relationship_nodes]
+    output["data"]["attributes"][
+      "nodes"] = output_nodes + special_relationship_nodes
     for node in special_relationship_nodes:
       added.add(node['entity_id'])
 
     special_relationship_links = self._get_special_relationship_links(added)
     output_links = output["data"]["attributes"]["links"]
-    output["data"]["attributes"]["links"] = [*output_links,
-                                             *special_relationship_links]
+    output["data"]["attributes"][
+      "links"] = output_links + special_relationship_links
 
-    group_nodes = self._get_groups_nodes(added)
-    output_nodes = output["data"]["attributes"]["nodes"]
-    output["data"]["attributes"]["nodes"] = [*output_nodes, *group_nodes]
-    for node in group_nodes:
-      added.add(node['entity_id'])
+  group_nodes = self._get_groups_nodes(added)
+  output_nodes = output["data"]["attributes"]["nodes"]
+  output["data"]["attributes"]["nodes"] = output_nodes + group_nodes
+  for node in group_nodes:
+    added.add(node['entity_id'])
 
-    final_links = output["data"]["attributes"]["links"]
-    group_links = self._get_groups_links(final_links)
-    output["data"]["attributes"]["links"] = [*final_links, *group_links]
+  final_links = output["data"]["attributes"]["links"]
+  group_links = self._get_groups_links(final_links)
+  output["data"]["attributes"]["links"] = final_links + group_links
 
-    self._push_graph_to_vt(output)
-    self._push_editors()
-    self._push_viewers()
-    self._index = 0
+  self._push_graph_to_vt(output)
+  self._push_editors()
+  self._push_viewers()
+  self._index = 0
 
-  def get_api_calls(self):
-    """Get api counter in thread safe mode."""
-    with self._api_calls_lock:
-      api_calls = self._api_calls
-    return api_calls
 
-  def get_ui_link(self):
-    """Return VirusTotal UI link for the graph.
+def get_api_calls(self):
+  """Get api counter in thread safe mode."""
+  with self._api_calls_lock:
+    api_calls = self._api_calls
+  return api_calls
 
-    Requires that save_graph was called.
 
-    Raises:
-      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
+def get_ui_link(self):
+  """Return VirusTotal UI link for the graph.
 
-    Returns:
-        str: VirusTotal UI link.
-    """
-    if not self.graph_id:
-      raise vt_graph_api.errors.SaveGraphError(
-          "`save_graph` has not been called yet!")
-    return "https://www.virustotal.com/graph/{graph_id}".format(
-        graph_id=self.graph_id)
+  Requires that save_graph was called.
 
-  def get_iframe_code(self):
-    """Return VirusTotal UI iframe for the graph.
+  Raises:
+    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
 
-    Requires that save_graph was called.
+  Returns:
+      str: VirusTotal UI link.
+  """
+  if not self.graph_id:
+    raise vt_graph_api.errors.SaveGraphError(
+        "`save_graph` has not been called yet!")
+  return "https://www.virustotal.com/graph/{graph_id}".format(
+      graph_id=self.graph_id)
 
-    Raises:
-      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
 
-    Returns:
-        str: VirusTotal UI iframe.
-    """
-    if not self.graph_id:
-      raise vt_graph_api.errors.SaveGraphError(
-          "`save_graph` has not been called yet!")
-    return (
-        "<iframe src=\"https://www.virustotal.com/graph/embed/" +
-        "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
-        .format(graph_id=self.graph_id))
+def get_iframe_code(self):
+  """Return VirusTotal UI iframe for the graph.
 
-  def download_screenshot(self, path="."):
-    """Downloads a screenshot of the graph.
+  Requires that save_graph was called.
 
-    Args:
-      path: Path where screenshot will be saved.
+  Raises:
+    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
 
-    Raises:
-      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
-      vt_graph_api.errors.DownloadScreenshotError: if screenshot can't be downloaded.
+  Returns:
+      str: VirusTotal UI iframe.
+  """
+  if not self.graph_id:
+    raise vt_graph_api.errors.SaveGraphError(
+        "`save_graph` has not been called yet!")
+  return (
+      "<iframe src=\"https://www.virustotal.com/graph/embed/" +
+      "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
+      .format(graph_id=self.graph_id))
 
-    """
-    if not self.graph_id:
-      raise vt_graph_api.errors.SaveGraphError(
-          "`save_graph` has not been called yet!")
 
-    url = "https://www.virustotal.com/api/v3/graphs/{graph_id}/screenshot".format(
-        graph_id=self.graph_id
+def download_screenshot(self, path="."):
+  """Downloads a screenshot of the graph.
+
+  Args:
+    path: Path where screenshot will be saved.
+
+  Raises:
+    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
+    vt_graph_api.errors.DownloadScreenshotError: if screenshot can't be downloaded.
+
+  """
+  if not self.graph_id:
+    raise vt_graph_api.errors.SaveGraphError(
+        "`save_graph` has not been called yet!")
+
+  url = "https://www.virustotal.com/api/v3/graphs/{graph_id}/screenshot".format(
+      graph_id=self.graph_id
+  )
+
+  r = requests.get(
+      url,
+      headers=self._get_headers(),
+      stream=True)
+
+  if r.status_code == 200:
+    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:
+      shutil.copyfileobj(r.raw, f)
+  else:
+    raise vt_graph_api.errors.DownloadScreenshotError(
+        "Couldn't download screenshot for graph {graph_id}".format(
+            graph_id=self.graph_id
+        )
     )
 
-    r = requests.get(
-        url,
-        headers=self._get_headers(),
-        stream=True)
-
-    if r.status_code == 200:
-      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:
-        shutil.copyfileobj(r.raw, f)
-    else:
-      raise vt_graph_api.errors.DownloadScreenshotError(
-          "Couldn't download screenshot for graph {graph_id}".format(
-              graph_id=self.graph_id
-          )
-      )
 
-  def _get_nodes_by_type(self, node_type):
-    return {"data": [
-        {'type': node.node_type, 'id': node.node_id}
-        for node in self.nodes.values() if node.node_type == node_type
-    ]}
+def _get_nodes_by_type(self, node_type):
+  return {"data": [
+      {'type': node.node_type, 'id': node.node_id}
+      for node in self.nodes.values() if node.node_type == node_type
+  ]}
 
-  def _get_groups_nodes(self, final_nodes):
-    """Returns all the loaded and generated group nodes.
 
-    Args:
-      final_nodes: Set with the ids of all the added nodes.
+def _get_groups_nodes(self, final_nodes):
+  """Returns all the loaded and generated group nodes.
 
-    Returns: Array of group nodes.
+  Args:
+    final_nodes: Set with the ids of all the added nodes.
 
-    """
-    group_nodes = []
-    current_index = len(final_nodes)
-    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'] = current_index + idx
-        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.
+  Returns: Array of group nodes.
 
-    Args:
-      final_links: list with all the links added in the Graph.
+  """
+  group_nodes = []
+  current_index = len(final_nodes)
+  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'] = current_index + idx
+      group_nodes.append(group_node)
+  return group_nodes
 
-    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):
-    """Returns a list with the special relationship nodes."""
-    return list(self.special_relationship_nodes.values())
-
-  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.
+def _get_groups_links(self, final_links):
+  """Generates all the links needed to represent groups in VTGraph.
 
-    Args:
-      final_nodes: Set with all the nodes added
+  Args:
+    final_links: list with all the links added in the Graph.
 
-    Returns: list of links related to special_relationship_nodes.
+  Returns: List with the links needed to represent groups in VTGraph.
 
-    """
-    return [link for link in self.special_relationship_links if
-            link['source'] in final_nodes and link['target'] in final_nodes]
+  """
+  map_node_group = {}
 
-  def create_collection(self, name=None, description=None):
-    """Creates a VT Collection taking entities from current Graph.
+  # 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']
 
-    Args:
-      name: Collection name.
-      description: Collection description
+  group_links = [{'source': node, 'target': group, 'connection_type': 'group'}
+                 for node, group in map_node_group.items()]
 
-    Raises:
-      vt_graph_api.errors.CreateCollectionError: if the collection couldn't be
-        created
+  # 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
 
-    Returns:
-      str: VirusTotal UI Collection link.
-    """
 
-    data = {
-        "type": "collection",
-        "attributes": {
-            "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')
-        }
-    }
+def _get_special_relationship_nodes(self):
+  """Returns a list with the special relationship nodes."""
+  return list(self.special_relationship_nodes.values())
 
-    if description:
-      data["attributes"]["description"] = description
-    elif self.graph_id:
-      data["attributes"]["description"] = (
-          "Collection created from graph {graph_id}").format(
-          graph_id=self.graph_id)
 
-    url = "https://www.virustotal.com/api/v3/collections"
-    response = requests.post(
-        url, headers=self._get_headers(), json={"data": data})
+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.
 
-    if (response.status_code != 200):
-      print(response.json())
-      raise vt_graph_api.errors.CreateCollectionError()
+  Args:
+    final_nodes: Set with all the nodes added
 
-    collection_id = response.json()["data"]["id"]
+  Returns: list of links related to special_relationship_nodes.
 
-    return "https://www.virustotal.com/gui/collection/{collection_id}".format(
-        collection_id=collection_id
-    )
+  """
+  return [link for link in self.special_relationship_links if
+          link['source'] in final_nodes and link['target'] in final_nodes]
 
-  def _generate_group_node_id(self, nodes_id):
-    return ",".join(nodes_id).replace("=", "")
 
-  def create_group(self, node_ids, group_name):
-    """
+def create_collection(self, name=None, description=None):
+  """Creates a VT Collection taking entities from current Graph.
 
-    Args:
-      node_ids:
-      group_name:
+  Args:
+    name: Collection name.
+    description: Collection description
 
-    Returns:
+  Raises:
+    vt_graph_api.errors.CreateCollectionError: if the collection couldn't be
+      created
 
-    Raises: CreateGroupError if the
+  Returns:
+    str: VirusTotal UI Collection link.
+  """
 
-    """
-    group_node_id = self._generate_group_node_id(node_ids)
-    node_ids_set = set(node_ids)
+  data = {
+      "type": "collection",
+      "attributes": {
+          "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')
+      }
+  }
 
-    # 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.
-    nodes_already_in_a_group = set()
-    for group_node in self.group_nodes.values():
-      for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
-        nodes_already_in_a_group.add(node_id)
-
-    intersection = nodes_already_in_a_group.intersection(node_ids)
-    if (len(intersection) > 0):
-      raise vt_graph_api.errors.CreateGroupError(
-          "Nodes {intersection} are already in groups.".format(
-              intersection=",".join(list(intersection))))
-
-    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
-    }
+  if description:
+    data["attributes"]["description"] = description
+  elif self.graph_id:
+    data["attributes"]["description"] = (
+        "Collection created from graph {graph_id}").format(
+        graph_id=self.graph_id)
 
-    self.group_nodes[group_node_id] = relationship_node
+  url = "https://www.virustotal.com/api/v3/collections"
+  response = requests.post(
+      url, headers=self._get_headers(), json={"data": data})
 
-  def set_representation(self, representation):
-    """Sets Graph representation.
+  if (response.status_code != 200):
+    print(response.json())
+    raise vt_graph_api.errors.CreateCollectionError()
 
-    Args:
-      representation: Graph representation. See :py:class::`RepresentationType`.
-    """
-    self.representation = representation
+  collection_id = response.json()["data"]["id"]
+
+  return "https://www.virustotal.com/gui/collection/{collection_id}".format(
+      collection_id=collection_id
+  )
+
+
+def _generate_group_node_id(self, nodes_id):
+  return ",".join(nodes_id).replace("=", "")
+
+
+def create_group(self, node_ids, group_name):
+  """
+
+  Args:
+    node_ids:
+    group_name:
+
+  Returns:
+
+  Raises: CreateGroupError if the
+
+  """
+  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.
+  nodes_already_in_a_group = set()
+  for group_node in self.group_nodes.values():
+    for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
+      nodes_already_in_a_group.add(node_id)
+
+  intersection = nodes_already_in_a_group.intersection(node_ids)
+  if (len(intersection) > 0):
+    raise vt_graph_api.errors.CreateGroupError(
+        "Nodes {intersection} are already in groups.".format(
+            intersection=",".join(list(intersection))))
+
+  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

From c092b2448832687d9371e7bf9219eff36271a796 Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 11:27:41 +0200
Subject: [PATCH 06/10] Concat lists

---
 vt_graph_api/graph.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index 558660a..b726599 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -1771,20 +1771,20 @@ def save_graph(self):
     output["data"]["attributes"][
       "links"] = output_links + special_relationship_links
 
-  group_nodes = self._get_groups_nodes(added)
-  output_nodes = output["data"]["attributes"]["nodes"]
-  output["data"]["attributes"]["nodes"] = output_nodes + group_nodes
-  for node in group_nodes:
-    added.add(node['entity_id'])
-
-  final_links = output["data"]["attributes"]["links"]
-  group_links = self._get_groups_links(final_links)
-  output["data"]["attributes"]["links"] = final_links + group_links
-
-  self._push_graph_to_vt(output)
-  self._push_editors()
-  self._push_viewers()
-  self._index = 0
+    group_nodes = self._get_groups_nodes(added)
+    output_nodes = output["data"]["attributes"]["nodes"]
+    output["data"]["attributes"]["nodes"] = output_nodes + group_nodes
+    for node in group_nodes:
+      added.add(node['entity_id'])
+
+    final_links = output["data"]["attributes"]["links"]
+    group_links = self._get_groups_links(final_links)
+    output["data"]["attributes"]["links"] = final_links + group_links
+
+    self._push_graph_to_vt(output)
+    self._push_editors()
+    self._push_viewers()
+    self._index = 0
 
 
 def get_api_calls(self):

From ee6ea280897b94bec98d6d3b7fa91d5f1c98c17b Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 11:28:26 +0200
Subject: [PATCH 07/10] Fix indentation

---
 vt_graph_api/graph.py | 438 +++++++++++++++++++++---------------------
 1 file changed, 219 insertions(+), 219 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index b726599..f44ea3e 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -1787,281 +1787,281 @@ def save_graph(self):
     self._index = 0
 
 
-def get_api_calls(self):
-  """Get api counter in thread safe mode."""
-  with self._api_calls_lock:
-    api_calls = self._api_calls
-  return api_calls
+  def get_api_calls(self):
+    """Get api counter in thread safe mode."""
+    with self._api_calls_lock:
+      api_calls = self._api_calls
+    return api_calls
 
 
-def get_ui_link(self):
-  """Return VirusTotal UI link for the graph.
+  def get_ui_link(self):
+    """Return VirusTotal UI link for the graph.
 
-  Requires that save_graph was called.
+    Requires that save_graph was called.
 
-  Raises:
-    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
+    Raises:
+      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
 
-  Returns:
-      str: VirusTotal UI link.
-  """
-  if not self.graph_id:
-    raise vt_graph_api.errors.SaveGraphError(
-        "`save_graph` has not been called yet!")
-  return "https://www.virustotal.com/graph/{graph_id}".format(
-      graph_id=self.graph_id)
+    Returns:
+        str: VirusTotal UI link.
+    """
+    if not self.graph_id:
+      raise vt_graph_api.errors.SaveGraphError(
+          "`save_graph` has not been called yet!")
+    return "https://www.virustotal.com/graph/{graph_id}".format(
+        graph_id=self.graph_id)
 
 
-def get_iframe_code(self):
-  """Return VirusTotal UI iframe for the graph.
+  def get_iframe_code(self):
+    """Return VirusTotal UI iframe for the graph.
 
-  Requires that save_graph was called.
+    Requires that save_graph was called.
 
-  Raises:
-    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
+    Raises:
+      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
 
-  Returns:
-      str: VirusTotal UI iframe.
-  """
-  if not self.graph_id:
-    raise vt_graph_api.errors.SaveGraphError(
-        "`save_graph` has not been called yet!")
-  return (
-      "<iframe src=\"https://www.virustotal.com/graph/embed/" +
-      "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
-      .format(graph_id=self.graph_id))
+    Returns:
+        str: VirusTotal UI iframe.
+    """
+    if not self.graph_id:
+      raise vt_graph_api.errors.SaveGraphError(
+          "`save_graph` has not been called yet!")
+    return (
+        "<iframe src=\"https://www.virustotal.com/graph/embed/" +
+        "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
+        .format(graph_id=self.graph_id))
 
 
-def download_screenshot(self, path="."):
-  """Downloads a screenshot of the graph.
+  def download_screenshot(self, path="."):
+    """Downloads a screenshot of the graph.
+
+    Args:
+      path: Path where screenshot will be saved.
 
-  Args:
-    path: Path where screenshot will be saved.
+    Raises:
+      vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
+      vt_graph_api.errors.DownloadScreenshotError: if screenshot can't be downloaded.
 
-  Raises:
-    vt_graph_api.errors.SaveGraphError: if `save_graph` was not called.
-    vt_graph_api.errors.DownloadScreenshotError: if screenshot can't be downloaded.
+    """
+    if not self.graph_id:
+      raise vt_graph_api.errors.SaveGraphError(
+          "`save_graph` has not been called yet!")
 
-  """
-  if not self.graph_id:
-    raise vt_graph_api.errors.SaveGraphError(
-        "`save_graph` has not been called yet!")
-
-  url = "https://www.virustotal.com/api/v3/graphs/{graph_id}/screenshot".format(
-      graph_id=self.graph_id
-  )
-
-  r = requests.get(
-      url,
-      headers=self._get_headers(),
-      stream=True)
-
-  if r.status_code == 200:
-    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:
-      shutil.copyfileobj(r.raw, f)
-  else:
-    raise vt_graph_api.errors.DownloadScreenshotError(
-        "Couldn't download screenshot for graph {graph_id}".format(
-            graph_id=self.graph_id
-        )
+    url = "https://www.virustotal.com/api/v3/graphs/{graph_id}/screenshot".format(
+        graph_id=self.graph_id
     )
 
+    r = requests.get(
+        url,
+        headers=self._get_headers(),
+        stream=True)
+
+    if r.status_code == 200:
+      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:
+        shutil.copyfileobj(r.raw, f)
+    else:
+      raise vt_graph_api.errors.DownloadScreenshotError(
+          "Couldn't download screenshot for graph {graph_id}".format(
+              graph_id=self.graph_id
+          )
+      )
+
 
-def _get_nodes_by_type(self, node_type):
-  return {"data": [
-      {'type': node.node_type, 'id': node.node_id}
-      for node in self.nodes.values() if node.node_type == node_type
-  ]}
+  def _get_nodes_by_type(self, node_type):
+    return {"data": [
+        {'type': node.node_type, 'id': node.node_id}
+        for node in self.nodes.values() if node.node_type == node_type
+    ]}
 
 
-def _get_groups_nodes(self, final_nodes):
-  """Returns all the loaded and generated group nodes.
+  def _get_groups_nodes(self, final_nodes):
+    """Returns all the loaded and generated group nodes.
 
-  Args:
-    final_nodes: Set with the ids of all the added nodes.
+    Args:
+      final_nodes: Set with the ids of all the added nodes.
 
-  Returns: Array of group nodes.
+    Returns: Array of group nodes.
 
-  """
-  group_nodes = []
-  current_index = len(final_nodes)
-  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'] = current_index + idx
-      group_nodes.append(group_node)
-  return group_nodes
+    """
+    group_nodes = []
+    current_index = len(final_nodes)
+    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'] = current_index + idx
+        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.
 
-def _get_groups_links(self, final_links):
-  """Generates all the links needed to represent groups in VTGraph.
+    """
+    map_node_group = {}
 
-  Args:
-    final_links: list with all the links added in the Graph.
+    # 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']
 
-  Returns: List with the links needed to represent groups in VTGraph.
+    group_links = [{'source': node, 'target': group, 'connection_type': 'group'}
+                   for node, group in map_node_group.items()]
 
-  """
-  map_node_group = {}
+    # 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
 
-  # 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()]
+  def _get_special_relationship_nodes(self):
+    """Returns a list with the special relationship nodes."""
+    return list(self.special_relationship_nodes.values())
 
-  # 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_links(self, final_nodes):
+    """Returns links needed to represent special relationship nodes.
+    It checks if the link connections exist or not.
 
-def _get_special_relationship_nodes(self):
-  """Returns a list with the special relationship nodes."""
-  return list(self.special_relationship_nodes.values())
+    Args:
+      final_nodes: Set with all the nodes added
 
+    Returns: list of links related to 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.
+    """
+    return [link for link in self.special_relationship_links if
+            link['source'] in final_nodes and link['target'] in final_nodes]
 
-  Args:
-    final_nodes: Set with all the nodes added
 
-  Returns: list of links related to special_relationship_nodes.
+  def create_collection(self, name=None, description=None):
+    """Creates a VT Collection taking entities from current Graph.
 
-  """
-  return [link for link in self.special_relationship_links if
-          link['source'] in final_nodes and link['target'] in final_nodes]
+    Args:
+      name: Collection name.
+      description: Collection description
 
+    Raises:
+      vt_graph_api.errors.CreateCollectionError: if the collection couldn't be
+        created
 
-def create_collection(self, name=None, description=None):
-  """Creates a VT Collection taking entities from current Graph.
+    Returns:
+      str: VirusTotal UI Collection link.
+    """
 
-  Args:
-    name: Collection name.
-    description: Collection description
+    data = {
+        "type": "collection",
+        "attributes": {
+            "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')
+        }
+    }
 
-  Raises:
-    vt_graph_api.errors.CreateCollectionError: if the collection couldn't be
-      created
+    if description:
+      data["attributes"]["description"] = description
+    elif self.graph_id:
+      data["attributes"]["description"] = (
+          "Collection created from graph {graph_id}").format(
+          graph_id=self.graph_id)
 
-  Returns:
-    str: VirusTotal UI Collection link.
-  """
+    url = "https://www.virustotal.com/api/v3/collections"
+    response = requests.post(
+        url, headers=self._get_headers(), json={"data": data})
 
-  data = {
-      "type": "collection",
-      "attributes": {
-          "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')
-      }
-  }
+    if (response.status_code != 200):
+      print(response.json())
+      raise vt_graph_api.errors.CreateCollectionError()
 
-  if description:
-    data["attributes"]["description"] = description
-  elif self.graph_id:
-    data["attributes"]["description"] = (
-        "Collection created from graph {graph_id}").format(
-        graph_id=self.graph_id)
+    collection_id = response.json()["data"]["id"]
 
-  url = "https://www.virustotal.com/api/v3/collections"
-  response = requests.post(
-      url, headers=self._get_headers(), json={"data": data})
+    return "https://www.virustotal.com/gui/collection/{collection_id}".format(
+        collection_id=collection_id
+    )
 
-  if (response.status_code != 200):
-    print(response.json())
-    raise vt_graph_api.errors.CreateCollectionError()
 
-  collection_id = response.json()["data"]["id"]
+  def _generate_group_node_id(self, nodes_id):
+    return ",".join(nodes_id).replace("=", "")
 
-  return "https://www.virustotal.com/gui/collection/{collection_id}".format(
-      collection_id=collection_id
-  )
 
+  def create_group(self, node_ids, group_name):
+    """
 
-def _generate_group_node_id(self, nodes_id):
-  return ",".join(nodes_id).replace("=", "")
+    Args:
+      node_ids:
+      group_name:
 
+    Returns:
 
-def create_group(self, node_ids, group_name):
-  """
+    Raises: CreateGroupError if the
 
-  Args:
-    node_ids:
-    group_name:
+    """
+    group_node_id = self._generate_group_node_id(node_ids)
+    node_ids_set = set(node_ids)
 
-  Returns:
+    # 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.
+    nodes_already_in_a_group = set()
+    for group_node in self.group_nodes.values():
+      for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
+        nodes_already_in_a_group.add(node_id)
+
+    intersection = nodes_already_in_a_group.intersection(node_ids)
+    if (len(intersection) > 0):
+      raise vt_graph_api.errors.CreateGroupError(
+          "Nodes {intersection} are already in groups.".format(
+              intersection=",".join(list(intersection))))
+
+    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
+    }
 
-  Raises: CreateGroupError if the
+    self.group_nodes[group_node_id] = relationship_node
 
-  """
-  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.")
+  def set_representation(self, representation):
+    """Sets Graph representation.
 
-  # 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.
-  nodes_already_in_a_group = set()
-  for group_node in self.group_nodes.values():
-    for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
-      nodes_already_in_a_group.add(node_id)
-
-  intersection = nodes_already_in_a_group.intersection(node_ids)
-  if (len(intersection) > 0):
-    raise vt_graph_api.errors.CreateGroupError(
-        "Nodes {intersection} are already in groups.".format(
-            intersection=",".join(list(intersection))))
-
-  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
+    Args:
+      representation: Graph representation. See :py:class::`RepresentationType`.
+    """
+    self.representation = representation

From d7c59374e6a8ebd3d51e7541252066363c44616e Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 15:41:37 +0200
Subject: [PATCH 08/10] PR fixes

---
 tests/test_create_group.py |   7 +-
 vt_graph_api/graph.py      | 161 ++++++++++++++++++++-----------------
 2 files changed, 91 insertions(+), 77 deletions(-)

diff --git a/tests/test_create_group.py b/tests/test_create_group.py
index 861ea92..25dfc73 100644
--- a/tests/test_create_group.py
+++ b/tests/test_create_group.py
@@ -30,7 +30,7 @@ def test_create_group_with_nodes_already_grouped():
   test_graph.create_group(['virustotal.com', 'google.com'], 'Group 1')
 
   with pytest.raises(vt_graph_api.errors.CreateGroupError,
-                     match=r"Nodes .+ are already in groups."):
+                     match=r"Node .+ is already in a group."):
     test_graph.create_group(['virustotal.com', 'google.com'], "Group 1")
 
 
@@ -43,7 +43,7 @@ def test_create_group_with_node_that_does_not_exist():
 
   with pytest.raises(vt_graph_api.errors.CreateGroupError,
                      match=r"Node hola.es is not in the Graph."):
-    test_graph.create_group(['virustotal.com', 'hola.es'], "Group 1")
+    test_graph.create_group(['hola.es'], "Group 1")
 
 
 def test_create_group(mocker):
@@ -62,7 +62,8 @@ def test_create_group(mocker):
 
   # 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 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
diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index f44ea3e..a8ae702 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -406,28 +406,29 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
       # It is necessary to clean the given links because they have relationship
       # nodes
 
-      non_special_relationship_links = [
+      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 non_special_relationship_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 non_special_relationship_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"]])
@@ -462,6 +463,23 @@ def _add_group_nodes_from_json_data(self, 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.
@@ -479,20 +497,10 @@ def _add_special_relationship_nodes_from_json_data(self,
           "y": ""
       }
     """
-    special_relationship_nodes = []
-
-    for node in json_graph_data_nodes:
-      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
-
-      if (is_relationship_node and
-          (relationship_type and relationship_type != "group") or
-          is_commonalities):
-        special_relationship_nodes.append((node['entity_id'], node))
-
-    self.special_relationship_nodes = dict(special_relationship_nodes)
+    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.
@@ -1759,41 +1767,35 @@ def save_graph(self):
       self._add_node_to_output(output, node_id)
       added.add(node_id)
 
-    special_relationship_nodes = self._get_special_relationship_nodes()
-    output_nodes = output["data"]["attributes"]["nodes"]
-    output["data"]["attributes"][
-      "nodes"] = output_nodes + special_relationship_nodes
+    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_links = output["data"]["attributes"]["links"]
-    output["data"]["attributes"][
-      "links"] = output_links + special_relationship_links
+    output["data"]["attributes"]["links"] += special_relationship_links
 
     group_nodes = self._get_groups_nodes(added)
-    output_nodes = output["data"]["attributes"]["nodes"]
-    output["data"]["attributes"]["nodes"] = output_nodes + group_nodes
     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"] = final_links + group_links
+    output["data"]["attributes"]["links"] += group_links
 
     self._push_graph_to_vt(output)
     self._push_editors()
     self._push_viewers()
     self._index = 0
 
-
   def get_api_calls(self):
     """Get api counter in thread safe mode."""
     with self._api_calls_lock:
       api_calls = self._api_calls
     return api_calls
 
-
   def get_ui_link(self):
     """Return VirusTotal UI link for the graph.
 
@@ -1811,7 +1813,6 @@ def get_ui_link(self):
     return "https://www.virustotal.com/graph/{graph_id}".format(
         graph_id=self.graph_id)
 
-
   def get_iframe_code(self):
     """Return VirusTotal UI iframe for the graph.
 
@@ -1831,7 +1832,6 @@ def get_iframe_code(self):
         "{graph_id}\" width=\"800\" height=\"600\"></iframe>"
         .format(graph_id=self.graph_id))
 
-
   def download_screenshot(self, path="."):
     """Downloads a screenshot of the graph.
 
@@ -1869,36 +1869,42 @@ def download_screenshot(self, path="."):
           )
       )
 
-
   def _get_nodes_by_type(self, node_type):
     return {"data": [
         {'type': node.node_type, 'id': node.node_id}
         for node in self.nodes.values() if node.node_type == node_type
     ]}
 
-
   def _get_groups_nodes(self, final_nodes):
-    """Returns all the loaded and generated group 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.
+      final_nodes: Set with the ids of all the added nodes to the Graph.
 
     Returns: Array of group nodes.
 
     """
     group_nodes = []
-    current_index = len(final_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):
+          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'] = current_index + idx
+        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.
 
@@ -1931,15 +1937,26 @@ def _get_groups_links(self, final_links):
         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.
 
-  def _get_special_relationship_nodes(self):
-    """Returns a list with the special relationship nodes."""
-    return list(self.special_relationship_nodes.values())
+    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.
+    It checks if the link connections exist or not. Only used for saving the
+    Graph.
 
     Args:
       final_nodes: Set with all the nodes added
@@ -1950,7 +1967,6 @@ def _get_special_relationship_links(self, final_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.
 
@@ -2000,11 +2016,9 @@ def create_collection(self, name=None, description=None):
         collection_id=collection_id
     )
 
-
   def _generate_group_node_id(self, nodes_id):
     return ",".join(nodes_id).replace("=", "")
 
-
   def create_group(self, node_ids, group_name):
     """
 
@@ -2014,7 +2028,10 @@ def create_group(self, node_ids, group_name):
 
     Returns:
 
-    Raises: CreateGroupError if the
+    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)
@@ -2036,27 +2053,23 @@ def create_group(self, node_ids, group_name):
     nodes_already_in_a_group = set()
     for group_node in self.group_nodes.values():
       for node_id in group_node["entity_attributes"]["grouped_node_ids"]:
-        nodes_already_in_a_group.add(node_id)
-
-    intersection = nodes_already_in_a_group.intersection(node_ids)
-    if (len(intersection) > 0):
-      raise vt_graph_api.errors.CreateGroupError(
-          "Nodes {intersection} are already in groups.".format(
-              intersection=",".join(list(intersection))))
-
-    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
+        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.

From 063cf5e2e4c98781ce15ef84022a074c402d13e7 Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 15:47:42 +0200
Subject: [PATCH 09/10] Fix test

---
 vt_graph_api/graph.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index a8ae702..068abf0 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -417,7 +417,7 @@ def _add_links_from_json_graph_data(self, json_graph_data_links):
       relationship_links = []
       non_relationship_links = []
       for link in non_special_relationship_links:
-        if link_["source"].startswith("relationship"):
+        if link["source"].startswith("relationship"):
           relationship_links.append(link)
         else:
           non_relationship_links.append(link)
@@ -2050,7 +2050,6 @@ def create_group(self, node_ids, group_name):
                 node_id=node_id))
 
     # Check if nodes are already in a group.
-    nodes_already_in_a_group = set()
     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:
@@ -2058,18 +2057,18 @@ def create_group(self, node_ids, group_name):
               "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
-          }
+    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
+    self.group_nodes[group_node_id] = relationship_node
 
   def set_representation(self, representation):
     """Sets Graph representation.

From cf50f07611fed7a84b88a7243bc5fd993df934ba Mon Sep 17 00:00:00 2001
From: Andres Ramirez <anramar@google.com>
Date: Thu, 28 Jul 2022 15:55:03 +0200
Subject: [PATCH 10/10] Use braces

---
 vt_graph_api/graph.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/vt_graph_api/graph.py b/vt_graph_api/graph.py
index 068abf0..b17f97d 100644
--- a/vt_graph_api/graph.py
+++ b/vt_graph_api/graph.py
@@ -458,10 +458,10 @@ def _add_group_nodes_from_json_data(self, json_graph_data_nodes):
       }
     """
 
-    self.group_nodes = dict([
-        (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"])
+    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.