Merge pull request #48 from VirusTotal/references-post

Update create_reference example
VirusTotal · May 10, 2021 · 8de8932 · 8de8932
2 parents 2f2b620 + 2fed7f6
commit 8de8932
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 104 deletions.
diff --git a/docs/_modules/vt/client.html b/docs/_modules/vt/client.html
@@ -655,7 +655,7 @@ <h1>Source code for vt.client</h1><div class="highlight"><pre>
 <span class="sd">    :param path: Path to API endpoint.</span>
 <span class="sd">    :param path_args: A variable number of arguments that are put into any</span>
 <span class="sd">      placeholders used in path.</span>
-<span class="sd">    :param obj: Instance :class:`Object` whith the type expected by the API</span>
+<span class="sd">    :param obj: Instance :class:`Object` with the type expected by the API</span>
 <span class="sd">      endpoint.</span>
 <span class="sd">    :type path: str</span>
 <span class="sd">    :type obj: :class:`Object`</span>

diff --git a/examples/create_reference.py b/examples/create_reference.py
@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 # Copyright © 2021 The vt-py authors. All Rights Reserved.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,82 +12,74 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-This example program explains how to generate a reference in VirusTotal
-and add IOCs to it.
+"""How to generate a reference in VirusTotal and add IOCs to it.
 
 NOTICE: In order to use this program you will need an API key that has
 privileges for creating References.
 """
 
 import argparse
-import vt
 import json
+import vt
 
-def create_reference(url, creation_date, title, author, iocs, client):
-  """ Creates a reference in VirusTotal.
+
+def create_reference(url, creation_date, title, author, client, iocs):
+  """Creates a reference in VirusTotal.
 
   Args:
     url: Reference url.
     creation_date: Reference creation date (YY-MM-DD HH:mm:ss).
     title: Reference title.
     author: Author
-    iocs: List of IOCs. Each IOC must be a dict with type and (id|url). E.g:
-      {'type': 'file', 'id': '4c3499f3cc4a4fdc7e67417e055891c78540282dccc57e37a01167dfe351b244'}
-      {'type': "url", 'url': "http://www.colasprint.com/_vti_log/upload.asp"},
-      {'type': "domain", 'id': "opsonew3org.sg"},
-      {'type': "ip_address", 'id': '8.8.8.8'}
     client: VirusTotal client.
+    iocs: Dict with the different IOCs to add to the reference.
 
+  Returns:
+    The new reference object.
   """
 
   # Generate url identifier
   payload = {
-      'data': {
-          'attributes': {
-              'url': url,
-              'creation_date': creation_date,
-              'title': title,
-              'author': author
-          },
-          'relationships': {},
-          'type': 'reference'
-      }
+      'attributes': {
+          'url': url,
+          'creation_date': creation_date,
+          'title': title,
+          'author': author
+      },
+      'relationships': {},
+      'type': 'reference',
+      'id': '',
   }
 
   # Add IOCs to Reference.
   add_iocs_to_reference_payload(iocs, payload)
 
-  # Post object
-  client.post('/references', data = json.dumps(payload))
+  reference_obj = vt.Object.from_dict(payload)
+  return client.post_object('/references', obj=reference_obj)
+
 
 def add_iocs_to_reference_payload(iocs, reference_payload):
   """Adds IOCs relationships to a given reference.
 
   Args:
-    iocs: List of IOCs. Each IOC must be a dict with type and (id|url). E.g:
-      {'type': 'file', 'id': '4c3499f3cc4a4fdc7e67417e055891c78540282dccc57e37a01167dfe351b244'}
-      {'type': "url", 'url': "http://www.colasprint.com/_vti_log/upload.asp"},
-      {'type': "domain", 'id': "opsonew3org.sg"},
-      {'type': "ip_address", 'id': '8.8.8.8'}
+    iocs: Dict with the different IOCs to add to the reference.
     reference_payload: Reference payload
