Extending support for additional feeds endpoints

examples/url_feed.py → examples/get_feed.py

@@ -43,18 +43,20 @@ def main():
 
   parser.add_argument("--apikey", required=True, help="your VirusTotal API key")
 
+  parser.add_argument("--feed_type", required=True, help="The feedtype to get", type=vt.FeedType, choices=list(vt.FeedType))
+
   parser.add_argument(
       "--cursor", required=False, help="cursor indicating where to start"
   )
 
   args = parser.parse_args()
 
   with vt.Client(args.apikey) as client:
-    # Iterate over the URL feed, one file at a time. This loop doesn't
+    # Iterate over the provided feed, one file at a time. This loop doesn't
     # finish, when the feed is consumed it will keep waiting for more files.
 
     try:
-      for url_obj in client.feed(vt.FeedType.URLS, cursor=args.cursor):
+      for url_obj in client.feed(args.feed_type, cursor=args.cursor):
         # process the url_obj
         process_item(url_obj)
     except KeyboardInterrupt:

vt/feed.py

@@ -17,7 +17,7 @@
 import asyncio
 import bz2
 from datetime import datetime
-from datetime import timedelta
+from datetime import timedelta, timezone
 import enum
 import io
 import json
@@ -37,7 +37,9 @@
 class FeedType(enum.Enum):
   """Feed types."""
 
+  DOMAINS = "domains"
   FILES = "files"
+  IPS = "ipaddresses"
   URLS = "urls"
   FILE_BEHAVIOURS = "file-behaviours"
 
@@ -87,7 +89,7 @@ def __init__(
       self._batch_time = datetime.strptime(batch_time, "%Y%m%d%H%M")
       self._batch_skip = int(batch_skip) if batch_skip else 0
     else:
-      self._batch_time = datetime.utcnow() - timedelta(minutes=70)
+      self._batch_time = datetime.now(timezone.utc) - timedelta(minutes=70)
       self._batch_skip = 0
 
     self._next_batch_time = self._batch_time