Merge pull request #45 from ioncodes/ida9-improvements

IDA 9 improvements

Author: mrexodia

README.md

@@ -9,7 +9,7 @@ IDACode is still in a very early state and bugs are to be expected. Please open
 * **Modularity**: IDACode does not make extensive use of safe wrappers for thread synchronization, this allows you to import any module from any path at any given time. Instead IDACode synchronizes the script execution thread with IDAs main thread to avoid performance and unexpected issues.
 * **Syncing**: As IDACode uses `debugpy` for communication, it syncs the output window naturally with VS Code's output panel.
 
-IDACode supports both Python 2 and Python 3!
+IDACode has been tested on Windows and macos with IDA 8.4/9.0 and Python 3.12 (older python versions have issues with debugging).
 
 ## Setup
 To set up the dependencies for the IDA plugin run:
@@ -22,12 +22,11 @@ python -m pip install --user debugpy tornado
 
 Either clone this repository or download a release package from [here](https://github.com/ioncodes/idacode/releases). `ida.zip` reflects the contents of the `ida` folder in this repository. Copy all files into IDAs plugin directory.  
 
-The next step is to configure your settings to match your environment. Edit `idacode_utils/settings.py` accordingly:
+The next step is to configure your settings to match your environment (optional). Edit `idacode_utils/settings.py` accordingly:
 
 * `HOST`: This is the host address. This is always `127.0.0.1` unless you want it to be accessible from a remote location. **Keep in mind that this plugin does not make use of authentication.**
 * `PORT`: This is the port you want IDA to listen to. This is used for websocket communication between IDA and VS Code.
 * `DEBUG_PORT`: This is the port you want to listen on for incoming debug sessions.
-* `PYTHON`: This is the absolute path to the Python distribution that your IDA setup uses.
 * `LOGGING`: Determines whether the debugger should log into files. This is especially useful when you are running into issues with IDACode. Please submit a new issue if you find anything. The files are always located in your temp directory (e.g. Windows: `%TEMP%`). The files are called `debugpy.*.log`.
 
 You can now start the plugin by clicking on `IDACode` in the plugins menu.  

ida/idacode_utils/hooks.py

@@ -4,17 +4,15 @@
 getcwd_original = os.getcwd
 
 def getcwd_hook():
-    global script_folder
-
-    cwd = getcwd_original()
-    if cwd.lower() in script_folder.lower() and script_folder.lower() != cwd.lower():
-        cwd = script_folder
-    return cwd
+    # NOTE: We return the script folder here, otherwise breakpoints fail in VSCode
+    if script_folder:
+        return script_folder
+    return getcwd_original()
 
 def set_script_folder(folder):
     global script_folder
 
     script_folder = folder
 
 def install():
-    os.getcwd = getcwd_hook
+    os.getcwd = getcwd_hook

ida/idacode_utils/plugin.py

@@ -1,44 +1,93 @@
-import socket, sys, os, threading, inspect, subprocess
+import sys, threading, subprocess, logging
+import idacode_utils.settings as settings
 try:
     import tornado, debugpy
 except ImportError:
-    print("[IDACode] Dependencies missing, run: python -m pip install --user debugpy tornado")
+    print("[IDACode] Dependencies missing, run:\n  \"{}\" -m pip install --user debugpy tornado".format(settings.PYTHON))
     sys.exit()
 import idaapi
-import idacode_utils.dbg as dbg
-import idacode_utils.hooks as hooks
-import idacode_utils.settings as settings
 from idacode_utils.socket_handler import SocketHandler
 
+# Source: https://github.com/OALabs/hexcopy-ida/blob/8b0b2a3021d7dc9010c01821b65a80c47d491b61/hexcopy.py#L30
+major, minor = map(int, idaapi.get_kernel_version().split("."))
+using_ida7api = (major > 6)
+using_pyqt5 = using_ida7api or (major == 6 and minor >= 9)
+
+if using_pyqt5:
+    import PyQt5.QtWidgets as QtWidgets
+else:
+    import PySide.QtGui as QtGui
+    QtWidgets = QtGui
+
 # Fix for https://github.com/tornadoweb/tornado/issues/2608
 if sys.version_info >= (3, 4):
     import asyncio
 
-VERSION = "0.3.0"
-initialized = False
-
-def setup_patches():
-    hooks.install()
-    #sys.executable = settings.PYTHON
-
-def create_socket_handler():
-    if sys.version_info >= (3, 4):
-        asyncio.set_event_loop(asyncio.new_event_loop())
-    app = tornado.web.Application([
-        (r"/ws", SocketHandler),
-    ])
-    server = tornado.httpserver.HTTPServer(app)
-    print("[IDACode] Listening on {address}:{port}".format(address=settings.HOST, port=settings.PORT))
-    server.listen(address=settings.HOST, port=settings.PORT)
-
-def start_server():
-    # Fix for https://github.com/tornadoweb/tornado/issues/2608
-    if sys.platform=='win32' and sys.version_info >= (3,8):
-        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-    
-    setup_patches()
-    create_socket_handler()
-    tornado.ioloop.IOLoop.current().start()
+def join_gui_thread(thread: threading.Thread, timeout=None):
+    iterations = 0
+    iteration_timeout = 0.1
+    while True:
+        if not thread.is_alive():
+            return True
+        thread.join(iteration_timeout)
+        QtWidgets.QApplication.processEvents()
+        if timeout is not None and iteration_timeout * iterations >= timeout:
+            return False
+        iterations += 1
+
+class Server:
+    def __init__(self):
+        self.started = False
+        self.server: tornado.httpserver.HTTPServer = None
+        self.thread: threading.Thread = None
+
+    def start(self):
+        self.stop()
+        self.thread = threading.Thread(target=self.server_thread)
+        self.thread.start()
+        self.started = True
+
+    def stop(self):
+        if not self.started:
+            return
+
+        if self.server is not None:
+            self.io_loop.add_callback(self.server.stop)
+            self.io_loop.add_callback(self.server.close_all_connections)
+            self.io_loop.add_callback(self.io_loop.stop)
+
+        if not join_gui_thread(self.thread, 1.0):
+            print("[IDACode] Waiting for server to stop...")
+            if not join_gui_thread(self.thread, 5.0):
+                print("[IDACode] deadlock while stopping server, please report an issue!\n")
+        self.thread = None
+        self.server = None
+        print("[IDACode] Server stopped")
+
+    def server_thread(self):
+        # Create a new event loop for the thread
+        # https://github.com/tornadoweb/tornado/issues/2308#issuecomment-372582005
+        loop = asyncio.new_event_loop()
+        loop.set_debug(False)
+        logging.getLogger("asyncio").setLevel(logging.CRITICAL)  # Remove some debug spam
+        asyncio.set_event_loop(loop)
+
+        # Before starting the event loop, instantiate a WebSocketClient and add a
+        # callback to the event loop to start it. This way the first thing the
+        # event loop does is to start the client.
+        self.io_loop = tornado.ioloop.IOLoop.current()
+        app = tornado.web.Application([
+            (r"/ws", SocketHandler),
+        ])
+        self.server = tornado.httpserver.HTTPServer(app)
+        print("[IDACode] Listening on {address}:{port}".format(address=settings.HOST, port=settings.PORT))
+        self.server.listen(address=settings.HOST, port=settings.PORT)
+
+        # Start the event loop.
+        self.io_loop.start()
+
+        # Signal that the service is finished
+        self.started = False
 
 def get_python_versions():
     settings_version = subprocess.check_output([settings.PYTHON, "-c", "import sys; print(sys.version + sys.platform)"])
@@ -47,36 +96,27 @@ def get_python_versions():
     return (settings_version, ida_version)
 
 class IDACode(idaapi.plugin_t):
-    def __init__(self):
-        self.flags = idaapi.PLUGIN_UNL
-        self.comment = "IDACode"
-        self.help = "IDACode"
-        self.wanted_name = "IDACode"
-        self.wanted_hotkey = ""
+    flags = idaapi.PLUGIN_KEEP
+    comment = "IDACode"
+    help = "IDACode"
+    wanted_name = "IDACode"
+    wanted_hotkey = "Ctrl-Shift-I"
 
     def init(self):
-        global initialized
-        if not initialized:
-            initialized = True
-            if os.path.isfile(settings.PYTHON):
-                settings_version, ida_version = get_python_versions()
-                if settings_version != ida_version:
-                    print("[IDACode] settings.PYTHON version mismatch, aborting load:")
-                    print("[IDACode] IDA interpreter: {}".format(ida_version))
-                    print("[IDACode] settings.PYTHON: {}".format(settings_version))
-                    return idaapi.PLUGIN_SKIP
-            else:
-                print("[IDACode] settings.PYTHON ({}) does not exist, aborting load".format(settings.PYTHON))
-                print("[IDACode] To fix this issue, modify idacode_utils/settings.py to point to the python executable")
-                return idaapi.PLUGIN_SKIP
-            print("[IDACode] Plugin version {}".format(VERSION))
-            print("[IDACode] Plugin loaded, use Edit -> Plugins -> IDACode to start the server")
-        return idaapi.PLUGIN_OK
+        settings_version, ida_version = get_python_versions()
+        if settings_version != ida_version:
+            print("[IDACode] settings.PYTHON version mismatch, aborting load:")
+            print("[IDACode] IDA interpreter: {}".format(ida_version))
+            print("[IDACode] settings.PYTHON: {}".format(settings_version))
+            return idaapi.PLUGIN_SKIP
+
+        self.server = Server()
+        print("[IDACode] Plugin version 0.4.0")
+        print("[IDACode] Plugin loaded, use Edit -> Plugins -> IDACode to start the server")
+        return idaapi.PLUGIN_KEEP
 
     def run(self, args):
-        thread = threading.Thread(target=start_server)
-        thread.daemon = True
-        thread.start()
+        self.server.start()
 
     def term(self):
-        pass
+        self.server.stop()

ida/idacode_utils/settings.py

@@ -1,5 +1,30 @@
 HOST = "127.0.0.1"
 PORT = 7065
 DEBUG_PORT = 7066
-PYTHON = "C:\\Program Files\\IDA 7.5\\python38\\python.exe"
-LOGGING = False
+ALLOW_UNSAFE_ORIGIN = False
+LOGGING = True
+
+# Heuristically detect the python executable path
+PYTHON = ""
+import sys
+import os
+import platform
+for path in sys.path:
+    if platform.system() == "Windows":
+        path = path.replace("/", "\\")
+
+    split = path.split(os.sep)
+    if split[-1].endswith(".zip"):
+        path = os.path.dirname(path)
+        if platform.system() == "Windows":
+            python_executable = os.path.join(path, "python.exe")
+        else:
+            python_executable = os.path.join(path, "..", "bin", "python3")
+        python_executable = os.path.abspath(python_executable)
+
+        if os.path.exists(python_executable):
+            print("[IDACode] Detected python executable: " + python_executable)
+            PYTHON = python_executable
+            break
+if len(PYTHON) == 0 or not os.path.exists(PYTHON):
+    raise FileNotFoundError("[IDACode] Could not find python executable, report an issue!")

ida/idacode_utils/socket_handler.py

@@ -12,22 +12,42 @@ def create_env():
         "__name__": "__main__"
     }
 
