Merge pull request #1 from wh0am1i/dev

Dev

pocsuite3/__init__.py

@@ -1,5 +1,5 @@
 __title__ = 'pocsuite3'
-__version__ = '2.0.5'
+__version__ = '2.0.6'
 __author__ = 'Knownsec 404 Team'
 __author_email__ = '[email protected]'
 __license__ = 'GPLv2'

pocsuite3/lib/core/common.py

@@ -458,6 +458,10 @@ def parse_target(address, additional_ports=[], skip_target_port=False):
     return targets
 
 
+def parse_poc_docker_name(name):
+    return name.lower().replace(' ', '_')
+
+
 def single_time_log_message(message, level=logging.INFO, flag=None):
     if flag is None:
         flag = hash(message)

pocsuite3/lib/core/docker_env.py

@@ -0,0 +1,64 @@
+from io import BytesIO
+from docker import client
+from docker import errors
+
+
+from pocsuite3.lib.core.data import logger
+
+
+class DockerEnv:
+
+    def __init__(self):
+        self.client = client.from_env()
+
+    def build(self, name, docker_file):
+        file_obj = BytesIO(docker_file.encode())
+        try:
+            logger.info("Building image...")
+            build_info = self.client.images.build(fileobj=file_obj, tag=name)
+            return build_info
+        except errors.BuildError as e:
+            logger.error(e)
+
+    def run(self, tag_name, docker_file, ports, envs, volumes):
+        try:
+            # if image exists run
+            self.client.images.get(tag_name)
+            logger.info("Image {} exists".format(tag_name))
+            run_info = self.client.containers.run(
+                tag_name,
+                detach=True,
+                ports=ports,
+                environment=envs,
+                volumes=volumes
+            )
+            return run_info
+        except errors.ImageNotFound:
+            # if image not exists, build image first
+            logger.info("Image {} does not exist".format(tag_name))
+            build_info = self.build(tag_name, docker_file)
+            if build_info[0].tags:
+                run_info = self.client.containers.run(
+                    tag_name,
+                    detach=True,
+                    ports=ports,
+                    environment=envs,
+                    volumes=volumes
+                )
+                return run_info
+
+
+if __name__ == "__main__":
+    docker_env = DockerEnv()
+    ports = {"8080/tcp": '8899', '8090/tcp': ("127.0.0.1", 8890)}
+    env = ["PORT=8899", "PORT=8890"]
+    volumes = ["/tmp:/home"]
+    dockerfile = "FROM ubuntu:latest"
+    image_tag = "ubuntu:pocsuite"
+    docker_env.run(
+        image_tag,
+        docker_file=dockerfile,
+        ports=ports,
+        envs=env,
+        volumes=volumes
+    )

pocsuite3/lib/core/option.py

@@ -8,16 +8,20 @@
 from queue import Queue
 from urllib.parse import urlsplit
 
+import docker.errors
 import socks
 import prettytable
 from termcolor import colored
+from dockerfile import parse_string
 from pocsuite3.lib.core.clear import remove_extra_log_message
 from pocsuite3.lib.core.common import boldify_message, check_file, get_file_items, parse_target, \
     get_public_type_members, data_to_stdout
 from pocsuite3.lib.core.common import check_path, extract_cookies
 from pocsuite3.lib.core.common import get_local_ip, mosaic, get_host_ip
 from pocsuite3.lib.core.common import single_time_warn_message
-from pocsuite3.lib.core.common import OrderedSet, get_file_text
+from pocsuite3.lib.core.common import OrderedSet, get_file_text, get_poc_name
+from pocsuite3.lib.core.common import index_modules, ltrim
+from pocsuite3.lib.core.common import parse_poc_docker_name
 from pocsuite3.lib.core.convert import stdout_encode
 from pocsuite3.lib.core.data import conf, cmd_line_options
 from pocsuite3.lib.core.data import kb
@@ -30,12 +34,14 @@
 from pocsuite3.lib.core.log import FORMATTER
 from pocsuite3.lib.core.register import load_file_to_module
 from pocsuite3.lib.core.settings import DEFAULT_LISTENER_PORT, CMD_PARSE_WHITELIST
