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