diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 3e6e191a..010dcd04 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -61,4 +61,5 @@ HomerQing * contributing to fix ipv6 compatibility issue in build_url Xxcdd -* Add support for requests session reuse \ No newline at end of file +* Add support for requests session reuse +* Add support for web hook diff --git a/pocsuite3/lib/core/option.py b/pocsuite3/lib/core/option.py index d35f35da..81e11060 100644 --- a/pocsuite3/lib/core/option.py +++ b/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): """ diff --git a/pocsuite3/lib/core/settings.py b/pocsuite3/lib/core/settings.py index fb36e30e..cec7b82e 100644 --- a/pocsuite3/lib/core/settings.py +++ b/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" ] diff --git a/pocsuite3/lib/parse/cmd.py b/pocsuite3/lib/parse/cmd.py index a97a581b..dd22a9d4 100644 --- a/pocsuite3/lib/parse/cmd.py +++ b/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, diff --git a/pocsuite3/plugins/web_hook.py b/pocsuite3/plugins/web_hook.py new file mode 100644 index 00000000..fa6dbfa5 --- /dev/null +++ b/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 = "×tamp={}&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)