Merge pull request knownsec#101 from fenixns/master

refactory --verify mode ppt option code, add --shell mode ppt option support

pocsuite3/lib/controller/controller.py

@@ -1,7 +1,7 @@
 import copy
 import time
 
-from pocsuite3.lib.core.common import data_to_stdout
+from pocsuite3.lib.core.common import data_to_stdout, desensitization
 from pocsuite3.lib.core.data import conf, cmd_line_options
 from pocsuite3.lib.core.data import kb
 from pocsuite3.lib.core.data import logger
@@ -36,7 +36,8 @@ def start():
         task_done()
 
     if conf.mode == "shell" and not conf.api:
-        info_msg = "connect back ip: {0}    port: {1}".format(conf.connect_back_host, conf.connect_back_port)
+        info_msg = "connect back ip: {0}    port: {1}".format(
+            desensitization(conf.connect_back_host) if conf.ppt else conf.connect_back_host, conf.connect_back_port)
         logger.info(info_msg)
         info_msg = "watting for shell connect to pocsuite"
         logger.info(info_msg)
@@ -96,13 +97,7 @@ def task_run():
 
         # for hide some infomations
         if conf.ppt:
-            length = len(target)
-            _target = target
-            if length > 15:
-                _target = "*" + _target[length - 9:]
-            else:
-                _target = "*" + _target[length - 3:]
-            info_msg = "running poc:'{0}' target '{1}'".format(poc_name, _target)
+            info_msg = "running poc:'{0}' target '{1}'".format(poc_name, desensitization(target))
         else:
             info_msg = "running poc:'{0}' target '{1}'".format(poc_name, target)
 
@@ -167,13 +162,7 @@ def task_run():
         output = AttribDict(result.to_dict())
         if conf.ppt:
             # hide some information
