import gradio as gr
import shutil
import subprocess
import tempfile
import os
import sys
import json

def get_functions(file):

    with tempfile.TemporaryDirectory() as TEMP_DIR:

        subprocess.run(f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {file} -postscript /home/user/app/scripts/dump_functions.py {TEMP_DIR}/funcs.json", shell=True)

        json_funcs = json.load(open(f"{TEMP_DIR}/funcs.json"))

        return json_funcs

with gr.Blocks() as demo:

    all_dis_state = gr.State()

    intro = gr.Markdown(
        """
    # DIRTY-Ghidra Testing
    First, upload a binary.
    """
    )

    file_widget = gr.File(label="Executable file")

    with gr.Column(visible=False) as col:
        #output = gr.Textbox("Output")

        gr.Markdown("""
        Great, you selected an executable!  Now pick the function you would like to analyze.                    
                    """)

        fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)

        gr.Markdown("""
        Below you can find the selected function's disassembly, and the model's
        prediction of whether the function is an object-oriented method or a
        regular function.
                    """)

        with gr.Row(visible=True) as result:
            disassembly = gr.Textbox(label="Disassembly", lines=20)
            with gr.Column():
                clazz = gr.Label()
                #interpret_button = gr.Button("Interpret (very slow)")
            #interpretation = gr.components.Interpretation(disassembly)

    def file_change_fn(file):

        if file is None:
            return {
                    col: gr.update(visible=False),
                    all_dis_state: None
                    }
        else:

            try:
                progress = gr.Progress()
                progress(0, desc="Analyzing binary...")
                fun_data = get_functions(file.name)
                #print(fun_data)

                addrs = [(f"{name} ({hex(int(addr))})", int(addr)) for addr, name in fun_data.items()]
            except:
                raise gr.Error("Unable to obtain functions")

            return {
                    col: gr.Column(visible=True),
                    fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
                    all_dis_state: fun_data
                    }

    # Need to put intro as output to get progress to work!
    file_widget.change(file_change_fn, file_widget, outputs=[intro, col, fun_dropdown, all_dis_state])

# spaces only shows stderr..
os.dup2(sys.stdout.fileno(), sys.stderr.fileno())

demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860)