-
   """
+  for relationship_name in ['files', 'domains', 'urls', 'ip_addresses']:
+    if relationship_name not in iocs:
+      continue
+    if relationship_name == 'urls':
+      descriptors = [{'type': 'url', 'url': u} for u in iocs['urls']]
+    else:
+      type_name = (
+          'ip_address' if relationship_name == 'ip_addresses'
+          else relationship_name[:-1])
+      descriptors = [
+          {'type': type_name, 'id': i} for i in iocs[relationship_name]]
+
+    reference_payload['relationships'][relationship_name] = {
+        'data': descriptors}
 
-  # Groups ioc by types
-  files = [ioc for ioc in iocs if ioc['type'] == 'file']
-  domains = [ioc for ioc in iocs if ioc['type'] == 'domain']
-  urls = [ioc for ioc in iocs if ioc['type'] == 'url']
-  ip_addresses = [ioc for ioc in iocs if ioc['type'] == 'ip_address']
-
-  relationship_items = [files, domains, urls, ip_addresses]
-  relationship_names = ['files', 'domains', 'urls', 'ip_addresses']
-
-  # Iterate through all relationship types.
-  for relationships, relationship_name in zip(
-      relationship_items, relationship_names):
-
-    reference_payload['data']['relationships'][relationship_name] = relationships
 
 def main():
   parser = argparse.ArgumentParser(
@@ -95,58 +88,67 @@ def main():
   parser.add_argument('--apikey', required=True, help='your VirusTotal API key')
 
   args = parser.parse_args()
-  API_KEY = args.apikey
-
-  url = 'https://blog.google/threat-analysis-group/new-campaign-targeting-security-researchers/'
-
-  client = vt.Client(API_KEY)
-
-  # IOCs must specify their type and ID.
-  # For urls, instead of "id", "url" field is specified.
-  # Allowed types : ['file', 'domain', 'url', ip_address']
-  iocs = [
-      {'type': 'file', 'id': '4c3499f3cc4a4fdc7e67417e055891c78540282dccc57e37a01167dfe351b244'},
-      {'type': 'file', 'id': '68e6b9d71c727545095ea6376940027b61734af5c710b2985a628131e47c6af7'},
-      {'type': 'file', 'id': '25d8ae4678c37251e7ffbaeddc252ae2530ef23f66e4c856d98ef60f399fa3dc'},
-      {'type': 'file', 'id': 'a75886b016d84c3eaacaf01a3c61e04953a7a3adf38acf77a4a2e3a8f544f855'},
-      {'type': 'file', 'id': 'a4fb20b15efd72f983f0fb3325c0352d8a266a69bb5f6ca2eba0556c3e00bd15'},
-      {'type': 'url', 'url': 'a4fb20b15efd72f983f0fb3325c0352d8a266a69bb5f6ca2eba0556c3e00bd15'},
-      {'type': "url", 'url': "https://angeldonationblog.com/image/upload/upload.php"},
-      {'type': "url", 'url': "https://codevexillium.org/image/download/download.asp"},
-      {'type': "url", 'url': "https://investbooking.de/upload/upload.asp"},
-      {'type': "url", 'url': "https://transplugin.io/upload/upload.asp"},
-      {'type': "url", 'url': "https://www.dronerc.it/forum/uploads/index.php"},
-      {'type': "url", 'url': "https://www.dronerc.it/shop_testbr/Core/upload.php"},
-      {'type': "url", 'url': "https://www.dronerc.it/shop_testbr/upload/upload.php"},
-      {'type': "url", 'url': "https://www.edujikim.com/intro/blue/insert.asp"},
-      {'type': "url", 'url': "https://www.fabioluciani.com/es/include/include.asp"},
-      {'type': "url", 'url': "http://trophylab.com/notice/images/renewal/upload.asp"},
-      {'type': "url", 'url': "http://www.colasprint.com/_vti_log/upload.asp"},
-      {'type': "domain", 'id': "angeldonationblog.com"},
-      {'type': "domain", 'id': "codevexillium.org"},
-      {'type': "domain", 'id': "investbooking.de"},
-      {'type': "domain", 'id': "krakenfolio.com"},
-      {'type': "domain", 'id': "opsonew3org.sg"},
-      {'type': "domain", 'id': "transferwiser.io"},
-      {'type': "domain", 'id': "transplugin.io"},
-      {'type': "domain", 'id': "trophylab.com"},
-      {'type': "domain", 'id': "www.colasprint.com"},
-      {'type': "domain", 'id': "www.dronerc.it"},
-      {'type': "domain", 'id': "www.edujikim.com"},
-      {'type': "domain", 'id': "www.fabioluciani.com"}
-  ]
+  client = vt.Client(args.apikey)
+
+  # Reference's URL.
+  url = (
+      'https://blog.google/threat-analysis-group/'
+      'new-campaign-targeting-security-researchers/')
+
+  # Fill in the reference's IOCs
+  # Allowed types: files, domains, urls, ip_addresses.
+  iocs = {
+      'files': [
+          '4c3499f3cc4a4fdc7e67417e055891c78540282dccc57e37a01167dfe351b244',
+          '68e6b9d71c727545095ea6376940027b61734af5c710b2985a628131e47c6af7',
+          '25d8ae4678c37251e7ffbaeddc252ae2530ef23f66e4c856d98ef60f399fa3dc',
+          'a75886b016d84c3eaacaf01a3c61e04953a7a3adf38acf77a4a2e3a8f544f855',
+          'a4fb20b15efd72f983f0fb3325c0352d8a266a69bb5f6ca2eba0556c3e00bd15',
+      ],
+      'urls': [
+          'https://angeldonationblog.com/image/upload/upload.php',
+          'https://codevexillium.org/image/download/download.asp',
+          'https://investbooking.de/upload/upload.asp',
+          'https://transplugin.io/upload/upload.asp',
+          'https://www.dronerc.it/forum/uploads/index.php',
+          'https://www.dronerc.it/shop_testbr/Core/upload.php',
+          'https://www.dronerc.it/shop_testbr/upload/upload.php',
+          'https://www.edujikim.com/intro/blue/insert.asp',
+          'https://www.fabioluciani.com/es/include/include.asp',
+          'http://trophylab.com/notice/images/renewal/upload.asp',
+          'http://www.colasprint.com/_vti_log/upload.asp',
+      ],
+      'domains': [
+          'angeldonationblog.com'
+          'codevexillium.org',
+          'investbooking.de',
+          'krakenfolio.com',
+          'opsonew3org.sg',
+          'transferwiser.io',
+          'transplugin.io',
+          'trophylab.com',
+          'www.colasprint.com',
+          'www.dronerc.it',
+          'www.edujikim.com',
+          'www.fabioluciani.com',
+      ],
+      'ip_addresses': [
+          '193.70.64.169',
+      ]
+  }
 
   # Create Reference
-  create_reference(
+  reference_obj = create_reference(
       url=url,
-      creation_date="2021-01-25 00:00:00",
-      title="New campaign targeting security researchers",
-      author="Google Threat Analysis Group",
+      creation_date='2021-01-25 00:00:00',
+      title='New campaign targeting security researchers',
+      author='Google Threat Analysis Group',
       iocs=iocs,
       client=client
   )
 
   client.close()
+  print(json.dumps(reference_obj.to_dict(), indent=2))
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   main()
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -11,49 +11,57 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Client tests."""
+
 import bz2
 import datetime
 import io
 import json
 
 import pytest
