From 1a85ae9f8d158a50a7a498b1e0ed03918b70586a Mon Sep 17 00:00:00 2001
From: dcore94 <m.lettere@gmail.com>
Date: Mon, 11 Oct 2021 17:43:59 +0200
Subject: [PATCH] moved multipart to streaming

---
 pyexecplugins/HttpBridge.py | 57 ++++++++++++++++++++++++++++++-------
 1 file changed, 47 insertions(+), 10 deletions(-)

diff --git a/pyexecplugins/HttpBridge.py b/pyexecplugins/HttpBridge.py
index 9752599..c3901a1 100644
--- a/pyexecplugins/HttpBridge.py
+++ b/pyexecplugins/HttpBridge.py
@@ -1,6 +1,7 @@
 import requests
 import json
 import logging
+from datetime import datetime
 from pyexecplugins.pyexecplugins import PyExecPlugin
 
 class Plugin(PyExecPlugin):
@@ -45,11 +46,12 @@ class Plugin(PyExecPlugin):
         logging.getLogger("pyexec").debug("%s - %s", ep["method"], ep["url"])
         self.request1 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], data = ep["body"], params = ep["params"]).prepare()
         self.response1 = self.session.send(self.request1, stream=True)
+        self.content_type = self.response1.headers["Content-type"]
         return self.response1
     
     def doRequest2(self):
         ep = self.ep2
-        logging.getLogger("pyexec").debug("%s - %s", ep["method"], ep["url"])
+        logging.getLogger("pyexec").debug("Entering request2 %s - %s", ep["method"], ep["url"])
         
         if self.datamap["place"] == "body":
         	self.request2 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], data = self.response1.content).prepare()
@@ -59,19 +61,54 @@ class Plugin(PyExecPlugin):
         	self.request2 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], params = params).prepare()
         	params.pop(self.datamap["name"])
         elif self.datamap["place"] == "files":
-        	files = self.ep2["files"] or {}
-        	files[self.datamap["name"]] = self.response1.content
-        	self.request2 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], files = files).prepare()
-        	files.pop(self.datamap["name"])
+                #encode custom multipart in order to exploit streaming
+                files = self.ep2["files"] or {}
+                logging.getLogger("pyexec").debug("Files are - %s", files)
+                headers = ep["headers"]
+                headers["Content-Type"] = "multipart/form-data; boundary=PYRESTBRIDGE_BOUNDARY"
+                boundary_line = "--PYRESTBRIDGE_BOUNDARY"
+                prolog = ""
+                for k in files:
+                    prolog += boundary_line + "\r\n"
+                    prolog += 'Content-Disposition: form-data; name="%s"' % k + "\r\n\r\n"
+                    prolog += files[k] + "\r\n"
+
+                prolog += boundary_line + "\r\n"
+                prolog += 'Content-Disposition: form-data; name="%s"' % "file" + "\r\n"
+                prolog += "Content-Type: " + self.response1.headers["Content-Type"] + "\r\n\r\n"
+
+                epilog = "\r\n" + boundary_line + "--"
+                logging.getLogger("pyexec").debug("Request2 preparing %s", files)
+                self.request2 = requests.Request(ep["method"], ep["url"], headers=headers, data = self.chunkedUpload(prolog, self.response1, epilog, 1000000)).prepare()
+                logging.getLogger("pyexec").debug("Request2 sending %s", datetime.now())
+                
+                #files[self.datamap["name"]] = (self.datamap["name"], self.response1.raw.stream())
+                #m = MultipartEncoder(fields=files)
+                #logging.getLogger("pyexec").debug("Files are - %s", files)
+                #headers = ep["headers"]
+                #headers["Content-Type"] = m.content_type
+                #self.request2 = requests.Request(ep["method"], ep["url"], headers=headers, data = m).prepare()
+                #self.request2 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], files = files).prepare()
+                #files.pop(self.datamap["name"])
         elif self.datamap["place"] == "headers":
-        	params = self.ep2["headers"] or {}
-        	params[self.datamap["name"]] = self.response1.content
-        	self.request2 = requests.Request(ep["method"], ep["url"], headers=headers, data = ep["body"], params=ep["params"], files = files).prepare()
+        	headers = self.ep2["headers"] or {}
+        	headers[self.datamap["name"]] = self.response1.content
+        	self.request2 = requests.Request(ep["method"], ep["url"], headers=headers, data = ep["body"], params=ep["params"]).prepare()
         	
         self.response2 = self.session.send(self.request2, stream=True)
+        logging.getLogger("pyexec").debug("Request2 sentg %s", datetime.now())
         self.session.close()
         return self.response2
-    
+   
+    def chunkedUpload(self, prolog, response1, epilog, chunksize=1024):
+        logging.getLogger("pyexec").debug("Sending prolog - %s", datetime.now())
+        yield str.encode(prolog)
+        logging.getLogger("pyexec").debug("Sending chunks - %s", datetime.now())
+        for chunk in response1.iter_content(chunksize):
+            yield chunk
+        logging.getLogger("pyexec").debug("Sending epilog - %s", datetime.now())
+        yield str.encode(epilog)
+
     def computeStatus(self, ep, response):
         if ep["expect"] == None: 
             return "COMPLETED" if response.ok else "FAILED"
@@ -101,7 +138,7 @@ class Plugin(PyExecPlugin):
                 raise Exception(msg)
             else:
                 return {
-                    "endpoint1" : { "status" : self.response1.status_code, "reason" : self.response1.reason, "headers" : hdrs1, "error" : msg},
+                    "endpoint1" : { "status" : self.response1.status_code, "reason" : self.response1.reason, "headers" : hdrs1},
                     "endpoint2" : { "status" : self.response2.status_code, "reason" : self.response2.reason, "headers" : hdrs2, "error" : msg}
                 }
         else: