From d272dcdb324fb72e5f019d8e0195eba524434af4 Mon Sep 17 00:00:00 2001 From: Rajat Gupta <35985127+rjt-gupta@users.noreply.github.com> Date: Wed, 24 Jul 2019 05:43:20 +0530 Subject: [PATCH] Aiodocker Helper (#341) * Aiodocker helper * Aiodocker tests * deleting used containers * close docker removed --- requirements.txt | 1 + tanner/emulators/cmd_exec.py | 19 +++--- tanner/emulators/lfi.py | 29 ++++----- tanner/session.py | 4 +- tanner/tests/test_aiodocker_helper.py | 73 ++++++++++++++++++++++ tanner/tests/test_base.py | 2 +- tanner/tests/test_cmd_exec_emulation.py | 9 ++- tanner/tests/test_lfi_emulator.py | 9 ++- tanner/utils/aiodocker_helper.py | 82 +++++++++++++++++++++++++ tanner/utils/docker_helper.py | 66 -------------------- 10 files changed, 192 insertions(+), 102 deletions(-) create mode 100644 tanner/tests/test_aiodocker_helper.py create mode 100644 tanner/utils/aiodocker_helper.py delete mode 100644 tanner/utils/docker_helper.py diff --git a/requirements.txt b/requirements.txt index ec0184a9..6e81b7f2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,3 +12,4 @@ pylibinjection jinja2 pycodestyle geoip2 +aiodocker diff --git a/tanner/emulators/cmd_exec.py b/tanner/emulators/cmd_exec.py index 807e0749..37a5aab4 100644 --- a/tanner/emulators/cmd_exec.py +++ b/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 diff --git a/tanner/emulators/lfi.py b/tanner/emulators/lfi.py index d6b9ebb5..68730830 100644 --- a/tanner/emulators/lfi.py +++ b/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 diff --git a/tanner/session.py b/tanner/session.py index ce5163bf..d57aca6a 100644 --- a/tanner/session.py +++ b/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) diff --git a/tanner/tests/test_aiodocker_helper.py b/tanner/tests/test_aiodocker_helper.py new file mode 100644 index 00000000..3ffe4f15 --- /dev/null +++ b/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() diff --git a/tanner/tests/test_base.py b/tanner/tests/test_base.py index 750fa8bf..334e2e62 100644 --- a/tanner/tests/test_base.py +++ b/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() diff --git a/tanner/tests/test_cmd_exec_emulation.py b/tanner/tests/test_cmd_exec_emulation.py index dc5ff920..421cd734 100644 --- a/tanner/tests/test_cmd_exec_emulation.py +++ b/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() diff --git a/tanner/tests/test_lfi_emulator.py b/tanner/tests/test_lfi_emulator.py index 93d58de3..a3b16898 100644 --- a/tanner/tests/test_lfi_emulator.py +++ b/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() diff --git a/tanner/utils/aiodocker_helper.py b/tanner/utils/aiodocker_helper.py new file mode 100644 index 00000000..fc4aae5d --- /dev/null +++ b/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) diff --git a/tanner/utils/docker_helper.py b/tanner/utils/docker_helper.py deleted file mode 100644 index 0ff993f8..00000000 --- a/tanner/utils/docker_helper.py +++ /dev/null @@ -1,66 +0,0 @@ -import logging -import docker - -from tanner.config import TannerConfig - -# TODO : Replace docker with aiodocker - - -class DockerHelper: - def __init__(self): - self.logger = logging.getLogger('tanner.docker_helper.DockerHelper') - try: - self.docker_client = docker.from_env(version='auto') - except docker.errors.APIError as docker_error: - self.logger.exception('Error while connecting to docker service %s', docker_error) - self.host_image = TannerConfig.get('DOCKER', 'host_image') - - async def setup_host_image(self): - try: - if not self.docker_client.images.list(self.host_image): - self.docker_client.images.pull(self.host_image) - except docker.errors.APIError 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_if_exists = self.docker_client.containers.list(all=True, - filters=dict(name=container_name) - ) - if container_if_exists: - container = container_if_exists[0] - except docker.errors.APIError as server_error: - self.logger.exception('Error while fetching container list %s', server_error) - return container - - async def create_container(self, container_name): - await self.setup_host_image() - container = await self.get_container(container_name) - if not container: - try: - container = self.docker_client.containers.create(image=self.host_image, - stdin_open=True, - name=container_name - ) - except (docker.errors.APIError, docker.errors.ImageNotFound) as docker_error: - self.logger.exception('Error while creating a container %s', docker_error) - return container - - async def delete_env(self, container_name): - container = await self.get_container(container_name) - try: - if container: - container.remove(force=True) - except docker.errors.APIError as server_error: - self.logger.exception('Error while removing container %s', server_error) - - async def execute_cmd(self, container, cmd): - execute_result = None - try: - container.start() - execute_result = container.exec_run(['sh', '-c', cmd]).decode('utf-8') - container.kill() - except docker.errors.APIError as server_error: - self.logger.exception('Error while executing command %s in container %s', cmd, server_error) - return execute_result