ida_babel.py

import idaapi
from http.server import BaseHTTPRequestHandler, HTTPServer
import cgi
import tempfile
import os
import threading

DEBUG = False
HOST = "localhost"
PORT = 4000


def debug(message):
    if DEBUG:
        print(message)


class IDAPythonRequestHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        try:
            content_type, pdict = cgi.parse_header(self.headers["Content-Type"])
            debug(f"Content-Type: {content_type}")

            if content_type == "multipart/form-data":
                pdict["boundary"] = bytes(pdict["boundary"], "utf-8")
                form_data = cgi.parse_multipart(self.rfile, pdict)

                code_content = form_data.get("code", [None])[0]

                if not code_content:
                    raise ValueError("No 'code' parameter provided.")

                if isinstance(code_content, bytes):
                    code_content = code_content.decode("utf-8")

                with tempfile.NamedTemporaryFile(
                    delete=False, suffix=".py"
                ) as temp_code_file:
                    temp_code_file.write(code_content.encode("utf-8"))
                    temp_code_file_path = temp_code_file.name

                try:
                    debug(f"Executing the following code:\n{code_content}")

                    def execute_code():
                        try:
                            exec(code_content, globals(), locals())
                            debug("Executed code successfully.")
                        except Exception as e:
                            debug(f"Error during exec: {e}")

                    idaapi.execute_sync(execute_code, idaapi.MFF_FAST)

                    debug("Output should be visible in the IDA console.")

                finally:
                    os.remove(temp_code_file_path)

                self.send_response(200)
                # self.send_header("Content-type", "text/plain")
                self.end_headers()
            else:
                raise ValueError(
                    "Invalid content type. Expected 'multipart/form-data'."
                )
        except Exception as e:
            self.send_response(500)
            self.end_headers()
            error_message = f"Execution failed: {e}"
            print(f"Error: {error_message}")
            self.wfile.write(error_message.encode("utf-8"))


class IDABabelServer:
    def __init__(self, host=HOST, port=PORT):
        self.host = host
        self.port = port
        self.server = None
        self.thread = None

    def start(self):
        if self.server is not None:
            debug("Server is already running.")
            return

        print(f"Starting IDA Babel Server at http://{self.host}:{self.port}")
        self.server = HTTPServer((self.host, self.port), IDAPythonRequestHandler)
        self.thread = threading.Thread(target=self.server.serve_forever)
        self.thread.daemon = True
        self.thread.start()
        print("IDA Babel Server started and running.")

    def stop(self):
        if self.server is None:
            print("Server is not running.")
            return

        print("Stopping IDA Babel Server.")
        self.server.shutdown()
        self.server.server_close()
        self.thread.join()
        self.server = None
        self.thread = None
        print("IDA Babel Server stopped.")


ida_babel_server = IDABabelServer()


class UiAction(idaapi.action_handler_t):
    def __init__(self, id, name, tooltip, menuPath, callback):
        super().__init__()
        self.id = id
        self.name = name
        self.tooltip = tooltip
        self.menuPath = menuPath
        self.callback = callback

    def registerAction(self):
        action_desc = idaapi.action_desc_t(
            self.id,
            self.name,
            self,
            "",
            self.tooltip,
            0,
        )
        if idaapi.register_action(action_desc):
            idaapi.attach_action_to_menu(self.menuPath, self.id, idaapi.SETMENU_APP)
            return True
        return False

    def unregisterAction(self):
        idaapi.detach_action_from_menu(self.menuPath, self.id)
        idaapi.unregister_action(self.id)

    def activate(self, ctx):
        self.callback(ctx)
        return 1

    def update(self, ctx):
        return idaapi.AST_ENABLE_ALWAYS


class IDAPlugin(idaapi.plugin_t):
    flags = idaapi.PLUGIN_UNL
    comment = "IDA Babel Server Plugin"
    help = "Runs an HTTP server to execute Python code in IDA from emacs"
    wanted_name = "IDA Babel Server"
    wanted_hotkey = ""

    def init(self):
        actions = [
            {
                "id": "idababel:start",
                "name": "Start IDA Babel Server",
                "tooltip": "Start the IDA Babel server.",
                "callback": self.start_server,
            },
            {
                "id": "idababel:stop",
                "name": "Stop IDA Babel Server",
                "tooltip": "Stop the IDA Babel server.",
                "callback": self.stop_server,
            },
        ]

        for action in actions:
            act = UiAction(
                id=action["id"],
                name=action["name"],
                tooltip=action["tooltip"],
                menuPath="File/Plugins/IDA Babel Server/",
                callback=action["callback"],
            )
            if not act.registerAction():
                debug(f"Failed to register action {action['id']}")
                return idaapi.PLUGIN_SKIP

        debug("Plugin initialised successfully.")
        return idaapi.PLUGIN_KEEP

    def start_server(self, ctx):
        if ida_babel_server.server is None:
            ida_babel_server.start()
        else:
            debug("Server is already running.")

    def stop_server(self, ctx):
        if ida_babel_server.server is not None:
            ida_babel_server.stop()
        else:
            debug("Server is not running.")

    def run(self, arg):
        pass

    def term(self):
        ida_babel_server.stop()
        debug("IDA Babel Server stopped")


def PLUGIN_ENTRY():
    return IDAPlugin()