-            length = len(target)
-            if length > 15:
-                target = "*" + target[length - 9:]
-            elif length > 8:
-                target = "*" + target[4:]
-            else:
-                target = "*" + target[1:]
+            target = desensitization(target)
 
         output.update({
             'target': target,

pocsuite3/lib/core/common.py

@@ -937,6 +937,17 @@ def get_all_nic_info():
     return nic_info
 
 
+def desensitization(s):
+    """ Hide sensitive information.
+    """
+    s = str(s)
+    return (
+        s[:len(s) // 4 if len(s) < 30 else 8] +
+        '***' +
+        s[len(s) * 3 // 4:]
+    )
+
+
 def get_host_ipv6(with_nic=True):
     nic_info = get_all_nic_info()
     ipv4 = get_host_ip()

pocsuite3/lib/core/option.py

@@ -11,7 +11,7 @@
 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
+from pocsuite3.lib.core.common import get_local_ip, desensitization
 from pocsuite3.lib.core.common import single_time_warn_message
 from pocsuite3.lib.core.convert import stdout_encode
 from pocsuite3.lib.core.data import conf, cmd_line_options
@@ -284,7 +284,7 @@ def _set_connect_back():
         data_to_stdout("[i] pocsusite is running in shell mode, you need to set connect back host:\n")
         message = '----- Local IP Address -----\n'
         for i, ip in enumerate(kb.data.local_ips):
-            message += "{0}    {1}\n".format(i, ip)
+            message += "{0}    {1}\n".format(i, desensitization(ip) if conf.ppt else ip)
         data_to_stdout(message)
         while True:
             choose = None
@@ -295,7 +295,8 @@ def _set_connect_back():
                 if choose.isdigit():
                     choose = int(choose)
                     conf.connect_back_host = kb.data.local_ips[choose]
-                    data_to_stdout("you choose {0}\n".format(conf.connect_back_host))
+                    data_to_stdout("you choose {0}\n".format(
+                        desensitization(conf.connect_back_host) if conf.ppt else conf.connect_back_host))
                     break
             except Exception:
                 data_to_stdout("wrong number, choose again\n")

pocsuite3/lib/core/poc.py

@@ -7,7 +7,7 @@
 from requests.exceptions import HTTPError
 from requests.exceptions import TooManyRedirects
 
-from pocsuite3.lib.core.common import parse_target_url
+from pocsuite3.lib.core.common import parse_target_url, desensitization
 from pocsuite3.lib.core.data import conf
 from pocsuite3.lib.core.data import logger
 from pocsuite3.lib.core.enums import OUTPUT_STATUS, CUSTOM_LOGGING, ERROR_TYPE_ID, POC_CATEGORY
@@ -195,7 +195,7 @@ def execute(self, target, headers=None, params=None, mode='verify', verbose=True
                     logger.debug('POC: {0} time-out retry failed!'.format(self.name))
                 conf.retry -= 1
             else:
-                msg = "connect target '{0}' failed!".format(target)
+                msg = "connect target '{0}' failed!".format(desensitization(target) if conf.ppt else target)
                 logger.error(msg)
                 output = Output(self)
 
@@ -206,7 +206,7 @@ def execute(self, target, headers=None, params=None, mode='verify', verbose=True
 
         except ConnectionError as e:
             self.expt = (ERROR_TYPE_ID.CONNECTIONERROR, e)
-            msg = "connect target '{0}' failed!".format(target)
+            msg = "connect target '{0}' failed!".format(desensitization(target) if conf.ppt else target)
             logger.error(msg)
             output = Output(self)
 
@@ -295,27 +295,11 @@ def show_result(self):
                 if isinstance(v, dict):
                     for kk, vv in v.items():
                         if (kk == "URL" or kk == "IP") and conf.ppt:
-                            length = len(vv)
-                            _target = vv
-                            if length > 15:
-                                _target = "*" + _target[length - 9:]
-                            elif length > 8:
-                                _target = "*" + _target[5:]
-                            else:
-                                _target = "*" + _target[3:]
-                            vv = _target
+                            vv = desensitization(vv)
                         logger.log(CUSTOM_LOGGING.SUCCESS, "%s : %s" % (kk, vv))
                 else:
                     if (k == "URL" or k == "IP") and conf.ppt:
-                        length = len(v)
-                        _target = v
-                        if length > 15:
-                            _target = "*" + _target[length - 9:]
-                        elif length > 8:
-                            _target = "*" + _target[5:]
-                        else:
-                            _target = "*" + _target[3:]
-                        v = _target
+                        v = desensitization(v)
                     logger.log(CUSTOM_LOGGING.SUCCESS, "%s : %s" % (k, v))
 
     def to_dict(self):

pocsuite3/modules/listener/reverse_tcp.py

@@ -3,7 +3,7 @@
 import threading
 import time
 
-from pocsuite3.lib.core.common import data_to_stdout, has_poll, get_unicode
+from pocsuite3.lib.core.common import data_to_stdout, has_poll, get_unicode, desensitization
 from pocsuite3.lib.core.data import conf, kb, logger
 from pocsuite3.lib.core.datatype import AttribDict
 from pocsuite3.lib.core.enums import AUTOCOMPLETE_TYPE, OS, CUSTOM_LOGGING
@@ -66,7 +66,8 @@ def listener_worker():
             client.conn = conn
             client.address = address
             kb.data.clients.append(client)
-            info_msg = "new connection established from {0}".format(address[0])
+            info_msg = "new connection established from {0}".format(
+                desensitization(address[0]) if conf.ppt else address[0])
             logger.log(CUSTOM_LOGGING.SUCCESS, info_msg)
         except Exception:
             pass
@@ -93,8 +94,15 @@ def list_clients():
             logger.exception(ex)
             del kb.data.clients[i]
             continue
-        results += str(i) + "   " + str(client.address[0]) + "    " + str(client.address[1]) + " ({0})".format(
-            system) + '\n'
+        results += (
+            str(i) +
+            "   " +
+            (desensitization(client.address[0]) if conf.ppt else str(client.address[0])) +
+            "    " +
+            str(client.address[1]) +
+            " ({0})".format(system) +
+            '\n'
+        )
     data_to_stdout("----- Remote Clients -----" + "\n" + results)
 
 
@@ -103,7 +111,8 @@ def get_client(cmd):
         target = cmd.replace("select ", "")
         target = int(target)
         client = kb.data.clients[target]  # Connect to the selected clients
-        data_to_stdout("Now Connected: {0}\n".format(str(kb.data.clients[target].address[0])))
+        data_to_stdout("Now Connected: {0}\n".format(
+            desensitization(client.address[0] if conf.ppt else client.address[0])))
         return client
     except Exception:
         data_to_stdout("Invalid Client\n")
@@ -146,7 +155,7 @@ def send_shell_commands(client):
         cmd = None
         try:
             address = client.address[0]
-            cmd = input("{0}>: ".format(address))
+            cmd = input("{0}>: ".format(desensitization(address) if conf.ppt else address))
             if not cmd:
                 continue