Increased coverage final

mushorg · Mar 24, 2019 · 127f99e · 127f99e
1 parent 8186a68
commit 127f99e
Showing 1 changed file with 60 additions and 18 deletions.
diff --git a/snare/tests/test_parse_tanner_response.py b/snare/tests/test_parse_tanner_response.py
@@ -24,12 +24,12 @@ def setUp(self):
             f.write(self.page_content)
         with open(os.path.join(self.main_page_path, "meta.json"), 'w') as f:
             json.dump(meta_content, f)
-        args = run_args.parse_args(['--page-dir', page_dir])
-        args.index_page = '/index.html'
-        args.no_dorks = True
-        args.tanner = "tanner.mushmush.org"
-        uuid = "test_uuid"
-        self.handler = TannerHandler(args, meta_content, uuid)
+        self.args = run_args.parse_args(['--page-dir', page_dir])
+        self.args.index_page = '/index.html'
+        self.args.no_dorks = True
+        self.args.tanner = "tanner.mushmush.org"
+        self.uuid = "test_uuid"
+        self.handler = TannerHandler(self.args, meta_content, self.uuid)
         self.requested_name = '/'
         self.loop = asyncio.get_event_loop()
         self.handler.html_handler.handle_content = AsyncMock(return_value=self.page_content)
@@ -41,6 +41,7 @@ def setUp(self):
         self.expected_content = None
         self.call_content = None
         self.expected_header = None
+        self.expected_status_code = None
 
     def test_parse_type_one(self):
         self.detection = {"type": 1}
@@ -55,6 +56,36 @@ async def test():
         expected_result = [self.page_content, self.content_type, {}, 200]
         self.assertCountEqual(real_result, expected_result)
 
+    def test_parse_type_one_query(self):
+        self.requested_name = '/?'
+        self.detection = {"type": 1}
+        self.call_content = b'<html><body></body></html>'
+
+        async def test():
+            (self.res1, self.res2,
+             self.res3, self.res4) = await self.handler.parse_tanner_response(self.requested_name, self.detection)
+
+        self.loop.run_until_complete(test())
+        real_result = [self.res1, self.res2, self.res3, self.res4]
+        expected_result = [self.page_content, self.content_type, {}, 200]
+        self.assertCountEqual(real_result, expected_result)
+
+    def test_parse_type_one_error(self):
+        self.detection = {"type": 1}
+        meta_content = {"/index.html": {}}
+        self.requested_name = 'something/'
+        self.handler = TannerHandler(self.args, meta_content, self.uuid)
+        self.expected_status_code = 404
+
+        async def test():
+            (self.res1, self.res2,
+             self.res3, self.res4) = await self.handler.parse_tanner_response(self.requested_name, self.detection)
+
+        self.loop.run_until_complete(test())
+        real_result = [self.res1, self.res2, self.res3, self.res4]
+        expected_result = [None, None, {}, self.expected_status_code]
+        self.assertCountEqual(real_result, expected_result)
+
     def test_parse_type_two(self):
         self.detection = {
             "type": 2,
@@ -99,6 +130,29 @@ async def test():
         expected_result = [self.expected_content, self.content_type, self.expected_header, 200]
         self.assertCountEqual(real_result, expected_result)
 
+    def test_parse_type_two_error(self):
+
+        self.detection = {
+            "type": 2,
+            "payload": {
+                "page": "/index.html",
+                "value": "test"
+            }
+        }
+        meta_content = {"/index.html": {}}
+        self.handler = TannerHandler(self.args, meta_content, self.uuid)
+        self.expected_content = b'<html><body><div>test</div></body></html>'
+        self.content_type = r'text\html'
+
+        async def test():
+            (self.res1, self.res2,
+             self.res3, self.res4) = await self.handler.parse_tanner_response(self.requested_name, self.detection)
+
+        self.loop.run_until_complete(test())
+        real_result = [self.res1, self.res2, self.res3, self.res4]
+        expected_result = [self.expected_content, self.content_type, {}, 200]
+        self.assertCountEqual(real_result, expected_result)
+
     def test_parse_type_three(self):
         self.detection = {
             "type": 3,
@@ -131,17 +185,5 @@ async def test():
         self.loop.run_until_complete(test())
         self.handler.html_handler.handle_content.assert_called_with(self.call_content)
 
-    def test_parse_exception(self):
-        self.detection = {}
-        self.call_content = b'<html><body></body></html>'
-        self.expected_content = self.page_content
-
-        async def test():
-            (self.res1, self.res2,
-             self.res3, self.res4) = await self.handler.parse_tanner_response(self.requested_name, self.detection)
-
-        with self.assertRaises(KeyError):
-            self.loop.run_until_complete(test())
-
     def tearDown(self):
         shutil.rmtree(self.main_page_path)