Aiodocker Helper (#341)

* Aiodocker helper

* Aiodocker tests

* deleting used containers

* close docker removed

requirements.txt

@@ -12,3 +12,4 @@ pylibinjection
 jinja2                                                                                                                                                                                                                     
 pycodestyle                                                                                                                                                                                                                
 geoip2
+aiodocker

tanner/emulators/cmd_exec.py

@@ -1,20 +1,15 @@
-from tanner.utils import docker_helper
+from tanner.utils import aiodocker_helper
 from tanner.utils import patterns
 
 
 class CmdExecEmulator:
     def __init__(self):
-        self.helper = docker_helper.DockerHelper()
+        self.helper = aiodocker_helper.AIODockerHelper()
 
-    async def create_attacker_env(self, session):
-        container_name = 'attacker_' + session.sess_uuid.hex
-        container = await self.helper.create_container(container_name)
-        if container:
-            session.associate_env(container_name)
-        return container
+    async def get_cmd_exec_results(self, payload):
+        cmd = ["sh", "-c", payload]
 
-    async def get_cmd_exec_results(self, container, cmd):
-        execute_result = await self.helper.execute_cmd(container, cmd)
+        execute_result = await self.helper.execute_cmd(cmd)
         result = dict(value=execute_result, page=True)
         return result
 
@@ -25,6 +20,6 @@ def scan(self, value):
         return detection
 
     async def handle(self, attack_params, session=None):
-        container = await self.create_attacker_env(session)
-        result = await self.get_cmd_exec_results(container, attack_params[0]['value'])
+
+        result = await self.get_cmd_exec_results(attack_params[0]['value'])
         return result

tanner/emulators/lfi.py

@@ -1,28 +1,25 @@
 import shlex
 
-from tanner.utils import docker_helper
+from tanner.utils import aiodocker_helper
 from tanner.utils import patterns
 
 
 class LfiEmulator:
     def __init__(self):
-        self.helper = docker_helper.DockerHelper()
+        self.helper = aiodocker_helper.AIODockerHelper()
 
-    async def get_lfi_result(self, container, file_path):
+    async def get_lfi_result(self, file_path):
         # Terminate the string with NULL byte
         if '\x00' in file_path:
             file_path = file_path[:file_path.find('\x00')]
 
-        cmd = 'cat {file}'.format(file=shlex.quote(file_path))
-        execute_result = await self.helper.execute_cmd(container, cmd)
-        # Nulls are not printable, so replace it with another line-ender
-        execute_result = execute_result.replace('\x00', '\n')
-        return execute_result
+        cmd = ["sh", "-c", "cat {file}".format(file=shlex.quote(file_path))]
+        execute_result = await self.helper.execute_cmd(cmd)
 
-    async def setup_virtual_env(self):
-        container_name = 'lfi_container'
-        container = await self.helper.create_container(container_name)
-        return container
+        if execute_result:
+            # Nulls are not printable, so replace it with another line-ender
+            execute_result = execute_result.replace('\x00', '\n')
+        return execute_result
 
     def scan(self, value):
         detection = None
@@ -31,9 +28,7 @@ def scan(self, value):
         return detection
 
     async def handle(self, attack_params, session=None):
-        result = None
-        container = await self.setup_virtual_env()
-        if container:
-            lfi_result = await self.get_lfi_result(container, attack_params[0]['value'])
-            result = dict(value=lfi_result, page=False)
+
+        lfi_result = await self.get_lfi_result(attack_params[0]['value'])
+        result = dict(value=lfi_result, page=False)
         return result

tanner/session.py

@@ -4,7 +4,7 @@
 from urllib.parse import urlparse
 
 from tanner.config import TannerConfig
-from tanner.utils import docker_helper
+from tanner.utils import aiodocker_helper
 from tanner.utils.mysql_db_helper import MySQLDBHelper
 from tanner.utils.sqlite_db_helper import SQLITEDBHelper
 
@@ -80,7 +80,7 @@ def associate_env(self, env):
         self.associated_env = env
 
     async def remove_associated_env(self):
-        await docker_helper.DockerHelper().delete_env(self.associated_env)
+        await aiodocker_helper.AIODockerHelper().delete_container(self.associated_env)
 
     def get_uuid(self):
         return str(self.sess_uuid)

tanner/tests/test_aiodocker_helper.py

@@ -0,0 +1,73 @@
+import asyncio
+import unittest
+
+from tanner.utils.aiodocker_helper import AIODockerHelper
+
+
+class TestAioDockerHelper(unittest.TestCase):
+    def setUp(self):
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+        self.handler = AIODockerHelper()
+        self.image = None
+        self.expected_result = None
+        self.returned_result = None
+
+    def test_setup_host_image(self):
+        self.image = 'busybox:latest'
+
+        async def test():
+            await self.handler.setup_host_image()
+            self.returned_result = await self.handler.docker_client.images.list(filter=self.image)
+
+        self.loop.run_until_complete(test())
+        self.expected_result = 1
+        self.assertEqual(len(self.returned_result), self.expected_result)
+
+    def test_get_container(self):
+        container_name = 'attacker_container'
+
+        async def test():
+            await self.handler.create_container(container_name)
+            self.returned_result = await self.handler.get_container(container_name)
+            await self.handler.delete_container(container_name)
+
+        self.loop.run_until_complete(test())
+        self.assertTrue(self.returned_result._id)
+
+    def test_create_container(self):
+        container_name = 'attacker'
+
+        async def test():
+            container = await self.handler.create_container(container_name=container_name)
+            await container.start()
+            self.returned_result = await container.show()
+            await self.handler.delete_container(container_name)
+
+        self.loop.run_until_complete(test())
+        self.assertTrue(self.returned_result["State"]["Running"])
+
+    def test_execute_cmd(self):
+        cmd = ["sh", "-c", "echo 'Hello!'"]
+
+        async def test():
+            self.returned_result = await self.handler.execute_cmd(cmd)
+
+        self.loop.run_until_complete(test())
+        self.expected_result = 'Hello!'
+        self.assertIn(self.expected_result, self.returned_result)
+
+    def test_delete_container(self):
+        container_name = 'attacker_z'
+
+        async def test():
+            await self.handler.create_container(container_name)
+            await self.handler.delete_container(container_name)
+            self.returned_result = await self.handler.get_container(container_name)
+
+        self.loop.run_until_complete(test())
+        self.assertEqual(self.returned_result, None)
+
+    def tearDown(self):
+        self.loop.run_until_complete(self.handler.docker_client.close())
+        self.loop.close()

tanner/tests/test_base.py

@@ -11,7 +11,7 @@
 class TestBase(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(None)
+        asyncio.set_event_loop(self.loop)
         self.session = mock.Mock()
         self.session.associate_db = mock.Mock()
         self.data = mock.Mock()

tanner/tests/test_cmd_exec_emulation.py

@@ -7,7 +7,7 @@
 class TestCmdExecEmulator(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(None)
+        asyncio.set_event_loop(self.loop)
         self.handler = cmd_exec.CmdExecEmulator()
         self.handler.helper.host_image = 'busybox:latest'
         self.sess = mock.Mock()
@@ -45,6 +45,11 @@ def test_handle_nested_commands(self):
 
     def test_handle_invalid_commands(self):
         attack_params = [dict(id='foo', value='foo')]
+
         result = self.loop.run_until_complete(self.handler.handle(attack_params, self.sess))
-        assert_result = 'foo: not found'
+        assert_result = 'sh: foo: not found'
         self.assertIn(assert_result, result['value'])
+
+    def tearDown(self):
+        self.loop.run_until_complete(self.handler.helper.docker_client.close())
+        self.loop.close()

tanner/tests/test_lfi_emulator.py

@@ -6,7 +6,7 @@
 class TestLfiEmulator(unittest.TestCase):
     def setUp(self):
         self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(None)
+        asyncio.set_event_loop(self.loop)
         self.handler = lfi.LfiEmulator()
         self.handler.helper.host_image = 'busybox:latest'
 
@@ -40,4 +40,9 @@ def test_handle_path_null_character(self):
     def test_handle_missing_lfi(self):
         attack_params = [dict(id='foo', value='../../../../../etc/bar')]
         result = self.loop.run_until_complete(self.handler.handle(attack_params))
-        self.assertIn('No such file or directory', result['value'])
+        assert_result = 'No such file or directory'
+        self.assertIn(assert_result, result['value'])
+
+    def tearDown(self):
+        self.loop.run_until_complete(self.handler.helper.docker_client.close())
+        self.loop.close()

tanner/utils/aiodocker_helper.py

@@ -0,0 +1,82 @@
+import aiodocker
+import logging
+
+from tanner.config import TannerConfig
+
+
+class AIODockerHelper:
+    def __init__(self):
+
+        self.logger = logging.getLogger('tanner.aiodocker_helper.AIODockerHelper')
+
+        self.docker_client = aiodocker.Docker()
+        self.host_image = TannerConfig.get('DOCKER', 'host_image')
+
+    async def setup_host_image(self, remote_path=None, tag=None):
+
+        try:
+            if remote_path and tag is not None:
+                params = {"tag": tag, "remote": remote_path}
+                await self.docker_client.images.build(**params)
+
+            image = await self.docker_client.images.list(filter=self.host_image)
+            if not image:
+                await self.docker_client.images.pull(self.host_image)
+
+        except aiodocker.exceptions.DockerError as docker_error:
+            self.logger.exception('Error while pulling %s image %s', self.host_image, docker_error)
+
+    async def get_container(self, container_name):
+        container = None
+        try:
+            container = await self.docker_client.containers.get(container=container_name)
+
+        except aiodocker.exceptions.DockerError as server_error:
+            self.logger.exception('Error while fetching %s container %s', container_name, server_error)
+        return container
+
+    async def create_container(self, container_name, cmd=None, image=None):
+        await self.setup_host_image()
+        container = None
+        if image is None:
+            image = self.host_image
+
+        config = {
+            "Cmd": cmd,
+            "Image": image,
+        }
+        try:
+            container = await self.docker_client.containers.create_or_replace(config=config, name=container_name)
+
+        except (aiodocker.exceptions.DockerError or aiodocker.exceptions.DockerContainerError) as docker_error:
+            self.logger.exception('Error while creating a container %s', docker_error)
+        return container
+
+    async def execute_cmd(self, cmd, image=None):
+        execute_result = None
+        try:
+            if image is None:
+                image = self.host_image
+
+            config = {"Cmd": cmd, "Image": image}
+            container = await self.docker_client.containers.run(config=config)
+
+            await container.wait()
+            result_exists = await container.log(stdout=True, stderr=True)
+            if result_exists:
+                execute_result = ''.join(result_exists)
+
+            # Deleting the used container
+            await container.delete(force=True)
+
+        except (aiodocker.exceptions.DockerError or aiodocker.exceptions.DockerContainerError) as server_error:
+            self.logger.error('Error while executing command %s in container %s', cmd, server_error)
+        return execute_result
+
+    async def delete_container(self, container_name):
+        container = await self.get_container(container_name)
+        try:
+            if container:
+                await container.delete(force=True)
+        except aiodocker.exceptions.DockerError as server_error:
+            self.logger.exception('Error while removing %s container %s', container_name, server_error)