+debugpy_host = ""
+debugpy_port = 0
+
 def start_debug_server():
+    # At most one instance of debugpy can ever be created per process.
+    # Reference: https://github.com/microsoft/debugpy/issues/297
+    global debugpy_host, debugpy_port
+    if debugpy_port and debugpy_port:
+        print("[IDACode] debugpy server is already listening on {}:{}".format(debugpy_host, debugpy_port))
+        return
+
+    # Install hook for os.getcwd
+    hooks.install()
+
+    # Start debugpy server
     if settings.LOGGING:
         tmp_path = tempfile.gettempdir()
         debugpy.log_to(tmp_path)
         print("[IDACode] Logging to {} with pattern debugpy.*.log".format(tmp_path))
-    debugpy.configure({ "python": settings.PYTHON })
-    debugpy.listen((settings.HOST, settings.DEBUG_PORT))
-    print("[IDACode] IDACode debug server listening on {address}:{port}".format(address=settings.HOST, port=settings.DEBUG_PORT))
+    debugpy.configure(python=settings.PYTHON)
+    debugpy_host, debugpy_port = debugpy.listen((settings.HOST, settings.DEBUG_PORT))
+    print("[IDACode] Started debugpy server on {}:{}".format(debugpy_host, debugpy_port))
 
 class SocketHandler(tornado.websocket.WebSocketHandler):
+    def check_origin(self, origin):
+        # NOTE: This is called when connecting from a browser
+        return settings.ALLOW_UNSAFE_ORIGIN
+
     def open(self):
         print("[IDACode] Client connected")
 
     def on_message(self, message):
-        message = json.loads(message.decode("utf8"))
-        
+        if not isinstance(message, str):
+            message = message.decode("utf-8")
+        message = json.loads(message)
+
         if message["event"] == "set_workspace":
             path = message["path"]
             hooks.set_script_folder(path)