From 9be0baca496e95bd84564b821dde36537e86a99f Mon Sep 17 00:00:00 2001
From: xxcdd <xxcdd0501@gmail.com>
Date: Sun, 31 Dec 2023 00:14:34 +0800
Subject: [PATCH] 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 <https://github.com/ekszz>
 * contributing to customize poc protocol and default port #321
 
 HomerQing <https://github.com/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 <https://github.com/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