+from pocsuite3.lib.core.docker_env import DockerEnv
 from pocsuite3.lib.core.statistics_comparison import StatisticsComparison
 from pocsuite3.lib.core.update import update
 from pocsuite3.lib.core.template import create_poc_plugin_template
 from pocsuite3.lib.parse.cmd import DIY_OPTIONS
 from pocsuite3.lib.parse.configfile import config_file_parser
 from pocsuite3.lib.parse.rules import regex_rule
+from pocsuite3.lib.parse.dockerfile import parse_dockerfile
 from pocsuite3.lib.request.patch import patch_all
 from pocsuite3.modules.listener import start_listener
 
@@ -523,6 +529,7 @@ def _set_conf_attributes():
     conf.mode = 'verify'
     conf.poc = None
     conf.poc_keyword = None
+    conf.poc_list = None
     conf.cookie = None
     conf.host = None
     conf.referer = None
@@ -589,6 +596,13 @@ def _set_conf_attributes():
     conf.show_options = False
     conf.enable_tls_listener = False
 
+    # docker args
+    conf.docker_start = False
+    conf.docker_port = list()
+    conf.docker_env = list()
+    conf.docker_volume = list()
+    conf.docker_only = False
+
 
 def _set_kb_attributes(flush_all=True):
     """
@@ -661,6 +675,46 @@ def _set_poc_options(input_options):
             DIY_OPTIONS.append(line)
 
 
+def _set_docker_options():
+    port_dict = {}
+    if conf.poc and conf.docker_start:
+        # parse port string
+        # in docker package ports need {"2222/tcp": 3333}
+        # will expose port 2222 inside the container as port 3333 on the host.
+        if len(conf.docker_port) > 0:
+            for item in conf.docker_port:
+                item_split = item.rsplit(':', 1)
+                key = '{}/tcp'.format(item_split[1])
+                # if user input 127.0.0.1:8080:8080
+                # need change to {'8080/tcp': ('127.0.0.1', 8080)}
+                if item_split[0].find(":") > 0:
+                    temp_port = item_split[0].split(':')
+                    host_port = (temp_port[0], int(temp_port[1]))
+                else:
+                    host_port = int(item_split[0])
+                port_dict[key] = host_port
+
+        for poc in conf.poc:
+            res = parse_dockerfile(poc)
+            poc_name = parse_poc_docker_name(res.get('name'))
+            tag_name = "{}:{}".format(poc_name, "pocsuite")
+            docker_file = res.get("dockerfile")
+            try:
+                dk = DockerEnv()
+                dk_info = dk.run(
+                    tag_name,
+                    docker_file,
+                    ports=port_dict,
+                    envs=conf.docker_env,
+                    volumes=conf.docker_volume,
+                )
+                logger.info("Run container {} successful!".format(dk_info.short_id))
+            except docker.errors.APIError as e:
+                logger.error(e)
+        if conf.docker_only:
+            exit(0)
+
+
 def init_options(input_options=AttribDict(), override_options=False):
     cmd_line_options.update(input_options)
     _set_conf_attributes()
@@ -754,6 +808,49 @@ def _show_pocs_modules_options():
     exit()
 
 
+def _show_pocs_form_local():
+    if not conf.poc_list:
+        return
+    _pocs = []
+    tb = prettytable.PrettyTable(["Index", "Path", "Name"])
+    if conf.pocs_path:
+        # parse user defined poc scripts path
+        if check_path(conf.pocs_path):
+            for root, dirs, files in os.walk(conf.pocs_path):
+                files = list(filter(lambda x: not x.startswith("__") and x.endswith(".py") or x.endswith(".yaml"), files))
+                conf.poc = [os.path.join(conf.pocs_path, f) for f in files]
+        moduels = index_modules(conf.pocs_path)
+        poc_parent_directory = os.sep.join(
+            conf.pocs_path.rstrip(os.sep).split(os.sep)) + os.sep
+        for poc in moduels:
+            _pocs.append(ltrim(poc, conf.pocs_path).lstrip(os.sep))
+
+    else:
+        # default poc path
+        conf.poc = [paths.POCSUITE_POCS_PATH]
+        moduels = index_modules(paths.POCSUITE_POCS_PATH)
+        poc_parent_directory = os.sep.join(
+            paths.POCSUITE_POCS_PATH.rstrip(os.sep).split(os.sep)[0:-1]) + os.sep
+        for poc in moduels:
+            _pocs.append(ltrim(poc, poc_parent_directory).lstrip(os.sep))
+    # show result table
+    try:
+        index = 0
+        for poc in _pocs:
+            file = os.path.join(poc_parent_directory, poc + ".py")
+            if not os.path.exists(file):
+                file = os.path.join(poc_parent_directory, poc + ".yaml")
+            code = get_file_text(file)
+            name = get_poc_name(code)
+            tb.add_row([str(index), poc, name])
+            index += 1
+        data_to_stdout("\n" + tb.get_string() + "\n")
+    except PocsuiteSystemException as ex:
+        logger.error(str(ex))
+    finally:
+        exit(0)
+
+
 def init():
     """
     Set attributes into both configuration and knowledge base singletons