-import pytest_httpserver
 
 from vt import Client
 from vt import FeedType
 from vt import Object
 
 
 def new_client(httpserver):
-  return Client('dummy_api_key',
+  return Client(
+      'dummy_api_key',
       host='http://' + httpserver.host + ':' + str(httpserver.port))
 
 
 def test_object_from_dict():
-
   obj = Object.from_dict({
       'type': 'dummy_type',
       'id': 'dummy_id',
       'attributes': {
           'attr1': 'foo',
           'attr2': 1,
+      },
+      'relationships': {
+          'foos': {'data': [
+              {'type': 'foo', 'id': 'foo_id'}
+          ]}
       }})
 
   assert obj.id == 'dummy_id'
   assert obj.type == 'dummy_type'
   assert obj.attr1 == 'foo'
   assert obj.attr2 == 1
+  assert obj.relationships['foos']['data'][0]['id'] == 'foo_id'
 
-  with pytest.raises(ValueError, match=r"Expecting dictionary, got: int"):
+  with pytest.raises(ValueError, match=r'Expecting dictionary, got: int'):
     Object.from_dict(1)
 
-  with pytest.raises(ValueError, match=r"Object type not found"):
+  with pytest.raises(ValueError, match=r'Object type not found'):
     Object.from_dict({})
 
-  with pytest.raises(ValueError, match=r"Object id not found"):
+  with pytest.raises(ValueError, match=r'Object id not found'):
     Object.from_dict({'type': 'dummy_type'})
 
-  with pytest.raises(ValueError, match=r'Object attributes must be a dictionary'):
+  with pytest.raises(
+      ValueError, match=r'Object attributes must be a dictionary'):
     Object.from_dict({'type': 'dummy_type', 'id': 'dummy_id', 'attributes': 1})
 
 
@@ -154,8 +162,7 @@ def test_get_object(httpserver):
           'type': 'dummy_type',
           'attributes': {
               'foo': 'foo',
-              'bar': 'bar'}
-  }})
+              'bar': 'bar'}}})
 
   with new_client(httpserver) as client:
     obj = client.get_object('/dummy_types/dummy_id')
