Cleanup unpacked files

Yes, this is ugly.

plugin/idaconnect/core/core.py

@@ -1,10 +1,13 @@
+import json
 import logging
+import os
 
 import ida_idp
 import ida_kernwin
 import idaapi
 
 from ..module import Module
+from ..utilities.misc import localResource
 from ..shared.commands import Subscribe, Unsubscribe
 from hooks import Hooks, IDBHooks, IDPHooks, HexRaysHooks
 
@@ -130,6 +133,43 @@ def branch(self, uuid):
         """
         self._branch = uuid
 
+    def loadState(self):
+        """
+        Load the state file if it exists.
+        """
+        statePath = localResource('files', 'state.json')
+        if os.path.isfile(statePath):
+            with open(statePath, 'rb') as stateFile:
+                state = json.loads(stateFile.read())
+                logger.debug("Loaded state: %s" % state)
+                if state['connected']:
+                    self._plugin.network.connect(state['host'], state['port'])
+                if 'cleanup' in state and state['cleanup']:
+                    # Remove unpacked files from parent instance
+                    idbFile, idbExt = os.path.splitext(state['cleanup'])
+                    for extension in ['.id0', '.id1', '.nam', '.seg', '.til']:
+                        if os.path.exists(idbFile + extension):
+                            os.remove(idbFile + extension)
+            os.remove(statePath)
+
+    def saveState(self, cleanup=None):
+        """
+        Save the state file.
+
+        :param cleanup: the path of the file to cleanup
+        """
+        statePath = localResource('files', 'state.json')
+        with open(statePath, 'wb') as stateFile:
+            state = {
+                'connected': self._plugin.network.connected,
+                'host': self._plugin.network.host,
+                'port': self._plugin.network.port,
+            }
+            if cleanup:
+                state['cleanup'] = cleanup
+            logger.debug("Saved state: %s" % state)
+            stateFile.write(json.dumps(state))
+
     def loadNetnode(self):
         """
         Load the netnode if it exists.

plugin/idaconnect/core/events.py

@@ -532,4 +532,3 @@ def __call__(self):
         idaapi.save_user_labels(self.ea, labels)
 
         refreshPseudocodeView(True)
-

plugin/idaconnect/interface/actions.py

@@ -276,9 +276,9 @@ def _databaseDownloaded(self, branch, progress, reply):
         # Save the old database
         idbPath = idc.GetIdbPath()
         if idbPath:
-            idc.save_database(idbPath, ida_loader.DBFL_KILL)
+            idc.save_database(idbPath, 0)
         # Save the current state
-        self._plugin.network.saveState()
+        self._plugin.core.saveState(idbPath)
         # Open the new database
         QProcess.startDetached(qApp.applicationFilePath(), [filePath])
         qApp.quit()  # FIXME: Find an alternative, if any

plugin/idaconnect/network/network.py

@@ -1,14 +1,11 @@
-import json
 import logging
-import os
 
 import qt5reactor
 qt5reactor.install()  # noqa
 
 from twisted.internet import reactor
 
 from ..module import Module
-from ..utilities.misc import localResource
 from client import ClientFactory
 
 logger = logging.getLogger('IDAConnect.Network')
@@ -107,30 +104,3 @@ def sendPacket(self, packet):
         if self.connected:
             return self._factory.sendPacket(packet)
         return None
-
-    def loadState(self):
-        """
-        Load the state file if it exists.
-        """
-        statePath = localResource('files', 'state.json')
-        if os.path.isfile(statePath):
-            with open(statePath, 'rb') as stateFile:
-                state = json.loads(stateFile.read())
-                logger.debug("Loaded state: %s" % state)
-                if state['connected']:
-                    self.connect(state['host'], state['port'])
-            os.remove(statePath)
-
-    def saveState(self):
-        """
-        Save the state file.
-        """
-        statePath = localResource('files', 'state.json')
-        with open(statePath, 'wb') as stateFile:
-            state = {
-                'connected': self.connected,
-                'host': self._host,
-                'port': self._port,
-            }
-            logger.debug("Saved state: %s" % state)
-            stateFile.write(json.dumps(state))

plugin/idaconnect/plugin.py

@@ -109,7 +109,7 @@ def _init(self):
         self._network.install()
 
         # Load the current state
-        self.network.loadState()
+        self.core.loadState()
 
     def _printBanner(self):
         """
@@ -138,7 +138,7 @@ def _term(self):
         Terminate the plugin and its modules.
         """
         # Save the current state
-        self.network.saveState()
+        self.core.saveState()
 
         self._core.uninstall()
         self._interface.uninstall()