diff options
| -rw-r--r-- | binaryninjaapi.h | 2 | ||||
| -rw-r--r-- | binaryninjacore.h | 4 | ||||
| -rw-r--r-- | downloadprovider.cpp | 2 | ||||
| -rw-r--r-- | python/downloadprovider.py | 99 |
4 files changed, 56 insertions, 51 deletions
diff --git a/binaryninjaapi.h b/binaryninjaapi.h index d1467d0f..5c3e0cde 100644 --- a/binaryninjaapi.h +++ b/binaryninjaapi.h @@ -4939,7 +4939,7 @@ __attribute__ ((format (printf, 1, 2))) int PerformRequest(const std::string& url, BNDownloadInstanceOutputCallbacks* callbacks); int PerformCustomRequest(const std::string& method, const std::string& url, const std::unordered_map<std::string, std::string>& headers, Response& response, BNDownloadInstanceInputOutputCallbacks* callbacks); - uint64_t ReadDataCallback(uint8_t* data, uint64_t len); + int64_t ReadDataCallback(uint8_t* data, uint64_t len); uint64_t WriteDataCallback(uint8_t* data, uint64_t len); bool NotifyProgressCallback(uint64_t progress, uint64_t total); void SetError(const std::string& error); diff --git a/binaryninjacore.h b/binaryninjacore.h index bb678720..80f38350 100644 --- a/binaryninjacore.h +++ b/binaryninjacore.h @@ -2168,7 +2168,7 @@ extern "C" struct BNDownloadInstanceInputOutputCallbacks { - uint64_t (*readCallback)(uint8_t* data, uint64_t len, void* ctxt); + int64_t (*readCallback)(uint8_t* data, uint64_t len, void* ctxt); void* readContext; uint64_t (*writeCallback)(uint8_t* data, uint64_t len, void* ctxt); void* writeContext; @@ -4783,7 +4783,7 @@ __attribute__ ((format (printf, 1, 2))) BINARYNINJACOREAPI void BNFreeDownloadInstanceResponse(BNDownloadInstanceResponse* response); BINARYNINJACOREAPI int BNPerformDownloadRequest(BNDownloadInstance* instance, const char* url, BNDownloadInstanceOutputCallbacks* callbacks); BINARYNINJACOREAPI int BNPerformCustomRequest(BNDownloadInstance* instance, const char* method, const char* url, uint64_t headerCount, const char* const* headerKeys, const char* const* headerValues, BNDownloadInstanceResponse** response, BNDownloadInstanceInputOutputCallbacks* callbacks); - BINARYNINJACOREAPI uint64_t BNReadDataForDownloadInstance(BNDownloadInstance* instance, uint8_t* data, uint64_t len); + BINARYNINJACOREAPI int64_t BNReadDataForDownloadInstance(BNDownloadInstance* instance, uint8_t* data, uint64_t len); BINARYNINJACOREAPI uint64_t BNWriteDataForDownloadInstance(BNDownloadInstance* instance, uint8_t* data, uint64_t len); BINARYNINJACOREAPI bool BNNotifyProgressForDownloadInstance(BNDownloadInstance* instance, uint64_t progress, uint64_t total); BINARYNINJACOREAPI char* BNGetErrorForDownloadInstance(BNDownloadInstance* instance); diff --git a/downloadprovider.cpp b/downloadprovider.cpp index ce703b81..d54a85e3 100644 --- a/downloadprovider.cpp +++ b/downloadprovider.cpp @@ -85,7 +85,7 @@ void DownloadInstance::PerformFreeResponse(void* ctxt, BNDownloadInstanceRespons } -uint64_t DownloadInstance::ReadDataCallback(uint8_t* data, uint64_t len) +int64_t DownloadInstance::ReadDataCallback(uint8_t* data, uint64_t len) { return BNReadDataForDownloadInstance(m_object, data, len); } diff --git a/python/downloadprovider.py b/python/downloadprovider.py index 2e43128f..b45bda19 100644 --- a/python/downloadprovider.py +++ b/python/downloadprovider.py @@ -110,16 +110,25 @@ class DownloadInstance(object): for i in range(header_count): headers[header_key_array[i]] = header_value_array[i] - # Read all data - data = b'' - while True: - read_buffer = ctypes.create_string_buffer(0x1000) - read_len = core.BNReadDataForDownloadInstance(self.handle, ctypes.cast(read_buffer, ctypes.POINTER(ctypes.c_uint8)), 0x1000) - if read_len == 0: - break - data += read_buffer[:read_len] + # Generator function that returns a chunk of data on every iteration + def data_generator(): + while True: + read_buffer = ctypes.create_string_buffer(0x1000) + read_len = core.BNReadDataForDownloadInstance(self.handle, ctypes.cast(read_buffer, ctypes.POINTER(ctypes.c_uint8)), 0x1000) + if read_len == 0: + break + if read_len < 0: + raise IOError() + yield read_buffer[:read_len] + + try: + py_response = self.perform_custom_request(method, url, headers, data_generator()) + except Exception as e: + out_response[0] = None + core.BNSetErrorForDownloadInstance(self.handle, e.__class__.__name__) + log.log_error(traceback.format_exc()) + return -1 - py_response = self.perform_custom_request(method, url, headers, data) if py_response is not None: # Assign to an instance variable so the memory stays live until the request is done self.bn_response = core.BNDownloadInstanceResponse() @@ -152,7 +161,7 @@ class DownloadInstance(object): raise NotImplementedError @abc.abstractmethod - def perform_custom_request(self, method, url, headers, data): + def perform_custom_request(self, method, url, headers, data_generator): raise NotImplementedError def _read_callback(self, data, len_, ctxt): @@ -404,37 +413,39 @@ try: return 0 - def perform_custom_request(self, method, url, headers, data): - try: - proxy_setting = settings.Settings().get_string('network.httpsProxy') - if proxy_setting: - proxies = {"https": proxy_setting} - else: - proxies = None + def perform_custom_request(self, method, url, headers, data_generator): + proxy_setting = settings.Settings().get_string('network.httpsProxy') + if proxy_setting: + proxies = {"https": proxy_setting} + else: + proxies = None - r = requests.request(pyNativeStr(method), pyNativeStr(url), headers=headers, data=data, proxies=proxies) - response = r.content - if len(response) == 0: - core.BNSetErrorForDownloadInstance(self.handle, "No data received from server!") - return None - raw_bytes = (ctypes.c_ubyte * len(response)).from_buffer_copy(response) + # Cannot have Content-Length if the body is chunked + if b"Content-Length" in headers: + del headers[b"Content-Length"] + + r = requests.request(pyNativeStr(method), pyNativeStr(url), headers=headers, data=data_generator, proxies=proxies, stream=True) + + total_size = 0 + for (key, value) in r.headers.items(): + if key.lower() == b'content-length': + total_size = int(value) + break + + bytes_sent = 0 + for chunk in r.iter_content(None): + raw_bytes = (ctypes.c_ubyte * len(chunk)).from_buffer_copy(chunk) bytes_wrote = core.BNWriteDataForDownloadInstance(self.handle, raw_bytes, len(raw_bytes)) if bytes_wrote != len(raw_bytes): core.BNSetErrorForDownloadInstance(self.handle, "Bytes written mismatch!") return None - continue_download = core.BNNotifyProgressForDownloadInstance(self.handle, bytes_wrote, bytes_wrote) + bytes_sent += bytes_wrote + continue_download = core.BNNotifyProgressForDownloadInstance(self.handle, bytes_sent, total_size) if continue_download is False: core.BNSetErrorForDownloadInstance(self.handle, "Download aborted!") return None - return DownloadInstance.Response(r.status_code, r.headers, None) - except requests.RequestException as e: - core.BNSetErrorForDownloadInstance(self.handle, e.__class__.__name__) - return None - except: - core.BNSetErrorForDownloadInstance(self.handle, "Unknown Exception!") - log.log_error(traceback.format_exc()) - return None + return DownloadInstance.Response(r.status_code, r.headers, None) class PythonDownloadProvider(DownloadProvider): name = "PythonDownloadProvider" @@ -521,7 +532,7 @@ if not _loaded and (sys.platform != "win32") and (sys.version_info >= (2, 7, 9)) return self._method return Request.get_method(self, *args, **kwargs) - def perform_custom_request(self, method, url, headers, data): + def perform_custom_request(self, method, url, headers, data_generator): result = None try: proxy_setting = settings.Settings().get_string('network.httpsProxy') @@ -529,23 +540,21 @@ if not _loaded and (sys.platform != "win32") and (sys.version_info >= (2, 7, 9)) opener = build_opener(ProxyHandler({'https': proxy_setting})) install_opener(opener) + # Cannot have Content-Length if the body is chunked if b"Content-Length" in headers: del headers[b"Content-Length"] - req = PythonDownloadInstance.CustomRequest(pyNativeStr(url), data=data, headers=headers, method=pyNativeStr(method)) + req = PythonDownloadInstance.CustomRequest(pyNativeStr(url), data=data_generator, headers=headers, method=pyNativeStr(method)) result = urlopen(req) except HTTPError as he: result = he - except URLError as e: - core.BNSetErrorForDownloadInstance(self.handle, e.__class__.__name__) - log.log_error(str(e)) - return None - except: - core.BNSetErrorForDownloadInstance(self.handle, "Unknown Exception!") - log.log_error(traceback.format_exc()) - return None - total_size = int(result.headers.get('content-length', 0)) + total_size = 0 + for (key, value) in result.headers.items(): + if key.lower() == b'content-length': + total_size = int(value) + break + bytes_sent = 0 while True: data = result.read(4096) @@ -562,10 +571,6 @@ if not _loaded and (sys.platform != "win32") and (sys.version_info >= (2, 7, 9)) core.BNSetErrorForDownloadInstance(self.handle, "Download aborted!") return None - if not bytes_sent: - core.BNSetErrorForDownloadInstance(self.handle, "Received no data!") - return None - return DownloadInstance.Response(result.getcode(), result.headers, None) class PythonDownloadProvider(DownloadProvider): |