def upload_file(plugin, packet): # Save the current database with a tick=0 since it is a new snapshot plugin.core.tick = 0 plugin.core.save_netnode() input_path = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) ida_loader.save_database(input_path, 0) with open(input_path, "rb") as input_file: uncompressed_content = input_file.read() packet.content = bz2.compress(uncompressed_content) # Create the upload progress dialog text = "Uploading database to server, please wait..." progress = QProgressDialog(text, "Cancel", 0, 1) progress.setCancelButton(None) # Remove cancel button progress.setModal(True) # Set as a modal dialog window_flags = progress.windowFlags() # Disable close button progress.setWindowFlags(window_flags & ~Qt.WindowCloseButtonHint) progress.setWindowTitle("Save to server") icon_path = plugin.plugin_resource("upload.png") progress.setWindowIcon(QIcon(icon_path)) # Send the packet to upload the file packet.upback = partial(SaveActionHandler._on_progress, progress) d = plugin.network.send_packet(packet) if d: d.add_callback( partial(SaveActionHandler.file_uploaded, plugin, progress)) d.add_errback(plugin.logger.exception) progress.show()
def _dialog_accepted(self, dialog): repo, branch = dialog.get_result() self._plugin.core.repo = repo.name self._plugin.core.branch = branch.name # Save the current database self._plugin.core.save_netnode() inputPath = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) ida_loader.save_database(inputPath, 0) # Create the packet that will hold the database packet = UploadDatabase.Query(repo.name, branch.name) with open(inputPath, 'rb') as inputFile: packet.content = inputFile.read() # Create the progress dialog text = "Uploading database to server, please wait..." progress = QProgressDialog(text, "Cancel", 0, 1) progress.setCancelButton(None) # Remove cancel button progress.setModal(True) # Set as a modal dialog windowFlags = progress.windowFlags() # Disable close button progress.setWindowFlags(windowFlags & ~Qt.WindowCloseButtonHint) progress.setWindowTitle("Save to server") iconPath = self._plugin.resource('upload.png') progress.setWindowIcon(QIcon(iconPath)) # Send the packet to upload the file packet.upback = partial(self._on_progress, progress) d = self._plugin.network.send_packet(packet) d.add_callback(partial(self._database_uploaded, repo, branch, progress)) d.add_errback(logger.exception) progress.show()
def _file_downloaded(self, database, progress, reply): """Called when the file has been downloaded.""" progress.close() # Get the absolute path of the file app_path = QCoreApplication.applicationFilePath() app_name = QFileInfo(app_path).fileName() file_ext = "i64" if "64" in app_name else "idb" file_name = "%s_%s.%s" % (database.project, database.name, file_ext) file_path = self._plugin.user_resource("files", file_name) # Write the file to disk with open(file_path, "wb") as output_file: output_file.write(reply.content) self._plugin.logger.info("Saved file %s" % file_name) # Save the old database database = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if database: ida_loader.save_database(database, ida_loader.DBFL_TEMP) # This is a very ugly hack used to open a database into IDA. We don't # have any function for this in the SDK, so I sorta hijacked the # snapshot functionality in this effect. # Get the library to call functions not present in the bindings idaname = "ida64" if "64" in app_name else "ida" if sys.platform == "win32": dllname, dlltype = idaname + ".dll", ctypes.windll elif sys.platform == "linux2": dllname, dlltype = "lib" + idaname + ".so", ctypes.cdll elif sys.platform == "darwin": dllname, dlltype = "lib" + idaname + ".dylib", ctypes.cdll dllpath = ida_diskio.idadir(None) if not os.path.exists(os.path.join(dllpath, dllname)): dllpath = dllpath.replace("ida64", "ida") dll = dlltype[os.path.join(dllpath, dllname)] # Close the old database using the term_database library function old_path = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if old_path: dll.term_database() # Open the new database using the init_database library function # This call only won't be enough because the user interface won't # be initialized, this is why the snapshot functionality is used for args = [app_name, file_path] argc = len(args) argv = (ctypes.POINTER(ctypes.c_char) * (argc + 1))() for i, arg in enumerate(args): arg = arg.encode("utf-8") argv[i] = ctypes.create_string_buffer(arg) v = ctypes.c_int(0) av = ctypes.addressof(v) pv = ctypes.cast(av, ctypes.POINTER(ctypes.c_int)) dll.init_database(argc, argv, pv) # Create a temporary copy of the new database because we cannot use # the snapshot functionality to restore the currently opened database file_ext = ".i64" if "64" in app_name else ".idb" tmp_file, tmp_path = tempfile.mkstemp(suffix=file_ext) shutil.copyfile(file_path, tmp_path) # This hook is used to delete the temporary database when all done class UIHooks(ida_kernwin.UI_Hooks): def database_inited(self, is_new_database, idc_script): self.unhook() os.close(tmp_file) if os.path.exists(tmp_path): os.remove(tmp_path) hooks = UIHooks() hooks.hook() # Call the restore_database_snapshot library function # This will initialize the user interface, completing the process s = ida_loader.snapshot_t() s.filename = tmp_path # Use the temporary database ida_kernwin.restore_database_snapshot(s, None, None)
def _database_downloaded(self, branch, progress, reply): """ Called when the file has been downloaded. :param branch: the branch :param progress: the progress dialog :param reply: the reply from the server """ # Close the progress dialog progress.close() # Get the absolute path of the file appPath = QCoreApplication.applicationFilePath() appName = QFileInfo(appPath).fileName() fileExt = 'i64' if '64' in appName else 'idb' fileName = '%s_%s.%s' % (branch.repo, branch.name, fileExt) filePath = local_resource('files', fileName) # Write the packet content to disk with open(filePath, 'wb') as outputFile: outputFile.write(reply.content) logger.info("Saved file %s" % fileName) # Save the old database database = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if database: ida_loader.save_database(database, ida_loader.DBFL_TEMP) # Get the dynamic library idaname = 'ida64' if '64' in appName else 'ida' if sys.platform == 'win32': dllname, dlltype = idaname + '.dll', ctypes.windll elif sys.platform == 'linux2': dllname, dlltype = 'lib' + idaname + '.so', ctypes.cdll elif sys.platform == 'darwin': dllname, dlltype = 'lib' + idaname + '.dylib', ctypes.cdll dllpath = ida_diskio.idadir(None) if not os.path.exists(os.path.join(dllpath, dllname)): dllpath = dllpath.replace('ida64', 'ida') dll = dlltype[os.path.join(dllpath, dllname)] # Close the old database oldPath = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) if oldPath: dll.term_database() # Open the new database LP_c_char = ctypes.POINTER(ctypes.c_char) args = [appName, filePath] argc = len(args) argv = (LP_c_char * (argc + 1))() for i, arg in enumerate(args): arg = arg.encode('utf-8') argv[i] = ctypes.create_string_buffer(arg) LP_c_int = ctypes.POINTER(ctypes.c_int) v = ctypes.c_int(0) av = ctypes.addressof(v) pv = ctypes.cast(av, LP_c_int) dll.init_database(argc, argv, pv) # Create a copy of the new database fileExt = '.i64' if '64' in appName else '.idb' tmpFile, tmpPath = tempfile.mkstemp(suffix=fileExt) shutil.copyfile(filePath, tmpPath) class UIHooks(ida_kernwin.UI_Hooks): def database_inited(self, is_new_database, idc_script): self.unhook() # Remove the tmp database os.close(tmpFile) if os.path.exists(tmpPath): os.remove(tmpPath) hooks = UIHooks() hooks.hook() # Open the tmp database s = ida_loader.snapshot_t() s.filename = tmpPath ida_kernwin.restore_database_snapshot(s, None, None)
def activate(self, context): """ Called when the action has been clicked by the user """ idb_path = ida_loader.get_path(ida_loader.PATH_TYPE_IDB) ida_loader.save_database(idb_path, 0) with open(idb_path, 'rb') as input_file: idb_data = input_file.read() # Create the progress bar progress_bar = QProgressDialog( 'Uploading your idb to the server, please wait...', 'Cancel', 0, 1 ) progress_bar.setCancelButton(None) # Remove the cancel button so the user won't be able to cancel progress_bar.setModal(True) # Set this as a modal dialog window_flags = progress_bar.windowFlags() # Disable close button progress_bar.setWindowFlags(window_flags & ~Qt.WindowCloseButtonHint) progress_bar.setWindowTitle('Upload to server') idb_data_len = len(idb_data) idb_data_stream = io.BytesIO(idb_data) idb_hash = hashlib.sha1(idb_data).digest() total_packets = (idb_data_len / self.CHUNK_SIZE) total_packets = total_packets + 1 if idb_data_len % self.CHUNK_SIZE != 0 else total_packets # correction # Build the initiation packet here self._plugin.network_manager.send_request( RequestType.UPLOAD_IDB_START, idb_name=Unicoder.decode(os.path.split(idb_path)[-1]), idb_hash=idb_hash, idb_size=idb_data_len, ) self._plugin.logger.debug('starting to send packets') self._plugin.logger.debug('The amount of packets needed to be sent: {}'.format(total_packets)) def _on_error(progress_bar): progress_bar.close() success = QMessageBox() success.setIcon(QMessageBox.Critical) success.setStandardButtons(QMessageBox.Ok) success.setText("Could not upload IDB") success.setWindowTitle("Upload to server FAILED") success.exec_() for i in range(total_packets): current_pkt_data = idb_data_stream.read(self.CHUNK_SIZE) self._plugin.network_manager.send_request( RequestType.IDB_CHUNK, callback=partial(self._update_progress, progress_bar, i, total_packets), err_callback=partial(_on_error, progress_bar), data=current_pkt_data, ) self._plugin.logger.debug('finished sending packets') self._plugin.logger.debug('sending upload_end') def _close_window(progress_bar): SaveMenuActionHandler._update_progress(progress_bar, 100, 100) progress_bar.close() success = QMessageBox() success.setIcon(QMessageBox.Information) success.setStandardButtons(QMessageBox.Ok) success.setText("IDB successfully uploaded!") success.setWindowTitle("Upload to server") success.exec_() self._plugin.network_manager.send_request( RequestType.IDB_END, callback=partial(_close_window, progress_bar) ) progress_bar.show()