VirtualCompiler / force_exec.py
elsagranger's picture
Add force execution and introduction
077a99a
# %%
from capstone import *
import json
from tqdm import tqdm
import random
from multiprocessing import Process, Queue
from unicorn.x86_const import *
from unicorn import *
from datasets import concatenate_datasets
from keystone import *
import re
from datasets import load_from_disk
def test_single(code):
ks = Ks(KS_ARCH_X86, KS_MODE_64)
try:
count = ks.asm(code)
except:
count = 0
return count
def convert_hex_format(assembly):
hex_pattern = re.compile(r'\b([0-9A-Fa-f]+)h')
converted_assembly = hex_pattern.sub(r'0x\1', assembly)
return converted_assembly
def get_name_value(name):
# match "VAR_num" label
if name.startswith("var_"):
match = re.match(r"var_(\d+)", name)
if match:
# get number
return int(match.group(1))
# else case
return None
class KeypatchAsm:
def __init__(self, arch=KS_ARCH_X86, mode=KS_MODE_64):
self.arch = arch
self.mode = mode
self.ks = Ks(self.arch, self.mode)
def fix_cmp_instruction_size(self, assembly):
lines = assembly.split('\n')
updated_lines = []
for line in lines:
if 'cmp' in line and '[' in line and ']' in line:
# add default size indicator 'dword ptr'
if ' ptr ' not in line:
line = line.replace('cmp', 'cmp dword ptr', 1)
elif 'cmp' in line and ':' in line:
line = 'nop'
updated_lines.append(line)
return '\n'.join(updated_lines)
def replace_calls_and_leas(self, assembly):
lines = assembly.split('\n')
update_lines = []
for line in lines:
if ('call' in line) and not any(x in line for x in ['0x', '0X']):
update_lines.append('nop')
elif ('lea' in line):
update_lines.append('nop')
else:
update_lines.append(line)
# print(line)
return '\n'.join(update_lines)
def remove_comments(self, assembly):
# remove ';'
lines = assembly.split('\n')
cleaned_lines = [line.split(';', 1)[0] for line in lines]
return '\n'.join(cleaned_lines).strip()
def replace_segment_register_references(self, assembly):
lines = assembly.split('\n')
updated_lines = []
for line in lines:
if 'cs:' in line:
updated_lines.append('nop')
else:
if test_single(line) == 0 and "INSTR" not in line:
updated_lines.append('nop')
else:
updated_lines.append(line)
return '\n'.join(updated_lines)
def ida_resolve(self, assembly, address):
def _resolve(_op, ignore_kw=True):
names = re.findall(r"[\$a-z0-9_:\.]+", _op, re.I)
for name in names:
# ingnore known key words
if ignore_kw and name in ('byte', 'near', 'short', 'word', 'dword', 'ptr', 'offset'):
continue
# use get_name_value fucntion
value = get_name_value(name)
if value is not None:
_op = _op.replace(name, '0x'+str(value))
return _op
# split the part and anylaize each oprand
_asm = assembly.partition(' ')
mnem = _asm[0]
opers = _asm[2].split(',')
for idx, op in enumerate(opers):
_op = list(op.partition('['))
ignore_kw = True
if _op[1] == '':
_op[2] = _op[0]
_op[0] = ''
else:
_op[0] = _resolve(_op[0], ignore_kw=True)
ignore_kw = False
_op[2] = _resolve(_op[2], ignore_kw=ignore_kw)
opers[idx] = ''.join(_op)
asm = "{0} {1}".format(mnem, ','.join(opers))
return asm
def assemble(self, assembly, address=0, syntax=KS_OPT_SYNTAX_INTEL):
assembly = assembly.replace("endbr64\n", "")
assembly = self.remove_comments(assembly)
assembly = self.ida_resolve(assembly, address)
assembly = self.replace_calls_and_leas(assembly)
assembly = self.fix_cmp_instruction_size(assembly)
assembly = self.replace_segment_register_references(assembly)
def fix_ida_syntax(assembly):
assembly = convert_hex_format(assembly)
assembly = assembly.upper()
assembly = assembly.replace("0X", " 0x")
if self.arch == KS_ARCH_X86:
if 'RETN' in assembly:
return assembly.replace('RETN', 'RET', 1)
if 'OFFSET ' in assembly:
return assembly.replace('OFFSET ', ' ')
return assembly
if syntax is None:
syntax = KS_OPT_SYNTAX_INTEL
# print(fix_ida_syntax(assembly))
try:
self.ks.syntax = syntax
encoding, count = self.ks.asm(fix_ida_syntax(assembly), address)
except KsError as e:
print(f"Error:{e}")
print(f"Assembly:\n{fix_ida_syntax(assembly)}")
print("-"*50)
print("")
encoding, count = None, 0
return (encoding, count)
UC_X86_REG_MAPPING = {
UC_X86_REG_RAX: "RAX", UC_X86_REG_RBX: "RBX", UC_X86_REG_RCX: "RCX",
UC_X86_REG_RDX: "RDX", UC_X86_REG_RSI: "RSI", UC_X86_REG_RDI: "RDI",
UC_X86_REG_RBP: "RBP", UC_X86_REG_RSP: "RSP", UC_X86_REG_R8: "R8",
UC_X86_REG_R9: "R9", UC_X86_REG_R10: "R10", UC_X86_REG_R11: "R11",
UC_X86_REG_R12: "R12", UC_X86_REG_R13: "R13", UC_X86_REG_R14: "R14",
UC_X86_REG_R15: "R15", UC_X86_REG_RIP: "RIP",
# FPU register, vector register and flag register
UC_X86_REG_XMM0: "XMM0", UC_X86_REG_XMM1: "XMM1", UC_X86_REG_XMM2: "XMM2",
UC_X86_REG_XMM3: "XMM3", UC_X86_REG_XMM4: "XMM4", UC_X86_REG_XMM5: "XMM5",
UC_X86_REG_XMM6: "XMM6", UC_X86_REG_XMM7: "XMM7", UC_X86_REG_XMM8: "XMM8",
UC_X86_REG_XMM9: "XMM9", UC_X86_REG_XMM10: "XMM10", UC_X86_REG_XMM11: "XMM11",
UC_X86_REG_XMM12: "XMM12", UC_X86_REG_XMM13: "XMM13", UC_X86_REG_XMM14: "XMM14",
UC_X86_REG_XMM15: "XMM15",
# YMM register
UC_X86_REG_YMM0: "YMM0", UC_X86_REG_YMM1: "YMM1", UC_X86_REG_YMM2: "YMM2",
UC_X86_REG_YMM3: "YMM3", UC_X86_REG_YMM4: "YMM4", UC_X86_REG_YMM5: "YMM5",
UC_X86_REG_YMM6: "YMM6", UC_X86_REG_YMM7: "YMM7", UC_X86_REG_YMM8: "YMM8",
UC_X86_REG_YMM9: "YMM9", UC_X86_REG_YMM10: "YMM10", UC_X86_REG_YMM11: "YMM11",
UC_X86_REG_YMM12: "YMM12", UC_X86_REG_YMM13: "YMM13", UC_X86_REG_YMM14: "YMM14",
UC_X86_REG_YMM15: "YMM15",
# EFLAGS register segment register
UC_X86_REG_EFLAGS: "EFLAGS",
UC_X86_REG_CS: "CS",
UC_X86_REG_DS: "DS",
UC_X86_REG_ES: "ES",
UC_X86_REG_FS: "FS",
UC_X86_REG_GS: "GS",
UC_X86_REG_SS: "SS"
}
class MemoryAccessLogger:
def __init__(self):
self.read_accesses = []
self.write_accesses = []
def hook_mem_read(self, uc, access, address, size, value, user_data):
self.read_accesses.append((address, size, value))
def hook_mem_write(self, uc, access, address, size, value, user_data):
self.write_accesses.append((address, size, value))
def hook_mem_invalid(uc, access, address, size, value, user_data):
if access == UC_MEM_WRITE_UNMAPPED or access == UC_MEM_READ_UNMAPPED or access == UC_MEM_FETCH_UNMAPPED:
print(">>> Missing memory is being WRITE at 0x%x, data size = %u, data value = 0x%x"
% (address, size, value))
start_map_addr = address & 0xfffffffffffff000
uc.mem_map(start_map_addr, start_map_addr+0x1000)
return True
return True
def instruction_hook(uc, address, size, user_data):
# get the current instruction
code = uc.mem_read(address, size)
rbp = uc.reg_read(UC_X86_REG_RBP)
rsp = uc.reg_read(UC_X86_REG_RSP)
# print(f"RBP: 0x{rbp:016x}, RSP: 0x{rsp:016x}")
# for i in md.disasm(code, address):
# print("0x%x:\t%s\t%s" %(i.address, i.mnemonic, i.op_str))
def assemble_wrapper(asm_code, code_address, result_queue):
"""
execute the function in new process and catch any exception to avoid crash the main process
"""
try:
keypatch_asm = KeypatchAsm()
encoding, count = keypatch_asm.assemble(asm_code, code_address)
result_queue.put((encoding, count))
except Exception as e:
result_queue.put((None, 0))
print("Error during assembly:", str(e))
def safe_assemble(asm_code, code_address, timeout=3):
result_queue = Queue()
p = Process(target=assemble_wrapper, args=(
asm_code, code_address, result_queue))
p.start()
p.join(timeout)
if p.is_alive():
p.terminate()
print("Terminated the process due to timeout.")
return None, 0
try:
result = result_queue.get_nowait()
return result
except Exception:
return None, 0
md = Cs(CS_ARCH_X86, CS_MODE_64)
def compile_run(asm_code, code_address, seed=0):
try:
random.seed(seed)
encoding, count = safe_assemble(asm_code, code_address)
if encoding is None or count == 0:
return "ERROR", [], []
CODE_SIZE = (count+0x1000) // 0x1000 * 0x1000
CODE_ADDRESS = code_address
STACK_ADDRESS = 0x7fff0000
STACK_SIZE = 0x2000
mu = Uc(UC_ARCH_X86, UC_MODE_64)
mu.mem_map(CODE_ADDRESS, CODE_ADDRESS+CODE_SIZE)
mu.mem_map(STACK_ADDRESS, STACK_ADDRESS+STACK_SIZE)
mu.mem_write(CODE_ADDRESS, bytes(encoding))
mu.reg_write(UC_X86_REG_RAX, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RBX, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RCX, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RDX, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RSI, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RDI, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_R8, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_R9, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_R10, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_R11, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_R12, random.randint(0, 0x2000))
mu.reg_write(UC_X86_REG_RSP, STACK_ADDRESS + STACK_SIZE)
mu.reg_write(UC_X86_REG_RBP, STACK_ADDRESS + STACK_SIZE)
mu.hook_add(UC_HOOK_MEM_INVALID |
UC_HOOK_MEM_UNMAPPED, hook_mem_invalid)
memory_logger = MemoryAccessLogger()
mu.hook_add(UC_HOOK_MEM_READ, memory_logger.hook_mem_read)
mu.hook_add(UC_HOOK_MEM_WRITE, memory_logger.hook_mem_write)
mu.emu_start(CODE_ADDRESS, CODE_ADDRESS +
len(bytes(encoding)), timeout=0, count=1000)
registers = {}
for reg_id, reg_name in UC_X86_REG_MAPPING.items():
registers[reg_name] = mu.reg_read(reg_id)
return registers, memory_logger.read_accesses, memory_logger.write_accesses
except Exception as e:
return "ERROR", [], []
# %%
ds = load_from_disk("./virtual_assembly_and_ground_truth")
# %%
all_results = {
'ground_truth': [],
'generated': []
}
test_index = []
cnt = 0
for idx, code in tqdm(enumerate(ds['asm'])):
print(idx, cnt)
regs, read_mem, write_mem = compile_run(code, 0x1000, cnt)
if regs == "ERROR":
pass
else:
test_index.append(idx)
all_results['ground_truth'].append(
{
'regs': regs,
'read_mem': read_mem,
'write_mem': write_mem
}
)
cnt += 1
for seed, index in tqdm(enumerate(test_index)):
code = ds[index]['generated_asm']
regs, read_mem, write_mem = compile_run(code, 0x1000, seed)
if regs != "ERROR":
all_results['generated'].append(
{
'regs': regs,
'read_mem': read_mem,
'write_mem': write_mem
}
)
else:
all_results['generated'].append(None)
evaluation_results = {
'regs': [],
'read_mem': [],
'write_mem': [],
}
for overall_index in tqdm(range(len(test_index))):
ground_truth = all_results['ground_truth'][overall_index]
compare = all_results['generated'][overall_index]
if compare is None:
continue
# compare regs
if len(compare['regs']) == 0:
continue
reg_name_list = [
'RAX', 'RSP', 'RBP'
]
count = 0
for reg_name in reg_name_list:
if ground_truth['regs'][reg_name] == compare['regs'][reg_name]:
count += 1
evaluation_results['regs'].append(
float(count) / len(reg_name_list))
# compare read_mem
if len(ground_truth['read_mem']) != 0:
if len(compare['read_mem']) == 0:
evaluation_results['read_mem'].append(0)
# calculate the matching score of read_mem
else:
matching_score = 0
for address, size, value in compare['read_mem']:
if (address, size, value) in ground_truth['read_mem']:
matching_score += 1
evaluation_results['read_mem'].append(
float(matching_score) / len(ground_truth['read_mem']))
# compare write_mem
if len(ground_truth['write_mem']) != 0:
if len(compare['write_mem']) == 0:
evaluation_results['write_mem'].append(0)
# calculate the matching score of write_mem
else:
matching_score = 0
for address, size, value in compare['write_mem']:
if (address, size, value) in ground_truth['write_mem']:
matching_score += 1
evaluation_results['write_mem'].append(
matching_score / len(ground_truth['write_mem']))
# %%
reg_score = sum(evaluation_results['regs']) / len(evaluation_results['regs'])
read_score = sum(evaluation_results['read_mem']) / len(
evaluation_results['read_mem'])
write_score = sum(evaluation_results['write_mem']) / len(
evaluation_results['write_mem'])
print(f"mean_score: {(reg_score + read_score + write_score) / 3}")