@@ -169,6 +176,7 @@ def test_get_object(httpserver):
   assert obj.get('bar') == 'bar'
   assert obj.get('baz') is None
 
+
 def test_patch_object(httpserver):
 
   obj = Object('dummy_type', 'dummy_id', {'foo': 1, 'bar': 2})
@@ -298,7 +306,7 @@ def test_scan_file(httpserver):
   })
 
   with new_client(httpserver) as client:
-    f = io.StringIO("dummy file")
+    f = io.StringIO('dummy file')
     analysis = client.scan_file(f)
 
   assert analysis.type == 'analysis'
@@ -333,7 +341,7 @@ def test_feed(httpserver):
       method='GET',
       headers={'X-Apikey': 'dummy_api_key'}
   ).respond_with_data(
-    bz2.compress(b'{\"type\": \"file\", \"id\": \"dummy_file_id_1\"}'))
+      bz2.compress(b'{\"type\": \"file\", \"id\": \"dummy_file_id_1\"}'))
 
   # The feed iterator should tolerate missing feed packages, so let's return
   # a NotFoundError for package 200102030406.
@@ -343,15 +351,14 @@ def test_feed(httpserver):
       headers={'X-Apikey': 'dummy_api_key'}
   ).respond_with_json({
       'error': {
-          'code': 'NotFoundError'
-    }}, status=404)
+          'code': 'NotFoundError'}}, status=404)
 
   httpserver.expect_ordered_request(
       '/api/v3/feeds/files/200102030407',
       method='GET',
       headers={'X-Apikey': 'dummy_api_key'}
   ).respond_with_data(
-    bz2.compress(b'{\"type\": \"file\", \"id\": \"dummy_file_id_2\"}'))
+      bz2.compress(b'{\"type\": \"file\", \"id\": \"dummy_file_id_2\"}'))
 
   with new_client(httpserver) as client:
     feed = client.feed(FeedType.FILES, cursor='200102030405')

diff --git a/vt/object.py b/vt/object.py
@@ -201,6 +201,9 @@ def to_dict(self, modified_attributes_only=False):
     if attributes:
       result['attributes'] = attributes
 
+    if self.relationships:
+      result['relationships'] = self.relationships
+
     if self.context_attributes:
       result['context_attributes'] = self.context_attributes
 

diff --git a/vt/version.py b/vt/version.py
@@ -1 +1 @@
-__version__ = '0.6.1'
+__version__ = '0.6.2'