File size: 3,285 Bytes
86a2bff
 
58773bb
f9437fe
9cbc665
58773bb
 
b1fd156
1d56a5d
8ef67fc
25d463a
2bce853
 
 
25d463a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1d56a5d
8ef67fc
cecf2b8
1d56a5d
f3fdb65
 
 
 
 
 
 
10ca2d7
 
 
 
f3fdb65
 
 
 
 
c485ef9
 
 
1d7fc0d
 
1f9e351
1d7fc0d
c485ef9
f3fdb65
 
 
 
 
 
 
8c2cd7c
2c26c26
f3fdb65
76d23b8
 
f3fdb65
76d23b8
c485ef9
 
0ed5800
8c2cd7c
1d7fc0d
 
96362ce
1d7fc0d
b4f4d50
4c721b6
b4f4d50
f3fdb65
 
 
1f9e351
0ed5800
f3fdb65
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import gradio as gr

import os
import re
import subprocess
import tempfile

model = gr.load("ejschwartz/oo-method-test-model-bylibrary", src="models")

def get_all_dis(bname, addrs=None):

    anafile = tempfile.NamedTemporaryFile(prefix=os.path.basename(bname) + "_", suffix=".bat_ana")
    ananame = anafile.name

    addrstr = ""
    if addrs is not None:
        addrstr = " ".join([f"--function-at {x}" for x in addrs])

    subprocess.check_output(f"bat-ana {addrstr} --no-post-analysis -o {ananame} {bname} 2>/dev/null", shell=True)


    output = subprocess.check_output(f"bat-dis --no-insn-address --no-bb-cfg-arrows --color=off {ananame} 2>/dev/null", shell=True)
    output = re.sub(b' +', b' ', output)

    func_dis = {}
    last_func = None
    current_output = []

    for l in output.splitlines():
        if l.startswith(b";;; function 0x"):
            if last_func is not None:
                func_dis[last_func] = b"\n".join(current_output)
            last_func = int(l.split()[2], 16)
            current_output.clear()

        if not b";;" in l:
            current_output.append(l)

    if last_func is not None:
        if last_func in func_dis:
            print("Warning: Ignoring multiple functions at the same address")
        else:
            func_dis[last_func] = b"\n".join(current_output)

    return func_dis

def get_funs(f):
    funs = get_all_dis(f.name)
    return "\n".join(("%#x" % addr) for addr in funs.keys())

with gr.Blocks() as demo:

    all_dis_state = gr.State()

    gr.Markdown(
        """
    # Function/Method Detector

    First, upload a binary. Then, select a function from the dropdown. The
    function's disassembly, and the model's prediction of whether the function
    represents a C++ method or a regular function will be displayed below.
    """
    )

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

    with gr.Column(visible=False) as col:
        #output = gr.Textbox("Output")
        fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)

        with gr.Row(visible=True) as result:
            disassembly = gr.Textbox(label="Disassembly", lines=20)
            clazz = gr.Label()

    def file_change_fn(file):

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

            #fun_data = {42: 2, 43: 3}
            fun_data = get_all_dis(file.name)

            addrs = ["%#x" % addr for addr in fun_data.keys()]

            return {col: gr.update(visible=True),
                    fun_dropdown: gr.Dropdown.update(choices=addrs, value=addrs[0]),
                    all_dis_state: fun_data
                    }
        
    def function_change_fn(selected_fun, fun_data):

        disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")
        load_results = model.fn(disassembly_str)

        return {disassembly: gr.Textbox.update(value=disassembly_str),
                clazz: gr.Label.update(load_results)
        }

    file_widget.change(file_change_fn, file_widget, [col, fun_dropdown, all_dis_state])

    fun_dropdown.change(function_change_fn, [fun_dropdown, all_dis_state], [disassembly, clazz])

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