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()