Fixed FireEye test

InQuest · Dec 9, 2022 · 232341d · 232341d
1 parent 06c2331
commit 232341d
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 37 deletions.
diff --git a/sandboxapi/fireeye.py b/sandboxapi/fireeye.py
@@ -144,25 +144,22 @@ def is_available(self):
         :rtype:  bool
         :return: True if service is available, False otherwise.
         """
-        # if the availability flag is raised, return True immediately.
-        # NOTE: subsequent API failures will lower this flag. we do this here
-        # to ensure we don't keep hitting FireEye with requests while
-        # availability is there.
-        if self.server_available:
-            return True
-
-        # otherwise, we have to check with the cloud.
-        else:
-            try:
-                response = self._request("/config")
 
-                # we've got fireeye.
-                if response.status_code == 200:
-                    self.server_available = True
-                    return True
+        try:
+            response = self._request("/config")
 
-            except sandboxapi.SandboxError:
-                pass
+            # Successfully connected to FireEye
+            if response.status_code == 200:
+                self.server_available = True
+                return True
+
+            # Unable to connect to FireEye
+            if response.status_code >= 500:
+                self.server_available = False
+                return False
+
+        except sandboxapi.SandboxError:
+            pass
 
         self.server_available = False
         return False

diff --git a/tests/test_fireeye.py b/tests/test_fireeye.py
@@ -7,11 +7,11 @@
     from mock import patch, ANY as MOCK_ANY
 
 import responses
-from sandboxapi.fireeye import FireEyeAPI
+import sandboxapi.fireeye
 from . import read_resource
 
 class TestFireEye(TestCase):
-    sandbox = FireEyeAPI('username', 'password', 'http://fireeye.mock', 'profile')
+    sandbox = sandboxapi.fireeye.FireEyeAPI('username', 'password', 'http://fireeye.mock', 'profile')
 
     @responses.activate
     def test_analyze(self):
@@ -37,14 +37,14 @@ def test_is_available(self):
                       json=read_resource('fireeye_config'))
         self.assertTrue(self.sandbox.is_available())
 
-    # @responses.activate
-    # def test_not_is_available(self):
-    #     self.assertFalse(self.sandbox.is_available())
-    #     responses.add(responses.POST, 'http://fireeye.mock/wsapis/v1.2.0/auth/login',
-    #                   headers={'X-FeApi-Token': 'MOCK'})
-    #     responses.add(responses.GET, 'http://fireeye.mock/wsapis/v1.2.0/config',
-    #                   status=500)
-    #     self.assertFalse(self.sandbox.is_available())
+    @responses.activate
+    def test_not_is_available(self):
+        self.assertFalse(self.sandbox.is_available())
+        responses.add(responses.POST, 'http://fireeye.mock/wsapis/v1.2.0/auth/login',
+                      headers={'X-FeApi-Token': 'MOCK'})
+        responses.add(responses.GET, 'http://fireeye.mock/wsapis/v1.2.0/config',
+                      status=500)
+        self.assertFalse(self.sandbox.is_available())
 
     @responses.activate
     def test_report(self):
@@ -77,7 +77,7 @@ def test_proxies_is_passed_to_requests(self, m_get, m_post):
             'https': 'http://10.10.1.10:1080',
         }
 
-        api = FireEyeAPI('username', 'password',
+        api = sandboxapi.fireeye.FireEyeAPI('username', 'password',
                                             self.sandbox.api_url, 'profile',
                                             proxies=proxies)
         api._request('/test')
@@ -114,7 +114,7 @@ def test_reauthenticates_if_logged_out_json_401(self):
         self.assertEqual(self.sandbox.check('1'), True)
 
 class TestFireEyeLegacy(TestCase):
-    legacy_sandbox = FireEyeAPI('username', 'password', 'http://fireeye.mock', 'profile', legacy_api=True)
+    legacy_sandbox = sandboxapi.fireeye.FireEyeAPI('username', 'password', 'http://fireeye.mock', 'profile', legacy_api=True)
 
     # Legacy API support
     @responses.activate
@@ -141,14 +141,14 @@ def legacy_test_is_available(self):
                       json=read_resource('fireeye_config'))
         self.assertTrue(self.legacy_sandbox.is_available())
 
-    # @responses.activate
-    # def legacy_test_not_is_available(self):
-    #     self.assertFalse(self.legacy_sandbox.is_available())
-    #     responses.add(responses.POST, 'http://fireeye.mock/wsapis/v1.1.0/auth/login',
-    #                   headers={'X-FeApi-Token': 'MOCK'})
-    #     responses.add(responses.GET, 'http://fireeye.mock/wsapis/v1.1.0/config',
-    #                   status=500)
-    #     self.assertFalse(self.legacy_sandbox.is_available())
+    @responses.activate
+    def legacy_test_not_is_available(self):
+        self.assertFalse(self.legacy_sandbox.is_available())
+        responses.add(responses.POST, 'http://fireeye.mock/wsapis/v1.1.0/auth/login',
+                      headers={'X-FeApi-Token': 'MOCK'})
+        responses.add(responses.GET, 'http://fireeye.mock/wsapis/v1.1.0/config',
+                      status=500)
+        self.assertFalse(self.legacy_sandbox.is_available())
 
     @responses.activate
     def legacy_test_report(self):