From 9be0baca496e95bd84564b821dde36537e86a99f Mon Sep 17 00:00:00 2001 From: xxcdd Date: Sun, 31 Dec 2023 00:14:34 +0800 Subject: [PATCH 1/7] Add support for requests session reuse --- CONTRIBUTORS.md | 5 ++++- pocsuite3/lib/core/option.py | 19 +++++++++++++++++++ pocsuite3/lib/core/settings.py | 5 ++++- pocsuite3/lib/parse/cmd.py | 5 +++++ pocsuite3/lib/request/patch/session_reuse.py | 19 +++++++++++++++++++ 5 files changed, 51 insertions(+), 2 deletions(-) create mode 100644 pocsuite3/lib/request/patch/session_reuse.py diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index e5986a89..3e6e191a 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -58,4 +58,7 @@ ekszz * contributing to customize poc protocol and default port #321 HomerQing -* contributing to fix ipv6 compatibility issue in build_url \ No newline at end of file +* contributing to fix ipv6 compatibility issue in build_url + +Xxcdd +* Add support for requests session reuse \ No newline at end of file diff --git a/pocsuite3/lib/core/option.py b/pocsuite3/lib/core/option.py index b10dc9f2..d35f35da 100644 --- a/pocsuite3/lib/core/option.py +++ b/pocsuite3/lib/core/option.py @@ -10,6 +10,8 @@ from http.client import HTTPConnection import docker.errors +import requests +from requests_toolbelt.adapters.socket_options import TCPKeepAliveAdapter import socks import prettytable from termcolor import colored @@ -43,6 +45,7 @@ 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.lib.request.patch.session_reuse import api_request from pocsuite3.modules.listener import start_listener @@ -202,6 +205,18 @@ def _set_network_proxy(): } +def _set_session_queue(): + requests.api.request = api_request + for _ in range(0, conf.requests_session_reuse_num): + session = requests.Session() + session.headers.update({'Connection': 'keep-alive'}) + # https://github.com/psf/requests/issues/6354 + keep_alive = TCPKeepAliveAdapter() + session.mount("http://", keep_alive) + session.mount("https://", keep_alive) + kb.session_queue.put(session) + + def _set_multiple_targets(): # set multi targets to kb if conf.url: @@ -605,6 +620,7 @@ def _set_conf_attributes(): conf.docker_env = list() conf.docker_volume = list() conf.docker_only = False + conf.requests_session_reuse = False def _set_kb_attributes(flush_all=True): @@ -648,6 +664,7 @@ def _set_kb_attributes(flush_all=True): kb.current_poc = None kb.registered_pocs = AttribDict() kb.task_queue = Queue() + kb.session_queue = Queue() kb.cmd_line = DIY_OPTIONS or [] kb.comparison = None @@ -898,3 +915,5 @@ def init(): _set_threads() _set_listener() remove_extra_log_message() + if conf.requests_session_reuse: + _set_session_queue() diff --git a/pocsuite3/lib/core/settings.py b/pocsuite3/lib/core/settings.py index 80402f60..fb36e30e 100644 --- a/pocsuite3/lib/core/settings.py +++ b/pocsuite3/lib/core/settings.py @@ -189,5 +189,8 @@ "mode", "api", "connect_back_host", - "connect_back_port" + "connect_back_port", + + "requests-session-reuse", + "requests-session-reuse-num" ] diff --git a/pocsuite3/lib/parse/cmd.py b/pocsuite3/lib/parse/cmd.py index 82c47f48..a97a581b 100644 --- a/pocsuite3/lib/parse/cmd.py +++ b/pocsuite3/lib/parse/cmd.py @@ -73,6 +73,11 @@ 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", + help="Enable requests session reuse") + request.add_argument("--requests-session-reuse-num", type=int, dest="requests_session_reuse_num", default=10, + help="Requests session reuse number") + # Account options group = parser.add_argument_group("Account", "Account options") group.add_argument("--ceye-token", dest="ceye_token", help="CEye token") diff --git a/pocsuite3/lib/request/patch/session_reuse.py b/pocsuite3/lib/request/patch/session_reuse.py new file mode 100644 index 00000000..9334974b --- /dev/null +++ b/pocsuite3/lib/request/patch/session_reuse.py @@ -0,0 +1,19 @@ +from pocsuite3.lib.core.data import kb + + +class ReuseSession: + def __init__(self): + self.session_queue = kb.session_queue + self.session = None + + def __enter__(self): + self.session = self.session_queue.get() + return self.session + + def __exit__(self, *args): + self.session_queue.put(self.session) + + +def api_request(method, url, **kwargs): + with ReuseSession() as session: + return session.request(method=method, url=url, **kwargs) \ No newline at end of file From 32b46433e9a818e58ee8214ac2d85a905cf56104 Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Tue, 2 Jan 2024 18:10:52 +0800 Subject: [PATCH 2/7] FIX: W292 no newline at end of file --- pocsuite3/lib/request/patch/session_reuse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pocsuite3/lib/request/patch/session_reuse.py b/pocsuite3/lib/request/patch/session_reuse.py index 9334974b..24ab17d5 100644 --- a/pocsuite3/lib/request/patch/session_reuse.py +++ b/pocsuite3/lib/request/patch/session_reuse.py @@ -16,4 +16,4 @@ def __exit__(self, *args): def api_request(method, url, **kwargs): with ReuseSession() as session: - return session.request(method=method, url=url, **kwargs) \ No newline at end of file + return session.request(method=method, url=url, **kwargs) From ae07716636fa4217d0c38e857329cecc5789ce1e Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Sat, 6 Jan 2024 14:58:20 +0800 Subject: [PATCH 3/7] Add web hook report --- pocsuite3/lib/core/option.py | 8 ++++ pocsuite3/lib/core/settings.py | 7 +++- pocsuite3/lib/parse/cmd.py | 6 +++ pocsuite3/plugins/web_hook.py | 69 ++++++++++++++++++++++++++++++++++ 4 files changed, 89 insertions(+), 1 deletion(-) create mode 100644 pocsuite3/plugins/web_hook.py diff --git a/pocsuite3/lib/core/option.py b/pocsuite3/lib/core/option.py index b10dc9f2..2ff8338d 100644 --- a/pocsuite3/lib/core/option.py +++ b/pocsuite3/lib/core/option.py @@ -467,6 +467,9 @@ def _cleanup_options(): if conf.output_path and 'file_record' not in conf.plugins: conf.plugins.append('file_record') + if (conf.dingtalk_token and conf.dingtalk_secret) or conf.wx_work_key: + 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) @@ -606,6 +609,11 @@ def _set_conf_attributes(): conf.docker_volume = list() conf.docker_only = 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 80402f60..b9d8aafa 100644 --- a/pocsuite3/lib/core/settings.py +++ b/pocsuite3/lib/core/settings.py @@ -189,5 +189,10 @@ "mode", "api", "connect_back_host", - "connect_back_port" + "connect_back_port", + + # web hook + "dingtalk-token", + "dingtalk-secret", + "wx-work-key" ] diff --git a/pocsuite3/lib/parse/cmd.py b/pocsuite3/lib/parse/cmd.py index 82c47f48..d22f95fc 100644 --- a/pocsuite3/lib/parse/cmd.py +++ b/pocsuite3/lib/parse/cmd.py @@ -168,6 +168,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..1ccd90be --- /dev/null +++ b/pocsuite3/plugins/web_hook.py @@ -0,0 +1,69 @@ +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 + + +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): + if conf.dingtalk_token and conf.dingtalk_secret: + dingding_send(msg, conf.dingtalk_token, conf.dingtalk_secret) + if conf.wx_work_key: + wx_work_send(msg, conf.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) From 0f493a31750832bc6d884b776e13374f42d56707 Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Sat, 6 Jan 2024 15:05:37 +0800 Subject: [PATCH 4/7] Update settings.py --- pocsuite3/lib/core/settings.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pocsuite3/lib/core/settings.py b/pocsuite3/lib/core/settings.py index c6d92d55..1524139c 100644 --- a/pocsuite3/lib/core/settings.py +++ b/pocsuite3/lib/core/settings.py @@ -198,3 +198,4 @@ "dingtalk-token", "dingtalk-secret", "wx-work-key" +] From 95d08fbd9f844989fbdac39fd8c7a7553299c72a Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Sat, 6 Jan 2024 15:09:21 +0800 Subject: [PATCH 5/7] Update settings.py FIX: > ./pocsuite3/lib/core/settings.py:196:1: W293 blank line contains whitespace --- pocsuite3/lib/core/settings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pocsuite3/lib/core/settings.py b/pocsuite3/lib/core/settings.py index 1524139c..ccd3c25b 100644 --- a/pocsuite3/lib/core/settings.py +++ b/pocsuite3/lib/core/settings.py @@ -193,7 +193,7 @@ "requests-session-reuse", "requests-session-reuse-num", - + # web hook "dingtalk-token", "dingtalk-secret", From 3249bc366fd849f98c197a0712f712347ed61f71 Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Sat, 6 Jan 2024 16:05:21 +0800 Subject: [PATCH 6/7] simplify requests session argument --- CONTRIBUTORS.md | 3 ++- pocsuite3/lib/core/settings.py | 4 ++-- pocsuite3/lib/parse/cmd.py | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 3e6e191a..b8518b7f 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 \ No newline at end of file diff --git a/pocsuite3/lib/core/settings.py b/pocsuite3/lib/core/settings.py index ccd3c25b..cec7b82e 100644 --- a/pocsuite3/lib/core/settings.py +++ b/pocsuite3/lib/core/settings.py @@ -191,8 +191,8 @@ "connect_back_host", "connect_back_port", - "requests-session-reuse", - "requests-session-reuse-num", + "session-reuse", + "session-reuse-num", # web hook "dingtalk-token", diff --git a/pocsuite3/lib/parse/cmd.py b/pocsuite3/lib/parse/cmd.py index 5297a37e..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 From cdd847c7695be699852b298339a1fda763e53747 Mon Sep 17 00:00:00 2001 From: xxcdd <42600601+xxcdd@users.noreply.github.com> Date: Wed, 10 Jan 2024 22:46:10 +0800 Subject: [PATCH 7/7] add web hook token local persistence --- pocsuite3/lib/core/option.py | 5 ++--- pocsuite3/plugins/web_hook.py | 15 +++++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/pocsuite3/lib/core/option.py b/pocsuite3/lib/core/option.py index bca86c2a..81e11060 100644 --- a/pocsuite3/lib/core/option.py +++ b/pocsuite3/lib/core/option.py @@ -482,9 +482,8 @@ def _cleanup_options(): if conf.output_path and 'file_record' not in conf.plugins: conf.plugins.append('file_record') - if (conf.dingtalk_token and conf.dingtalk_secret) or conf.wx_work_key: - if 'web_hook' not in conf.plugins: - conf.plugins.append('web_hook') + 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) diff --git a/pocsuite3/plugins/web_hook.py b/pocsuite3/plugins/web_hook.py index 1ccd90be..fa6dbfa5 100644 --- a/pocsuite3/plugins/web_hook.py +++ b/pocsuite3/plugins/web_hook.py @@ -10,6 +10,10 @@ 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) @@ -43,10 +47,13 @@ def wx_work_send(msg, key): def web_hook_send(msg): - if conf.dingtalk_token and conf.dingtalk_secret: - dingding_send(msg, conf.dingtalk_token, conf.dingtalk_secret) - if conf.wx_work_key: - wx_work_send(msg, conf.wx_work_key) + 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):