@@ -770,9 +867,13 @@ def init():
     update()
     _set_multiple_targets()
     _set_user_pocs_path()
+    # show poc form pac path
+    _show_pocs_form_local()
     # The poc module module must be in front of the plug-in module,
     # and some parameters in the poc option call the plug-in
     _set_pocs_modules()
+    # run docker from poc
+    _set_docker_options()
     _set_plugins()
     _init_targets_plugins()
     _init_pocs_plugins()

pocsuite3/lib/core/settings.py

@@ -118,6 +118,7 @@
     "skip_target_port",
     "file",
     "poc_keyword",
+    "poc_list",
     "verify",
     "attack",
     "shell",
@@ -174,6 +175,13 @@
     "no-check",
     "options",
 
+    # docker
+    "docker-start",
+    "docker-port",
+    "docker-env",
+    "docker-volume",
+    "docker-only",
+
     # other
     "poc",
     "verbose",

pocsuite3/lib/parse/cmd.py

@@ -46,6 +46,7 @@ def cmd_line_parser(argv=None):
         target.add_argument("-r", dest="poc", nargs='+', help="Load PoC file from local or remote from seebug website")
         target.add_argument("-k", dest="poc_keyword", help="Filter PoC by keyword, e.g. ecshop")
         target.add_argument("-c", dest="configFile", help="Load options from a configuration INI file")
+        target.add_argument("-l", dest="poc_list", action="store_true", help="Show all PoC file from local")
 
         # Mode options
         mode = parser.add_argument_group("Mode", "Pocsuite running mode options")
@@ -152,6 +153,20 @@ def cmd_line_parser(argv=None):
                                   help="Specify the name of the export rule file")
         optimization.add_argument("--no-check", dest="no_check", action="store_true", default=False,
                                   help="Disable URL protocol correction and honeypot check")
+
+        # docker options
+        docker_environment = parser.add_argument_group("Docker Environment", "Docker Environment options")
+        docker_environment.add_argument("--docker-start", dest="docker_start", action="store_true",
+                                        default=False, help="Run the docker for PoC")
+        docker_environment.add_argument("--docker-port", dest="docker_port", action="append",
+                                        default=[], help="Publish a container's port(s) to the host")
+        docker_environment.add_argument("--docker-volume", dest="docker_volume", action="append",
+                                        default=[], help="Bind mount a volume")
+        docker_environment.add_argument("--docker-env", dest="docker_env", action="append", default=[],
+                                        help="Set environment variables")
+        docker_environment.add_argument("--docker-only", dest="docker_only", action="store_true",
+                                        default=False, help="Only run docker environment")
+
         # Diy options
         diy = parser.add_argument_group("Poc options", "definition options for PoC")
         diy.add_argument("--options", dest="show_options", action="store_true", default=False,

pocsuite3/lib/parse/dockerfile.py

@@ -0,0 +1,27 @@
+import re
+from pocsuite3.lib.core.data import conf
+from pocsuite3.lib.core.data import logger
+from pocsuite3.lib.core.common import get_file_text
+
+
+def parse_dockerfile(file):
+        regx_rules = [
+            "name = '(.*)'",
+            "vulID = '(.*)'",
+            r"dockerfile = '''([\s\S]*?)'''",
+        ]
+        result = {
+            "name": 0,
+            "vulID": 1,
+            "dockerfile": 2,
+        }
+        st = get_file_text(file)
+        for k, v in result.items():
+            pattern = re.compile(regx_rules[v])
+            match = pattern.findall(st)
+            if match is not None:
+                result[k] = match[0]
+
+        return result
+
+

requirements.txt

@@ -15,3 +15,4 @@ dacite >= 1.6.0
 PyYAML >= 6.0
 lxml >= 4.6.0
 mmh3 >= 3.0.0
+docker >=  6.1.3