-
Notifications
You must be signed in to change notification settings - Fork 53
/
uefi_retool.py
executable file
·155 lines (125 loc) · 4.32 KB
/
uefi_retool.py
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
# SPDX-License-Identifier: MIT
import json
import os
import tempfile
from concurrent.futures import ProcessPoolExecutor, as_completed
import click
import uefi_firmware
from tqdm import tqdm
from tools import utils
from tools.get_efi_images import get_efi_images
CONFIG_FILE = "config.json"
DONE = click.style("DONE", fg="green")
ERROR = click.style("ERROR", fg="red")
# read configuration data
with open(CONFIG_FILE, "rb") as cfile:
CONFIG = json.load(cfile)
# ida and ida64 executables
IDA_PATH = '"{}"'.format(CONFIG["IDA_PATH"])
IDA64_PATH = '"{}"'.format(CONFIG["IDA64_PATH"])
# directories with logs
PP_GUIDS_LOGS = os.path.join(tempfile.gettempdir(), "uefi-retool-pp-guids")
ALL_INFO_LOGS = os.path.join(tempfile.gettempdir(), "uefi-retool-all-info")
def error():
print(
f"{ERROR} check your config.json file or move ida_plugin/uefi_analyser directory to IDA plugins directory"
)
exit()
def analyse_module(module, scr_name):
module_path = os.path.join(CONFIG["PE_DIR"], module)
machine_type = utils.get_machine_type(module_path)
ida_exe = IDA64_PATH
if machine_type == utils.IMAGE_FILE_MACHINE_I386:
ida_exe = IDA_PATH
analyser = os.path.join("plugins", "uefi_analyser", scr_name)
cmd = " ".join([ida_exe, f"-c -A -S{analyser}", module_path])
# analyse module in batch mode
os.system(cmd)
if not (
os.path.isfile(f"{module_path}.i64") or os.path.isfile(f"{module_path}.idb")
):
error()
return True
def analyse_all(scr_name, max_workers):
files = os.listdir(CONFIG["PE_DIR"])
# check first module
analyse_module(files[0], scr_name)
with ProcessPoolExecutor(max_workers=max_workers) as executor:
futures = [
executor.submit(analyse_module, module, scr_name) for module in files[1:]
]
params = {
"total": len(futures),
"unit": "module",
"unit_scale": True,
"leave": True,
}
for _ in tqdm(as_completed(futures), **params):
pass
def clear(dirname):
for root, dirs, files in os.walk(dirname, topdown=False):
for name in files:
os.remove(os.path.join(root, name))
for name in dirs:
os.rmdir(os.path.join(root, name))
def clear_all():
clear(CONFIG["DUMP_DIR"])
clear(CONFIG["PE_DIR"])
clear(PP_GUIDS_LOGS)
clear(ALL_INFO_LOGS)
def get_log(command, firmware_path):
suf = "all-info.json"
tmp_dir = ALL_INFO_LOGS
if command == "get-pp":
suf = "pp-guids.json"
tmp_dir = PP_GUIDS_LOGS
# collect all in one log
if not os.path.isdir(CONFIG["LOGS_DIR"]):
os.mkdir(CONFIG["LOGS_DIR"])
_, fw_name = os.path.split(firmware_path)
log_fname = os.path.join(CONFIG["LOGS_DIR"], f"{fw_name}-{suf}")
info = list()
logs = os.listdir(tmp_dir)
for log in logs:
with open(os.path.join(tmp_dir, log), "r") as f:
info.append(json.load(f))
with open(log_fname, "w") as f:
json.dump(info, f, indent=4)
print(f"{DONE} check {log_fname} file")
@click.group()
def cli():
pass
@click.command()
@click.argument("firmware_path")
@click.option("-w", "--workers", help="Number of workers (8 by default).", type=int)
def get_info(firmware_path, workers):
"""Analyze the entire UEFI firmware. The analysis result is saved to .json file."""
if not workers:
workers = 8
clear_all()
get_efi_images(firmware_path)
analyse_all("log_all.py", workers)
get_log("get-info", firmware_path)
@click.command()
@click.argument("firmware_path")
@click.option("-w", "--workers", help="Number of workers (8 by default).", type=int)
def get_pp(firmware_path, workers):
"""Get a list of proprietary protocols in the UEFI firmware. The result is saved to .json file."""
if not workers:
workers = 8
clear_all()
get_efi_images(firmware_path)
analyse_all("log_pp_guids.py", workers)
get_log("get-pp", firmware_path)
@click.command()
@click.argument("firmware_path")
def get_images(firmware_path):
"""Get executable images from UEFI firmware. Images are stored in "modules" directory."""
clear_all()
get_efi_images(firmware_path)
print(f"{DONE} check .{os.sep}modules directory")
cli.add_command(get_info)
cli.add_command(get_pp)
cli.add_command(get_images)
if __name__ == "__main__":
cli()