159 lines
8.3 KiB
Python
159 lines
8.3 KiB
Python
import requests
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from pyexecplugins.pyexecplugins import PyExecPlugin
|
|
|
|
class Plugin(PyExecPlugin):
|
|
name = "HttpBridge"
|
|
|
|
taskdef = {
|
|
"name" : "pyrestbridge",
|
|
"description" : "Make a Request to an endpoint and directly stream result to another endpoint. endpoint1 and endpoint2 inputs are JSON objects containing the same keys as PyRest tasks. endpoint1 and endpoint2 outputs contain status, reason and headers of the respective endpoints. Method defaults to GET for endpoint 1 and to POST for endpoint2. Content-type for endpoint2 request will be forwarded from endpoint1 Accepts if present. If fail is set to false the operation will return as completed but will carry error information.",
|
|
"inputKeys" : ["endpoint1", "endpoint2", "fail"],
|
|
"outputKeys" : ["endpoint1", "endpoint2"],
|
|
"ownerEmail" : "m.lettere@gmail.com"
|
|
}
|
|
|
|
def __init__(self, data, config=None):
|
|
super().__init__(data, config)
|
|
self.ep1 = {}
|
|
self.ep2 = {}
|
|
logging.getLogger("pyexec").debug("%s", data)
|
|
self.ep1["method"] = data.get("endpoint1").get("method") or "get"
|
|
self.ep1["url"] = data.get("endpoint1").get("url")
|
|
self.ep1["headers"] = data.get("endpoint1").get("headers") or {}
|
|
self.ep1["body"] = data.get("endpoint1").get("body")
|
|
self.ep1["params"] = data.get("endpoint1").get("params")
|
|
self.ep1["files"] = data.get("endpoint1").get("files")
|
|
self.ep1["expect"] = data.get("endpoint1").get("expect")
|
|
|
|
self.datamap = data.get("datamap") or { "place" : "body" }
|
|
|
|
self.ep2["method"] = data.get("endpoint2").get("method") or "post"
|
|
self.ep2["url"] = data.get("endpoint2").get("url")
|
|
self.ep2["headers"] = data.get("endpoint2").get("headers") or {}
|
|
self.ep2["body"] = data.get("endpoint2").get("body")
|
|
self.ep2["params"] = data.get("endpoint2").get("params")
|
|
self.ep2["files"] = data.get("endpoint2").get("files")
|
|
self.ep2["expect"] = data.get("endpoint2").get("expect")
|
|
|
|
self.fail = data.get("fail") != False
|
|
|
|
def doRequest1(self):
|
|
self.session = requests.Session()
|
|
ep = self.ep1
|
|
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("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()
|
|
elif self.datamap["place"] == "params":
|
|
params = self.ep2["params"] or {}
|
|
params[self.datamap["name"]] = self.response1.content
|
|
self.request2 = requests.Request(ep["method"], ep["url"], headers=ep["headers"], params = params).prepare()
|
|
params.pop(self.datamap["name"])
|
|
elif self.datamap["place"] == "files":
|
|
#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":
|
|
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"
|
|
else:
|
|
if type(ep.expect) == list:
|
|
return "COMPLETED" if (response.status_code in ep["expect"]) else "FAILED"
|
|
else:
|
|
return "COMPLETED" if (response.status_code == ep["expect"]) else "FAILED"
|
|
|
|
def buildOutput(self, status1, status2):
|
|
msg = None
|
|
hdrs1 = {}
|
|
hdrs2 = {}
|
|
if status1 == "FAILED":
|
|
msg = "{} to {} failed with status {} ({}) - {}".format(self.ep1["method"], self.ep1["url"], self.response1.status_code, self.response1.reason, self.response1.text)
|
|
if self.fail:
|
|
raise Exception(msg)
|
|
else:
|
|
return {
|
|
"endpoint1" : { "status" : self.response1.status_code, "reason" : self.response1.reason, "headers" : hdrs1, "error" : msg},
|
|
"endpoint2" : { }
|
|
}
|
|
|
|
elif status2 == "FAILED":
|
|
msg = "{} to {} failed with status {} ({}) - {}".format(self.ep2["method"], self.ep2["url"], self.response2.status_code, self.response2.reason, self.response2.text)
|
|
if self.fail:
|
|
raise Exception(msg)
|
|
else:
|
|
return {
|
|
"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:
|
|
return {
|
|
"endpoint1" : { "status" : self.response1.status_code, "reason" : self.response1.reason, "headers" : hdrs1},
|
|
"endpoint2" : { "status" : self.response2.status_code, "reason" : self.response2.reason, "headers" : hdrs2}
|
|
}
|
|
|
|
|
|
def execute(self):
|
|
self.doRequest1()
|
|
status1 = self.computeStatus(self.ep1, self.response1)
|
|
status2 = "FAILED"
|
|
if status1 == "COMPLETED":
|
|
self.doRequest2()
|
|
status2 = self.computeStatus(self.ep2, self.response2)
|
|
return self.buildOutput(status1, status2)
|