-
Notifications
You must be signed in to change notification settings - Fork 0
/
data_preparation_telock.py
55 lines (47 loc) · 1.65 KB
/
data_preparation_telock.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
import glob
import os.path
import networkx as nx
import pefile
from networkx.drawing.nx_agraph import read_dot
from utils.oep_utils import verify_offset
data_folder = "/home/hungpt/workspace/research/oep-detection/data/asm_cfg/telock"
dot_files = glob.glob(data_folder + "/*.dot")
log_oep = open("oep_data/telock.txt", "w")
def find_address(address, file):
with open(file, "r") as f:
for line in f:
if line.strip().startswith("0x"):
if address.upper() in line[:10].upper():
# print(line)
return True
return False
for dot_file in dot_files:
if "test" in dot_file:
continue
# print(asm_file)
# print("dot file: {}".format(dot_file))
original_name = os.path.basename(dot_file)[7:-10]
print(original_name)
file = os.path.join("/media/hungpt/SSD-HUNG/original_telock", original_name)
if not os.path.exists(file):
continue
pe = pefile.PE(file)
entry_point = "{:X}".format(pe.OPTIONAL_HEADER.AddressOfEntryPoint)
try:
cfg = nx.DiGraph(read_dot(path=dot_file))
# cfg = BPCFG(dot_file)
for node in cfg.nodes:
# print(node)
if not node.startswith("a"):
continue
address = node[1:11]
if verify_offset(entry_point, address):
print("OEP: {}".format(node))
log_oep.writelines("{}_{},{}\n".format("telock", original_name, node))
except Exception as e:
print(e)
log_oep.close()
# break
# print(asm_file)
# print(entry_point_address)
# print("Found: {}".format(find_address(entry_point_address, asm_file)))