Merge remote-tracking branch 'upstream/master'

Author: mbengit

qpylib/ariel.py

@@ -108,6 +108,37 @@ def results(self, search_id, start=0, end=0, api_version='latest'):
                              .format(search_id, response.content))
         return response.json()
 
+    def delete(self, search_id, api_version='latest'):
+        ''' Deletes a previous Ariel search.
+              search_id: Ariel search ID
+              api_version: QRadar API version to use, defaults to latest.
+            Returns a tuple containing search status and search ID.
+            Raises ArielError if the search could not be deleted.
+        '''
+        response = qpylib.REST('DELETE', ArielSearch.SEARCH_ENDPOINT.format(search_id),
+                               headers=self._build_headers(api_version))
+        if response.status_code != 200:
+            raise ArielError('Ariel search {0} could not be deleted: HTTP {1} was returned'
+                             .format(search_id, response.status_code))
+        response_json = response.json()
+        return (response_json['status'], response_json['search_id'])
+
+    def cancel(self, search_id, api_version='latest'):
+        ''' Cancels an ongoing Ariel search.
+              search_id: Ariel search ID.
+              api_version: QRadar API version to use, defaults to latest.
+            Returns a tuple containing search status and search ID.
+            Raises ArielError if the search could not be cancelled.
+        '''
+        response = qpylib.REST('POST', ArielSearch.SEARCH_ENDPOINT.format(search_id),
+                               headers=self._build_headers(api_version),
+                               params={'status': 'CANCELLED'})
+        if response.status_code != 200:
+            raise ArielError('Ariel search {0} could not be cancelled: HTTP {1} was returned'
+                             .format(search_id, response.status_code))
+        response_json = response.json()
+        return (response_json['status'], response_json['search_id'])
+
     @staticmethod
     def _build_headers(api_version):
         headers = {'Accept': 'application/json'}

test/test_ariel.py

@@ -13,6 +13,8 @@
 POST_SEARCH = '{0}?query_expression=select+stuff+from+db'.format(ARIEL_URL)
 GET_SEARCH = '{0}/{1}'.format(ARIEL_URL, SEARCH_ID)
 GET_RESULTS = '{0}/{1}/results'.format(ARIEL_URL, SEARCH_ID)
+DELETE_SEARCH = '{0}/{1}'.format(ARIEL_URL, SEARCH_ID)
+CANCEL_SEARCH = '{0}/{1}?status=CANCELLED'.format(ARIEL_URL, SEARCH_ID)
 
 @pytest.fixture(scope='module', autouse=True)
 def pre_testing_setup():
@@ -153,3 +155,37 @@ def test_results_range_start_and_end():
     assert results_json['result3'] == 42
     assert results_json['result4'] == 99
     assert responses.calls[0].request.headers['Range'] == 'items=3-4'
+
+@responses.activate
+def test_search_delete_failure():
+    responses.add('DELETE', DELETE_SEARCH, status=404,
+                  json={})
+    with pytest.raises(ArielError, match='Ariel search {0} could not be deleted'
+                       .format(SEARCH_ID)):
+        ArielSearch().delete(SEARCH_ID)
+
+@responses.activate
+def test_search_delete_success():
+    responses.add('DELETE', DELETE_SEARCH, status=200,
+                  json={'status': 'COMPLETED', 'search_id': SEARCH_ID})
+    status, search_id = ArielSearch().delete(SEARCH_ID)
+
+    assert status == 'COMPLETED'
+    assert search_id == SEARCH_ID
+
+@responses.activate
+def test_search_cancel_failure():
+    responses.add('POST', CANCEL_SEARCH, status=404,
+                  json={})
+    with pytest.raises(ArielError, match='Ariel search {0} could not be cancelled'
+                       .format(SEARCH_ID)):
+        ArielSearch().cancel(SEARCH_ID)
+
+@responses.activate
+def test_search_cancel_success():
+    responses.add('POST', CANCEL_SEARCH, status=200,
+                  json={'status': 'COMPLETED', 'search_id': SEARCH_ID})
+    status, search_id = ArielSearch().cancel(SEARCH_ID)
+
+    assert status == 'COMPLETED'
+    assert search_id == SEARCH_ID