examples, src: Merge ThAPI and ThFeedAPI

woodruffw · woodruffw · commit b524be3dbdb1 · 2019-03-06T11:12:46.000-05:00
diff --git a/examples/threathunter/create_feed.py b/examples/threathunter/create_feed.py
@@ -6,7 +6,7 @@
 import validators
 import hashlib
 
-from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_feed_object
+from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_object
 from cbapi.psc.threathunter import Feed
 
 
@@ -35,7 +35,7 @@ def main():
     parser.add_argument("--rep_visibility", type=str, help="Report visibility")
 
     args = parser.parse_args()
-    cb = get_cb_threathunter_feed_object(args)
+    cb = get_cb_threathunter_object(args)
 
     feed_info = {
         "name": args.name,
diff --git a/examples/threathunter/feed_operations.py b/examples/threathunter/feed_operations.py
@@ -3,7 +3,7 @@
 
 import sys
 from cbapi.psc.threathunter.models import Feed, Report
-from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_feed_object
+from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_object
 import logging
 import json
 
@@ -206,7 +206,7 @@ def main():
     specifier.add_argument("-f", "--feedname", type=str, help="Feed Name")
 
     args = parser.parse_args()
-    cb = get_cb_threathunter_feed_object(args)
+    cb = get_cb_threathunter_object(args)
 
     if args.command_name == "list":
         return list_feeds(cb, parser, args)
diff --git a/examples/threathunter/watchlist_operations.py b/examples/threathunter/watchlist_operations.py
@@ -3,7 +3,7 @@
 
 import sys
 from cbapi.psc.threathunter.models import Watchlist, Report, Feed
-from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_feed_object
+from cbapi.example_helpers import build_cli_parser, get_cb_threathunter_object
 from cbapi.errors import ObjectNotFoundError
 import logging
 import json
@@ -329,7 +329,7 @@ def main():
     commands.add_parser("import", help="Import a previously exported watchlist")
 
     args = parser.parse_args()
-    cb = get_cb_threathunter_feed_object(args)
+    cb = get_cb_threathunter_object(args)
 
     if args.command_name == "list":
         return list_watchlists(cb, parser, args)
diff --git a/src/cbapi/example_helpers.py b/src/cbapi/example_helpers.py
@@ -12,7 +12,7 @@
 
 from cbapi.protection import CbEnterpriseProtectionAPI
 from cbapi.psc.defense import CbDefenseAPI
-from cbapi.psc.threathunter import CbThreatHunterAPI, CbThreatHunterFeedAPI
+from cbapi.psc.threathunter import CbThreatHunterAPI
 from cbapi.response import CbEnterpriseResponseAPI
 
 log = logging.getLogger(__name__)
@@ -102,21 +102,6 @@ def get_cb_threathunter_object(args):
     return cb
 
 
-def get_cb_threathunter_feed_object(args):
-    if args.verbose:
-        import logging
-        logging.basicConfig()
-        logging.getLogger("cbapi").setLevel(logging.DEBUG)
-        logging.getLogger("__main__").setLevel(logging.DEBUG)
-
-    if args.cburl and args.apitoken:
-        cb = CbThreatHunterFeedAPI(url=args.cburl, token=args.apitoken, ssl_verify=(not args.no_ssl_verify))
-    else:
-        cb = CbThreatHunterFeedAPI(profile=args.profile)
-
-    return cb
-
-
 def get_object_by_name_or_id(cb, cls, name_field="name", id=None, name=None, force_init=True):
     clsname = cls.__name__
     try:
diff --git a/src/cbapi/psc/threathunter/__init__.py b/src/cbapi/psc/threathunter/__init__.py
@@ -7,4 +7,3 @@
 from .rest_api import CbThreatHunterAPI
 from cbapi.psc.threathunter.models import Process, Event, Tree, Feed, Report, IOCs, IOC_V2, Watchlist
 from cbapi.psc.threathunter.query import QueryBuilder
-from .feed_api import CbThreatHunterFeedAPI
diff --git a/src/cbapi/psc/threathunter/feed_api.py b/src/cbapi/psc/threathunter/feed_api.py
diff --git a/src/cbapi/psc/threathunter/rest_api.py b/src/cbapi/psc/threathunter/rest_api.py
@@ -1,5 +1,6 @@
 from cbapi.psc.threathunter.query import Query
 from cbapi.connection import BaseAPI
+from cbapi.psc.threathunter.models import ReportSeverity
 import logging
 
 log = logging.getLogger(__name__)
@@ -26,6 +27,55 @@ def _perform_query(self, cls, **kwargs):
         else:
             return Query(cls, self, **kwargs)
 
+    def create(self, cls, data=None):
+        """Creates a new model.
+
+        >>> feed = cb.create(Feed, feed_data)
+
+        :param cls: The model being created
+        :param data: The data to pre-populate the model with
+        :type data: dict(str, object)
+        :return: an instance of `cls`
+        """
+        return cls(self, initial_data=data)
+
+    def validate_query(self, query):
+        """Validates the given IOC query.
+
+        >>> cb.validate_query("process_name:chrome.exe") # True
+
+        :param str query: the query to validate
+        :return: whether or not the query is valid
+        :rtype: bool
+        """
+        args = {"q": query}
+        resp = self.get_object("/pscr/query/v1/validate", query_parameters=args)
+
+        return resp.get("valid", False)
+
+    def convert_query(self, query):
+        """Converts a legacy CB Response query to a ThreatHunter query.
+
+        :param str query: the query to convert
+        :return: the converted query
+        :rtype: str
+        """
+        args = {"query": query}
+        resp = self.get_object("/threathunter/feedmgr/v1/query/translate", query_parameters=args)
+
+        return resp.get("query")
+
+    @property
+    def custom_severities(self):
+        """Returns a list of active :py:class:`ReportSeverity` instances
+
+        :rtype: list[:py:class:`ReportSeverity`]
+        """
+        # TODO(ww): There's probably a better place to put this.
+        resp = self.get_object("/threathunter/watchlistmgr/v1/severity")
+        items = resp.get("results", [])
+        return [self.create(ReportSeverity, item) for item in items]
+
     def queries(self):
         """Retrieves a list of queries, active or complete, known by
         the ThreatHunter server.
