Merge pull request #402 from xxcdd/master

Add support for web hook

CONTRIBUTORS.md

@@ -61,4 +61,5 @@ HomerQing <https://github.com/HomerQing>
 * contributing to fix ipv6 compatibility issue in build_url
 
 Xxcdd <https://github.com/xxcdd>
-* Add support for requests session reuse
+* Add support for requests session reuse
+* Add support for web hook

pocsuite3/lib/core/option.py

@@ -482,6 +482,8 @@ def _cleanup_options():
 
     if conf.output_path and 'file_record' not in conf.plugins:
         conf.plugins.append('file_record')
+    if 'web_hook' not in conf.plugins:
+        conf.plugins.append('web_hook')
 
     if conf.connect_back_port:
         conf.connect_back_port = int(conf.connect_back_port)
@@ -622,6 +624,11 @@ def _set_conf_attributes():
     conf.docker_only = False
     conf.requests_session_reuse = False
 
+    # web hook
+    conf.dingtalk_token = ""
+    conf.dingtalk_secret = ""
+    conf.wx_work_key = ""
+
 
 def _set_kb_attributes(flush_all=True):
     """

pocsuite3/lib/core/settings.py

@@ -191,6 +191,11 @@
     "connect_back_host",
     "connect_back_port",
 
-    "requests-session-reuse",
-    "requests-session-reuse-num"
+    "session-reuse",
+    "session-reuse-num",
+
+    # web hook
+    "dingtalk-token",
+    "dingtalk-secret",
+    "wx-work-key"
 ]

pocsuite3/lib/parse/cmd.py

@@ -73,9 +73,9 @@ def cmd_line_parser(argv=None):
         request.add_argument("--delay", dest="delay", help="Delay between two request of one thread")
         request.add_argument("--headers", dest="headers", help="Extra headers (e.g. \"key1: value1\\nkey2: value2\")")
         request.add_argument("--http-debug", dest="http_debug", type=int, default=0, help="HTTP debug level (default 0)")
-        request.add_argument("--requests-session-reuse", dest="requests_session_reuse", action="store_true",
+        request.add_argument("--session-reuse", dest="requests_session_reuse", action="store_true",
                              help="Enable requests session reuse")
-        request.add_argument("--requests-session-reuse-num", type=int, dest="requests_session_reuse_num", default=10,
+        request.add_argument("--session-reuse-num", type=int, dest="requests_session_reuse_num", default=10,
                              help="Requests session reuse number")
 
         # Account options
@@ -173,6 +173,12 @@ def cmd_line_parser(argv=None):
         docker_environment.add_argument("--docker-only", dest="docker_only", action="store_true",
                                         default=False, help="Only run docker environment")
 
+        # web hook options
+        web_hook = parser.add_argument_group('Web Hook', "Web Hook Options")
+        web_hook.add_argument("--dingtalk-token", dest="dingtalk_token", help="Dingtalk access token")
+        web_hook.add_argument("--dingtalk-secret", dest="dingtalk_secret", help="Dingtalk secret")
+        web_hook.add_argument("--wx-work-key", dest="wx_work_key", help="Weixin Work key")
+
         # 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/plugins/web_hook.py

@@ -0,0 +1,76 @@
+import hmac
+import hashlib
+import base64
+import urllib.parse
+import requests
+import time
+
+from pocsuite3.api import PLUGIN_TYPE, get_results
+from pocsuite3.api import PluginBase
+from pocsuite3.api import logger
+from pocsuite3.api import register_plugin, conf
+
+DINGTALK_TOKEN = ""
+DINGTALK_SECRET = ""
+WX_WORK_KEY = ""
+
+
+def dingding_send(msg, access_token, secret, msgtype="markdown", title="pocsuite3消息推送"):
+    ding_url = "https://oapi.dingtalk.com/robot/send?access_token={}".format(access_token)
+    timestamp = str(round(time.time() * 1000))
+    secret_enc = secret.encode('utf-8')
+    string_to_sign = '{}\n{}'.format(timestamp, secret)
+    string_to_sign_enc = string_to_sign.encode('utf-8')
+    hmac_code = hmac.new(secret_enc, string_to_sign_enc, digestmod=hashlib.sha256).digest()
+    sign = urllib.parse.quote_plus(base64.b64encode(hmac_code))
+    param = "&timestamp={}&sign={}".format(timestamp, sign)
+    ding_url = ding_url + param
+    send_json = {
+        "msgtype": msgtype,
+        "markdown": {
+            "title": title,
+            "text": "# pocsuite3消息推送\n\n" + msg
+        }
+    }
+    requests.post(ding_url, json=send_json)
+
+
+def wx_work_send(msg, key):
+    webhook_url = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=" + key
+    send_data = {
+        "msgtype": "markdown",
+        "markdown": {
+            "content": "# pocsuite3消息推送\n\n" + msg
+        }
+    }
+    requests.post(webhook_url, json=send_data)
+
+
+def web_hook_send(msg):
+    dingtalk_token = conf.dingtalk_token or DINGTALK_TOKEN
+    dingtalk_secret = conf.dingtalk_secret or DINGTALK_SECRET
+    wx_work_key = conf.wx_work_key or WX_WORK_KEY
+    if dingtalk_token and dingtalk_secret:
+        dingding_send(msg, dingtalk_token, dingtalk_secret)
+    if wx_work_key:
+        wx_work_send(msg, wx_work_key)
+
+
+class WebHook(PluginBase):
+    category = PLUGIN_TYPE.RESULTS
+
+    def init(self):
+        debug_msg = "[PLUGIN] web hook plugin init..."
+        logger.debug(debug_msg)
+
+    def start(self):
+        push_info = ""
+        for result in get_results():
+            if result.status == "success":
+                poc_name = result.get("poc_name")
+                target = result.get("target")
+                push_info += "- {} found vuln: {}".format(target, poc_name)
+        web_hook_send(push_info)
+
+
+register_plugin(WebHook)