From 7ed8209491b334614f887f7d55f7ea41f875d8c2 Mon Sep 17 00:00:00 2001 From: yonghwi Date: Wed, 9 Jun 2021 20:09:00 -0400 Subject: [PATCH] add harness generator --- harnessgen/.gitignore | 14 + harnessgen/README.md | 166 +++ harnessgen/common.py | 792 +++++++++++ harnessgen/dominator.py | 426 ++++++ harnessgen/harconf.py | 35 + harnessgen/harnessor.py | 200 +++ harnessgen/lib/.gitignore | 385 ++++++ harnessgen/lib/Tracer/Tracer.sln | 31 + harnessgen/lib/Tracer/Tracer.vcxproj | 265 ++++ harnessgen/lib/Tracer/library_trace.cpp | 1597 +++++++++++++++++++++++ harnessgen/logger.py | 25 + harnessgen/requirements.txt | 3 + harnessgen/syn-multi.py | 332 +++++ harnessgen/synthesizer.py | 147 +++ harnessgen/template.py | 90 ++ harnessgen/util.py | 49 + harnessgen/util/ida_func_type.py | 54 + 17 files changed, 4611 insertions(+) create mode 100644 harnessgen/.gitignore create mode 100644 harnessgen/README.md create mode 100644 harnessgen/common.py create mode 100644 harnessgen/dominator.py create mode 100644 harnessgen/harconf.py create mode 100644 harnessgen/harnessor.py create mode 100644 harnessgen/lib/.gitignore create mode 100644 harnessgen/lib/Tracer/Tracer.sln create mode 100644 harnessgen/lib/Tracer/Tracer.vcxproj create mode 100644 harnessgen/lib/Tracer/library_trace.cpp create mode 100644 harnessgen/logger.py create mode 100644 harnessgen/requirements.txt create mode 100644 harnessgen/syn-multi.py create mode 100644 harnessgen/synthesizer.py create mode 100644 harnessgen/template.py create mode 100644 harnessgen/util.py create mode 100644 harnessgen/util/ida_func_type.py diff --git a/harnessgen/.gitignore b/harnessgen/.gitignore new file mode 100644 index 0000000..7886c72 --- /dev/null +++ b/harnessgen/.gitignore @@ -0,0 +1,14 @@ +*.pyc +*.swp +*.bak +temptrace* +domitrace* +.vscode + +# exclude pin binaries +lib/META-INF +lib/pin +lib/pin*zip + +# exclude cache directory +cache/ \ No newline at end of file diff --git a/harnessgen/README.md b/harnessgen/README.md new file mode 100644 index 0000000..5128a70 --- /dev/null +++ b/harnessgen/README.md @@ -0,0 +1,166 @@ +# Harness Generator + +## Library installation (tested with VS2017 and VS2019) + +- Install required python packages + +```sh +python3 -m pip install -r requirements.txt +``` + +- Install pintools + +```sh +cd {repo}/lib +wget https://software.intel.com/sites/landingpage/pintool/downloads/pin-3.13-98189-g60a6ef199-msvc-windows.zip +unzip pin-3.13-98189-g60a6ef199-msvc-windows.zip +mv pin-3.13-98189-g60a6ef199-msvc-windows pin +rm -f pin-3.13-98189-g60a6ef199-msvc-windows.zip +``` + +- Copy the tracer source code to pintools + +```sh +cd {repo}/lib +cp -r Tracer pin/source/tools +``` + +- Compile the tracer (with pintools) + - Open Tracer.sln with visual studio + - Open library_trace.cpp and modify `_WINDOWS_H_PATH_` to a proper path (line 2) + - If you are using VS2017, you should modify platform toolset (under Property - General - Platform Toolset) + - Build the solution + - Test + +## Collect Dynamic Run Traces & Harness generation + +Then we run a program with PIN-based tracer to +- infer (LCA analysis) + +### One trace + +```sh +# Trace API calls +pin.exe -t source/tools/Tracer/Release/Tracer.dll ^ + -logdir "cor1_1" -trace_mode "all" ^ + -only_to_target "test.exe" -only_to_lib "test.dll" ^ + -- test.exe input1 + +# Run harness synthesizer on single trace which starts from START_FUNCTION(...) +python3 synthesizer.py harness -t cor1_1/drltrace.PID.log -d cor1_1/memdump -s START_FUNCTION +``` + +#### Result example + +```c++ +#include +... +typedef int (__stdcall *IDP_Init_func_t)(int); +typedef int (__stdcall *IDP_GetPlugInInfo_func_t)(int); +... + +void fuzz_me(char* filename){ + + IDP_Init_func_t IDP_Init_func; + IDP_GetPlugInInfo_func_t IDP_GetPlugInInfo_func; +... + + /* Harness function #0 */ + int* c0_a0 = (int*) calloc (4096, sizeof(int)); + LOAD_FUNC(dlllib, IDP_Init); + int IDP_Init_ret = IDP_Init_func(&c0_a0); + dbg_printf("IDP_Init, ret = %d\n", IDP_Init_ret); + + /* Harness function #1 */ + int* c1_a0 = (int*) calloc (4096, sizeof(int)); + LOAD_FUNC(dlllib, IDP_GetPlugInInfo); + int IDP_GetPlugInInfo_ret = IDP_GetPlugInInfo_func(&c1_a0); + dbg_printf("IDP_GetPlugInInfo, ret = %d\n", IDP_GetPlugInInfo_ret); + +... + /* Harness function #66 */ + int* c66_a0 = (int*) calloc (4096, sizeof(int)); + LOAD_FUNC(dlllib, IDP_CloseImage); + int IDP_CloseImage_ret = IDP_CloseImage_func(&c66_a0); + dbg_printf("IDP_CloseImage, ret = %d\n", IDP_CloseImage_ret); + +} + + +int main(int argc, char ** argv) +{ + if (argc < 2) { + printf("Usage %s: \n", argv[0]); + printf(" e.g., harness.exe input\n"); + exit(1); + } + + dlllib = LoadLibraryA("%s"); + if (dlllib == NULL){ + dbg_printf("failed to load library, gle = %d\n", GetLastError()); + exit(1); + } + + char * filename = argv[1]; + fuzz_me(filename); + return 0; +} +``` + +### One correct + one incorrect + +- Run the same tracer with the same input, but with `-logdir "cor1_2"`. +- Then run the same tracer with a different input, but with `-logdir "cor2_1"`. + +```sh +# Trace API calls +pin.exe -t source/tools/Tracer/Release/Tracer.dll ^ + -logdir "cor1_2" -trace_mode "all" ^ + -only_to_target "test.exe" -only_to_lib "test.dll" ^ + -- test.exe input1 + +# Trace API calls +pin.exe -t source/tools/Tracer/Release/Tracer.dll ^ + -logdir "cor2_1" -trace_mode "all" ^ + -only_to_target "test.exe" -only_to_lib "test.dll" ^ + -- test.exe input2 + +# Run harness synthesizer on single trace which starts from START_FUNCTION(...) +python3 syn-multi.py harness -t ./ -s START_FUNCTION +``` + +### LCA Analysis + +```sh +# Trace API calls +# - from test.exe -> test.dll (API) +# - file-related APIs (CreateFile, ...) +mkdir dom +pin.exe -t source/tools/Tracer/Release/Tracer.dll ^ + -logdir "dom" -trace_mode "dominator" ^ + -only_to_target "test.exe" -only_to_lib "test.dll" ^ + -- test.exe + +# Do LCA analysis with API call traces between CreateFileW ~ CloseHandle +python3 dominator.py -s CreateFileW -e CloseHandle -sample "" -t ./dom/drltrace.PID.log -d ./dom/memdump/ +``` + +#### Result example + +``` +[*] Displaying Most Frequent Address (Dominator candidates) + >> Total unique harness functions: 400 + >> Total number of function address identified: 223 + >> Total number of candidate address(es): 4 + +[*] Dominator analysis + >> Bad candidate (called multiple times): 0xc7c925 + >> Good candidate (called only once): 0xc579fb, 0xc5857f, 0xc56820 + >> Candidate address (sorted by the distance from harness): 0xc579fb, 0xc5857f, 0xc56820 +``` + +### Troubleshooting + +- The three synthesizers (synthesizer.py, syn-multi.py, dominator.py) utilizes static analyzer (e.g., IDA Pro). + If it fails to find IDA Pro, try adjusting `IDA_PATH` in harconf.py. + diff --git a/harnessgen/common.py b/harnessgen/common.py new file mode 100644 index 0000000..dbb486d --- /dev/null +++ b/harnessgen/common.py @@ -0,0 +1,792 @@ +import subprocess +import typing +import re +import json +import hashlib +import os +import bisect +from harconf import * +from template import * +from util import strings, u32 + +### PARSER: COMMON + + +def get_baseaddr(chunk: bytes, modulename: bytes): + lines = chunk.split(b"\n") + for line in lines: + if modulename in line: + return int(line.split(b',')[2], 16) + raise Exception("No modulename in the entry?") + + +def ret_start_point(pn: str, keyword: bytes): + """ + 1) return cid and tid from this example line + CALLID[3] TID[3756] IJ T2M 0x63621040->0x65cf6450(avformat-gp-57.dll!avformat_open_input+0x0) + 2) for now, this function is case sensitive + """ + with open(pn, 'rb') as f: + lines = f.readlines() + for line in lines: + if keyword in line and b"0x0" in line: + cid = int(line.split(b"CALLID[")[1].split(b"]")[0]) + tid = int(line.split(b"TID[")[1].split(b"]")[0]) + return cid, tid + + raise Exception("Cannot find the starting function from the trace file") + + +### PARSER: FUNCTION + +class Args: + def __init__(self, name: bytes, addr: int, ret_type: str, args: typing.List[str], convention: str): + """ + e.g., ret : void + e.g., args: ["int", "*int", int] + """ + self.name = name + self.addr = addr + self.ret_type = ret_type + self.args = args + self.convention = convention # std or cdecl + + def arg(self, index): + if index < len(self.args): + return self.args[index].replace("*", "") + raise Exception("Index out of bound") + + def argtype(self, index): + if index < len(self.args): + arg = self.args[index] + if "*" in arg: + return "data" + else: + return "pointer" + raise Exception("Index out of bound") + + @property + def argsize(self): + return len(self.args) + + @property + def rettype(self): + return self.ret_type + + +class Functype: + def __init__(self, functype_pn): + self.functypes_by_name = {} + self.functypes_by_addr = {} + self.functype_pn = functype_pn + self.sorted_addr = [] + self.parse() + + def parse(self): + with open(self.functype_pn, 'rb') as f: + lines = f.readlines() + for line in lines: + if b"|" not in line: + continue + + payload = [x.strip() for x in line.strip().split(b"|", 2)] + if payload[2] == "None": + continue + + # e.g., 0x000|aix_read_header|int __cdecl(int)|{"args": [...], "cc": "__stdcall", "ret_type": "..."} + addr = int(payload[0], 16) + funcname = payload[1] + funcinfo = json.loads(payload[2]) + ret_type = funcinfo['ret_type'] + convention = funcinfo['cc'] + args = [arg['type'] for arg in funcinfo['args']] + + self.functypes_by_addr[addr] = \ + self.functypes_by_name[funcname] = \ + Args(funcname, addr, ret_type, args, convention) + self.sorted_addr.append(addr) + + self.sorted_addr = sorted(self.sorted_addr) + + def by_name(self, funcname) -> Args: + return self.functypes_by_name[funcname] + + def by_addr(self, funcaddr) -> Args: + fi = self.functypes_by_addr.get(funcaddr) + if not fi: + return Args(b'', funcaddr, 'int', ['int'] * 9, '__cdecl') + else: + return fi + + def by_addr_near(self, funcaddr) -> Args: + index = bisect.bisect_right(self.sorted_addr, funcaddr) - 1 + assert index != -1, "Function not found at %s!" % hex(funcaddr) + addr = self.sorted_addr[index] + return self.functypes_by_addr[addr] + +### PARSER: TRACE + + +class TraceElement: + # An argument looks like this: [(0x14334333, "DP"), (0x000000, "D")] + args: typing.List[typing.List[typing.Tuple[int, str]]] + + # get dump from bin files + args_dump: typing.List[typing.Tuple[bytes, bytes]] + + def __init__(self, tracetype: str, src_addr: int, dst_addr: int, src_module: bytes, dst_module: bytes): + self.tracetype = tracetype # either "call" or "ret" + self.src_addr = src_addr + self.dst_addr = dst_addr # this is also return_address + self.args = [] + self.args_dump = [] + self.ret_val = None + + # Optional symbols + self.has_symbols = bool(dst_module or src_module) + self.src_module = (src_module or '').lower() + self.dst_module = (dst_module or '').lower() + self.src_symbol = None + self.dst_symbol = None + + def update_calltrace(self, dst_sym, args, args_dump, pointer): + self.dst_symbol = dst_sym + self.args = args + self.args_dump = args_dump + self.pointer = pointer + + def update_rettrace(self, ret_val): + self.ret_val = ret_val + + def __repr__(self): + return "TraceElement(%r, %s, %s)" % ((self.tracetype, hex(self.src_addr), hex(self.dst_addr))) + + +class FunctypeManager: + cache = {} + + def __init__(self): + pass + + @staticmethod + def get(path) -> Functype: + res = FunctypeManager.cache.get(path) + if res: + return res + + cache_out = FunctypeManager.dest_filename(path) + cache_idb = cache_out + '.idb' + command = [IDA_PATH, '-A', '-S' + IDA_SCRIPT] + command += ['-o' + cache_idb, path] if not os.path.isfile(cache_idb) else [cache_idb] + + if not os.path.isfile(cache_out): + print(f"Generating function types for {path.decode()} ...") + subprocess.run(command, + env={'TVHEADLESS': '1', 'DESTPATH': cache_out, **os.environ}) + else: + print(f"Found cached function types for {path.decode()} !") + + if not os.path.isfile(cache_out): + raise Exception(f"Generating function type for {path} failed!\nCommand: {command}") + + res = FunctypeManager.cache[path] = Functype(cache_out) + return res + + @staticmethod + def dest_filename(path) -> str: + return os.path.join(FUNCTYPE_CACHE_PATH, hashlib.sha256(path.lower()).hexdigest()) + + +class Trace: + # from print_address; address, [module, symbol, offset]. []: optional + CALL_ENTRY = r'({address})(?:\(({module})!({symbol})\+({address})\))?'\ + .format(address=r'0x[0-9a-fA-F]+', module=r'[^!]*', symbol=r'[^+]*')\ + .encode() + + cid_sequence: typing.List[int] + + # TODO: parse all traces, now we are tracing specified threadID with starting point + def __init__(self, trace_pn: str, dumpdir: str, tid=None, start_cid=None, build=True): + self.trace_pn = trace_pn + self.functype_manager = FunctypeManager() + self.tid = tid + self.start_cid = start_cid + self.cid_sequence = [] + self.dumpdir = dumpdir + + if build: + self.build() + + #self.module_baseaddr = None + #self.caller_baseasddr = None + + """ + print hex(self.calltrace[0].src_addr) + print hex(self.calltrace[0].dst_addr) + print self.calltrace[0].dst_symbol + print self.calltrace[0].args + """ + + def build(self): + """ + Analyze trace with function_type information (limit the number of arguments) + """ + calltrace: typing.Dict[int, TraceElement] + rettrace: typing.Dict[int, TraceElement] + + calltrace = {} + rettrace = {} + + self.modules = modules = {} + + with open(self.trace_pn, 'rb') as f: + entries = f.read().split(b"==\n") + entries, module_trace = entries[:-1], entries[-1] + + for line in module_trace.strip().split(b'\n')[1:]: + line = map(lambda x: x.strip(), line.split(b',', 9)) + idx1, idx2, base, end, entry, a, b, c, path = line + modules[path.lower()] = int(base, 16), int(end, 16) + + for x in range(len(entries) - 1): + chunk = entries[x + 1] + tid, cid, tracetype = self.get_tid(chunk) + + if self.tid is not None and tid != self.tid: + continue + + if tracetype == "CALL": + calltrace[cid] = self.parse_call(chunk, tid, cid) + self.cid_sequence.append(cid) + + elif tracetype == "RET": + rettrace[cid] = self.parse_ret(chunk) + + # store (for just in case) + first_call = calltrace[self.cid_sequence[0]] + self.module_baseaddr = modules[first_call.dst_module][0] + self.caller_baseaddr = modules[first_call.src_module][0] + + # remove if None + calltrace = {k: v for k, v in calltrace.items() if v is not None} + rettrace = {k: v for k, v in rettrace.items() if v is not None} + + self.calltrace, self.rettrace = calltrace, rettrace + + def parse_arg(self, line: bytes, index, tid, cid): + """ example + -A9: 0x002eeec8[DP] > 0x002eef08[DP] > 0x002eef1c[DP] > 0x002eef3c + """ + out: typing.List[typing.Tuple[int, str]] + chain: typing.List[bytes] + + out = [] + chain = line.split(b":")[1].split(b">") + + for value in chain: + if b"[" in value: + pointer_type = value.split(b"[")[1].split(b"]")[0].decode() + actual_value = int(value.split(b"[")[0], 16) + else: + pointer_type = "D" + actual_value = int(value.strip().split(b" ")[0], 16) + out.append((actual_value, pointer_type)) + + return out + + def find_module(self, address: int) -> typing.Tuple[typing.Union[bytes, None], int, int]: + for path, (base, end) in self.modules.items(): + if base <= address <= end: + return path, base, end + + return None, 0, 0 + + def find_function(self, address: int) -> typing.Union[Args, None]: + mod, mod_base, _ = self.find_module(address) + if not mod: + return None + + ft = self.functype_manager.get(mod) + fi = ft.by_addr(address - mod_base) + return fi + + def parse_call(self, chunk: bytes, tid, cid, parse_args=True) -> TraceElement: + """ example + CALLID[0] TID[3756] IJ T2M 0x63621000->0x65cad480(avformat-gp-57.dll!avformat_get_riff_audio_tags+0x0) + -A0: 0x636350a5[DP] > 0xa382c4a3 + ... + -A9: 0x002eeec8[DP] > 0x002eef08[DP] > 0x002eef1c[DP] > 0x002eef3c + + CALLID[0] TID[3664] IC T2M @0x008f10de(math3.exe!fuzz_me+0xde)->0x735a1100(MathLibrary.dll!test+0x0) + """ + lines = chunk.split(b"\n") + arg_lines = list(filter(lambda x: x.startswith(b' -A'), lines[1:])) + + args = [] + args_dump = [] + pointer = [] + + src, dst = re.findall(Trace.CALL_ENTRY, lines[0]) + src_addr = int(src[0], 16) + dst_addr = int(dst[0], 16) + src_module = src[1] + dst_module = dst[1] + + dst_sym = None + numargs = len(arg_lines) + + mod, _, _ = self.find_module(src_addr) + if mod: + src_module = mod + + mod, _, _ = self.find_module(dst_addr) + if mod: + dst_module = mod + + fi = self.find_function(dst_addr) + if fi: + dst_sym = fi.name + numargs = fi.argsize + + # XXX: Tracer decides the maximum argument size + numargs = min(numargs, len(arg_lines)) + + if parse_args: + # parse from A1 to A9 (a0 is ret addr (stack)) + for x in range(numargs): + current_arg = self.parse_arg(arg_lines[x], x, tid, cid) + args.append(current_arg) + + # FIXME: what if we meet code pointer? should we dump code? + # FIXME: we should consider multi-level, multi-element + if current_arg[0][1] == "DP": + dumpread, same, dump_pointer = self.read_dump(x, tid, cid) + args_dump.append(dumpread) + pointer.append(dump_pointer) + else: + args_dump.append(None) + pointer.append(None) + + te = TraceElement("call", src_addr, dst_addr, src_module, dst_module) + te.update_calltrace(dst_sym, args, args_dump, pointer) + + return te + + def read_dump(self, index, tid, cid): + # filename e.g., t788-c99-a8.post + + pre_pn = os.path.join( + self.dumpdir, "t%d-c%d-a%d.pre" % (tid, cid, index)) + post_pn = os.path.join( + self.dumpdir, "t%d-c%d-a%d.post" % (tid, cid, index)) + + with open(pre_pn, 'rb') as f: + pre = f.read(BINREAD) + f.seek(BINREAD) + pre_pointer = f.read(1000) + + if not os.path.isfile(post_pn): # unmap? + post_pn = pre_pn + + with open(post_pn, 'rb') as f: + post = f.read(BINREAD) + f.seek(BINREAD) + post_pointer = f.read(1000) + + return (pre, post), pre == post, (pre_pointer, post_pointer) + + def parse_ret(self, chunk: bytes) -> TraceElement: + """ example + RETID[2] TID[3756] RET2T 0x65be6913(avformat-gp-57.dll!av_register_all+0xcb3)->0x63641000(MediaSource.ax!libssh2_session_abstract+0x4b50) + RETVAL: 0x00000001 + """ + lines = chunk.split(b"\n") + assert len(lines) >= 2 + + keyword = lines[0].split(b' TID')[1].split(b' ')[1] + assert keyword in (b'RET2M', b'RET2T', b'RETFR'), 'Unknown keyword: %r' % keyword + + src, dst = re.findall(Trace.CALL_ENTRY, lines[0]) + src_addr = int(src[0], 16) + dst_addr = int(dst[0], 16) + + ret_val = int(chunk.split(b"RETVAL:")[1].split(b'\n')[0], 16) + + te = TraceElement("ret", src_addr, dst_addr, src[1], dst[1]) + te.update_rettrace(ret_val) + + return te + + def get_tid(self, chunk: bytes): + tracetype = "" + cid = -1 + tid = -1 + if b" DC " in chunk or b" IC " in chunk or b" IJ " in chunk or b" FR " in chunk: + tracetype = "CALL" + cid = int(chunk.split(b"[")[1].split(b"]")[0]) + elif b"RET" in chunk: + tracetype = "RET" + cid = int(chunk.split(b"[")[1].split(b"]")[0]) + if b" TID[" in chunk: + tid = int(chunk.split(b" TID[")[1].split(b"]")[0]) + + assert(tracetype != "") + + # no tid information from the chunk + return tid, cid, tracetype + + +class SimpleTrace(Trace): + def __init__(self, trace_pn): + self.trace_pn = trace_pn + self.cid_sequence = [] + self.unique_call = {} + self.calltrace = self.build() + + def build(self): + calltrace = {} + + with open(self.trace_pn, 'rb') as f: + fdata = f.read().split(b"==\n") + for x in range(len(fdata) - 1): + chunk = fdata[x + 1] + tid, cid, tracetype = self.get_tid(chunk) + + if tracetype == "CALL": + calltrace[cid] = self.parse_call(chunk, tid, cid, parse_args=False) + self.cid_sequence.append(cid) + + # remove if None + calltrace = {k: v for k, v in calltrace.items() if v is not None} + return calltrace + + +# This is the important part; program synthesizer +class Synthesizer: + def __init__(self, trace_pn: str, dump_pn: str, functype_pn: str, start_func=None, sample_name: str = None): + self.dump_pn = dump_pn + self.trace_pn = trace_pn + self.start_func = start_func + self.functype_pn = functype_pn + self.sample_name = sample_name + + if self.start_func != None: + self.start_cid, self.trace_tid = ret_start_point(self.trace_pn, self.start_func.encode()) + else: + self.start_cid, self.trace_tid = None, None + + self.functype_manager = FunctypeManager() + self.trace = Trace(self.trace_pn, self.dump_pn, self.trace_tid, self.start_cid) + self.defined_types, self.defined_funcs = self.typedef() + + self.defined_variables = [] + self.defined_pointer = {} # {address:variable_name} + self.body = [] + self.defined_func = [] + self.history = {} + + def emit_code(self, out_pn=None): + header = HEADER.replace("{typedef}", '\n'.join(self.defined_types)) + fuzzme = FUZZME.replace("{funcdef}", '\n'.join(self.defined_funcs)) + fuzzme = fuzzme.replace("{harness}", '\n'.join(self.body)) + + print(header) + print(fuzzme) + print(MAIN) + + def typedef(self): + """ + 1) for each trace, we identify unique function + 2) we prepare type information (ready to emit) + 3) emit to global typedef + e.g., typedef int (__stdcall *avformat_get_riff_audio_tags_func_t)(); + 4) emit to fuzzme() + e.g., avformat_get_riff_audio_tags_func_t avformat_get_riff_audio_tags_func; + """ + defined_types = [] + defined_funcs = [] + + for cid in self.trace.cid_sequence: + te = self.trace.calltrace[cid] + funcname = te.dst_symbol + mod, mod_base, _ = self.trace.find_module(te.dst_addr) + assert mod + funcinfo = self.functype_manager.get(mod).by_addr(te.dst_addr - mod_base) + args = funcinfo.args + args_str = ', '.join(args) if len(args) > 0 else '' + convention = funcinfo.convention + ret_type = funcinfo.ret_type + + _types = self.ret_typedef_func( + funcname, args_str, convention, ret_type) + _funcs = self.ret_defined_func(funcname) + if _types not in defined_types: + defined_types.append(_types) + if _funcs not in defined_funcs: + defined_funcs.append(_funcs) + + """ + for x in range(len(defined_types)): + print defined_types[x] + + for x in range(len(defined_types)): + print defined_funcs[x] + """ + + return defined_types, defined_funcs + + def analyze(self): + # 1) infer argument for filepath or input data + + # 2) diff pre and post function + + # 3) discover used pointer + pass + + def dig_userinput(self, cid, args, args_dump, args_ptr): + + input_pn = os.path.join(os.path.dirname(self.functype_pn), INPUT1) + + if os.path.exists(input_pn): + userinput = open(input_pn, 'rb').read() + else: + return None + + # for x in range(len(args)): + # if args[x][0][1] == 'DP': + # print userinput in args_dump[x][0] + + def build_body(self, *args): + raise NotImplementedError() + + def search_pointer(self, addr, passed_cid=10000, internal_use=False): + """ + utility function to pinpoint the location of address in memory dump + """ + + result = [] + + # for each callid + for cid in self.trace.cid_sequence: + + if cid > passed_cid: + continue + + variables = [] # defined variables: e.g., int a=0 + arguments = [] # used arguments: func(&a) + calltrace = self.trace.calltrace[cid] + + args = calltrace.args + args_dump = calltrace.args_dump + args_ptr = calltrace.pointer + + for x in range(len(args)): + current_arg = args[x][0][0] + if current_arg == addr: + if not internal_use: + print("[*] Passed argument at [cid:%d] [arg:%dth]" % (cid, x)) + else: + result.append(("arg", cid, x)) + + for x in range(len(args)): + if args[x][0][1] == 'DP': + postdump = args_dump[x][1] + for y in range(0, len(postdump), 4): + if addr == u32(postdump[y:y + 4]): + if not internal_use: + print("[*] Memory dump at [cid:%d] [arg:%dth] [idx:%d]" % (cid, x, y)) + else: + result.append(("dump", cid, x, y)) + + if not internal_use: + return None + else: + return result + + def ret_pointer_at_dump(self, cid, arg, idx): + # e.g., *((int*)c3_a0[3] + # FIXME: assuming integer type + return "*((int*)c%d_a%d[%d])" % (cid, arg, idx) + + def ret_addr_of_var(self, orig_val): + return "&(%s)" % orig_val + + def check_searched_result(self, _result, query): + # query: either "arg" or "dump" + # return oldest one + # print _result, "|", query + if _result is None: + return None + + for result in _result: + keyword = result[0] + if keyword == query: + return result + + return None + + def ret_arg_code(self, cid, args, args_dump, args_type, args_ptr): + """ + - cid: call id + - args: actual argument values + - args_dump: followed result from pointer array[0]=pre, array[1]=post + - args_type: inferred type for each argument + """ + need_to_define = [] + arguments = [] + pointer_defined_flag = False + + # 1) will use raw value (basically) + # 2) if pointer, we define variable and pass the address + # 3) if pointer indicates 0, we allocate heap with 1000 size + for x in range(len(args)): + pointer_defined_flag = False + # data pointer + if args[x][0][1] == 'DP': + # TODO: consider data-type when unpack + + # 1) infer filename argument (if the string contains filename information) + first_string = next(strings(args_dump[x][0])) + if self.sample_name.encode() in first_string: + arguments.append("filename") + continue + else: + dumped = hex(u32(args_dump[x][0])) + _type = args_type[x].replace("*", "") + + # 1-1) infer chuck of actual sample is used in the function + # TODO + + # 2) we allocate heap if pointed value is 0 + if dumped == '0x0': + # we always allocate enough space for pointer to zero (could be initialization) + need_to_define.append("%s* c%d_a%d = (%s*) calloc (%d, sizeof(%s));" % + (_type, cid, x, _type, BINREAD, _type)) + arguments.append("&c%d_a%d" % (cid, x)) + + # 3) Check pre-defined pointer + # If there is, we reuse the pointer + else: + # print args[x][0][0] + + # Is the address is already referenced from the previous pointer? + # print cid + # print args[x][0][0] + # print self.defined_pointer.keys() + if args[x][0][0] not in self.defined_pointer: + # print hex(args[x][0][0]) + result = self.search_pointer(args[x][0][0], cid, internal_use=True) + result_arg = self.check_searched_result(result, "arg") + result_dump = self.check_searched_result(result, "dump") + + # print result_arg + # print result_dump + + # 3-1) searches for the address from the previous operation + # if there exist address in the dump (e.g., assigned after function call), + # we try to use that (only if the dump exist previously) + if result is not None and result_dump is not None: + _cid = result_dump[1] + _arg = result_dump[2] + _idx = result_dump[3] + ptrname = self.ret_pointer_at_dump(_cid, _arg, _idx) + self.defined_pointer[args[x][0][0]] = ptrname + need_to_define.append('') + arguments.append(ptrname) + continue + + # 3-2) what if the address is used by another arguments? + elif False: + # elif result is not None and result_arg is not None: + # print result + self.defined_pointer[args[x][0] + [0]] = "&c%d_a%d" % (cid, x) + need_to_define.append( + "%s c%d_a%d = %s;" % (_type, cid, x, dumped)) + + # 3-3) if not, we define new one + else: + self.defined_pointer[args[x][0] + [0]] = "&c%d_a%d" % (cid, x) + need_to_define.append( + "%s c%d_a%d = %s;" % (_type, cid, x, dumped)) + pointer_defined_flag = True + + # if it is pre-defined, we do nothing + else: + need_to_define.append('') + + # If we don't have choice, we define new pointer + arguments.append(self.defined_pointer[args[x][0][0]]) + + # 4) Check whether referenced value (from pointer) is defined as another pointer + # e.g., arg1|A --> 0x1000, arg1|B --> A -> 0x1000 + # ==> B = &A (not just raw value of A) + if pointer_defined_flag == True: + # now, we are selecting the referenced value (this is also address) + __result = self.search_pointer( + args[x][1][0], cid, internal_use=True) + # print "pointer", __result + __result_arg = self.check_searched_result(__result, "arg") + # print __result_arg + + if __result_arg is not None: + + result_cid = __result_arg[1] + result_arg = __result_arg[2] + + # history = {cid: (need_to_define, arguments)} + # print "DEBUG", result_cid, result_arg, self.history, self.history[result_cid][1] + previous_argument = self.history[result_cid][1][result_arg] + addr_previous_argument = self.ret_addr_of_var( + previous_argument) + + # rollback + del self.defined_pointer[args[x][0][0]] + arguments = arguments[:-1] + need_to_define = need_to_define[:-1] + + # append arguments + self.defined_pointer[args[x][1] + [0]] = addr_previous_argument + arguments.append(addr_previous_argument) + need_to_define.append('') + + elif args[x][0][1] == 'CP': + """ failed trial + code_pointer = args[x][0][0] + self.defined_pointer[args[x][0][0]] = "&c%d_a%d" % (cid, x) + need_to_define.append("%s c%d_a%d = %s;" % (_type, cid, x, code_pointer)) + arguments.append(self.defined_pointer[args[x][0][0]]) + """ + + # print self.trace.caller_baseasddr + + raw_value = args[x][0][0] + _type = args_type[x].replace("*", "") + + append_str = " /* Possible code pointer offset: %s */" % hex( + int(raw_value) - self.trace.caller_baseaddr) + # print append_str + + # we provide the information about the code pointer + need_to_define.append("") + arguments.append(hex(raw_value) + append_str) + + # raw data + elif args[x][0][1] == 'D': + # TODO: consider data-type when unpack + raw_value = args[x][0][0] + _type = args_type[x].replace("*", "") + + need_to_define.append("") + arguments.append(hex(raw_value)) + + return need_to_define, arguments + + def ret_typedef_func(self, funcname, args_str, convention, ret_type): + # e.g., typedef int (__stdcall *avformat_get_riff_audio_tags_func_t)(); + return "typedef %s (%s *%s_func_t)(%s);" % (ret_type, convention, funcname.decode(), args_str) + + def ret_defined_func(self, funcname): + return (b" %s_func_t %s_func;" % (funcname, funcname)).decode() diff --git a/harnessgen/dominator.py b/harnessgen/dominator.py new file mode 100644 index 0000000..496982f --- /dev/null +++ b/harnessgen/dominator.py @@ -0,0 +1,426 @@ +# For section IV-B: call-sequence recovery +# From specified targets (leaf nodes; e.g. CreateFile, ReadFile), +# This tries to find a dominator node (or LCA) that will call all of these targets + +#!/usr/bin/env python2 +import re +import typing +from common import Trace, TraceElement +import os +import signal +import argparse +import networkx as nx +import matplotlib.pyplot as plt + +from tqdm import tqdm +from logger import * +from harconf import * +from template import * +from util import exit_gracefully + +""" +./dominator.py -t domitrace_alzipcon/trace.log -d domitrace_alzipcon +./dominator.py -t domitrace_alzipcon/trace.log -d domitrace_alzipcon --start Format_QueryCollectionInfo +./dominator.py -t domitrace_alzipcon/trace.log -d domitrace_alzipcon --start Format_QueryCollectionInfo --end Format_Release --sample-name test.egg + +alzip: answer is 0x4343d0 (4408272) +""" + +logger = getlogger("Dominator") +DUMPDIR = "" +MERGE_BOUNDARY = True + + +class Dominator(object): + def __init__(self, trace_pn, dump_pn, start_func=None, end_func=None, sample_name=None): + self.dump_pn = dump_pn + self.trace_pn = trace_pn + self.start_func = start_func.encode() + self.end_func = end_func.encode() + self.sample_name = sample_name.encode() + + """ + self.trace_tid = -1 + self.trace_tid2 = -1 + + if self.start_func != None: + self.start_cid, self.trace_tid = self.ret_start_point(self.trace_pn, self.start_func) + if self.end_func != None: + self.end_cid, self.trace_tid2 = self.ret_end_point(self.trace_pn, self.end_func) + assert(self.trace_tid == self.trace_tid2) + """ + + self.start_cid, self.end_cid, self.interesting_tid, _ = self.ret_interesting_locations() + self.trace = DominatorTrace(self.trace_pn, self.start_cid, + self.end_cid, self.interesting_tid) + + self.defined_variables = [] + self.defined_pointer = {} # {address:variable_name} + self.body = [] + self.history = {} + self.har_addr = {} + + """ ALL functions which used in the harness """ + for addr in list(self.trace.all_callers.keys()): + # print hex(addr), self.trace.all_callers[addr] + self.har_addr[addr] = self.trace.all_callers[addr] + + self.dominator() + # print self.trace.callgraph.edges + # print self.trace.func_boundary + + def dominator(self): + + # Calculate dominator using DFS and common address (Lowest common ancestor) + unique_code_size = len(self.trace.node_list) + all_nodes = self.trace.callgraph.nodes + + print("[*] Processing Directed-graph to find dominator") + storage = [] + for harness_addr in tqdm(list(self.trace.all_callers)): + harness_addr_start = self.trace.get_func_start(harness_addr, merge_boundary=MERGE_BOUNDARY) + out = [] + + for i in range(unique_code_size - 1, 0, -1): + # print(harness_addr, self.trace.node_list[i]) + if self.trace.node_list[i] == harness_addr_start or self.trace.node_list[i] == 0: + continue + if self.trace.node_list[i] not in all_nodes: + continue + + # print harness_addr_start, self.trace.node_list[i] + for path in nx.all_simple_paths(self.trace.callgraph, source=self.trace.node_list[i], target=harness_addr_start, cutoff=50): + out = out + path + + out = list(dict.fromkeys(out)) + # print out + storage = storage + out + + # Print out the report + print("[*] Displaying Most Frequent Address (Dominator candidates)") + func_counter: typing.Dict[int, int] = {} + for func_addr in storage: + func_counter[func_addr] = func_counter.get(func_addr, 0) + 1 + + popular_words = sorted(func_counter, key=func_counter.get, reverse=True) + most_func_count = func_counter[popular_words[0]] + + report_addr = [] + for addr in func_counter.keys(): + if func_counter[addr] == most_func_count: + report_addr.append(addr) + + print(" >> Total unique harness functions: %d" % (len(self.trace.all_callers))) + print(" >> Total number of function address identified: %d" % most_func_count) + print(" >> Total number of candidate address(es): %d" % len(report_addr)) + # print(" >> Total number of candidate address(es): %d" % ', '.join(report_addr)) + + # for debug + # report_addr = ['0x421000', '0x430200', '0x4333c0', '0x424410', '0x41a590', '0x424e20', '0x44d0b0', '0x4477d0', '0x421440', '0x424230', '0x4335c0', '0x433de0', '0x424a10', '0x430e40', '0x452a60', '0x4243f0', '0x44b8c0', '0x447850', '0x4331a0', '0x40fdd0', '0x436860', '0x430610', '0x41ac80', '0x433340', '0x452e70', '0x4339c0', '0x430880', '0x4242c0', '0x4343d0', '0x44ccf0', '0x40fc90', '0x44d200', '0x44baa0', '0x42eea0', '0x41a570', '0x4346b0', '0x4160c0', '0x432920', '0x4330d0', '0x4260d1', '0x423d90', '0x420ac0', '0x424ac0', '0x4308e0', '0x4179d0', '0x423f90', '0x4483a0', '0x424cf0', '0x423360', '0x42e720', '0x457500', '0x447480', '0x417f80', '0x422000', '0x422920', '0x430710', '0x41a610', '0x44d2c0', '0x44d920', '0x433530', '0x422c20', '0x424890', '0x420f30', '0x430680', '0x421ee0', '0x452b10', '0x41ed40', '0x448ae0', '0x421370', '0x424330', '0x447120', '0x430260', '0x44d6b0', '0x430d60', '0x422c90', '0x424c80', '0x422d00', '0x42eb70', '0x422e00', '0x408e80', '0x4326e0', '0x40e180', '0x447940', '0x430f80', '0x422b90', '0x4329e0', '0x44bf90', '0x446fa0', '0x424af0', '0x406ce0', '0x423910', '0x44d9b0', '0x433b10', '0x434330', '0x4231c0', '0x4211a0', '0x433800', '0x44c020', '0x41f0e0', '0x44dbd0', '0x4129c0', '0x44be10', '0x420a40', '0x44e2a0', '0x430de0', '0x418050', '0x434600', '0x415df0', '0x424e40'] + + # Heuristics + """ + 1. display report address with CID (CALLID) + 2. display observed number of that address in the trace + (less number is desirable, should be 1?) + 3. display upper address of that function + """ + + print("\n[*] Dominator analysis") + candidate = {} + candidate["good"] = [] + candidate["bad"] = [] + for addr in report_addr: + count = self.ret_addr_count_trace(addr) # how many times called? + + if count == 1: + candidate["good"].append(addr) + else: + candidate["bad"].append(addr) + + final_report = {} + for good_addr in candidate["good"]: + final_report[good_addr] = self.distance_from_startcid(good_addr) + + final_report = sorted(final_report, key=final_report.get, reverse=False) + + print(" >> Bad candidate (called multiple times): %s" % ', '.join(hex(addr) for addr in candidate["bad"])) + print(" >> Good candidate (called only once): %s" % ', '.join(hex(addr) for addr in candidate["good"])) + print(" >> Candidate address (sorted by the distance from harness): %s" % ', '.join(hex(addr) for addr in final_report)) + + """ how to find this address? + [*] Dominator analysis + {'0x40fc90': 3618, '0x44b8c0': 3493, '0x41a610': 3737, '0x447940': 3242, '0x4335c0': 3615, '0x42eea0': 3504, '0x447120': 3616, '0x446fa0': 3606, '0x452e70': 3601, '0x40e180': 3746, '0x432920': 3604, '0x4346b0': 3168} + >> Candidate address (sorted by the distance from harness): + ['0x4346b0', '0x447940', '0x44b8c0', '0x42eea0', '0x452e70', '0x432920', '0x446fa0', '0x4335c0', '0x447120', '0x40fc90', '0x41a610', '0x40e180'] + + (answer) + 4346b0 <- 41a610 <- 40e180 + 41a610 <- 40fc90 + """ + + def distance_from_startcid(self, good_addr): + current_cid = 0 + for cid in list(self.trace.calltrace.keys()): + if good_addr == self.trace.calltrace[cid].dst_addr: + current_cid = cid + return self.start_cid - current_cid + + def ret_addr_count_trace(self, dst_addr): + counter = 0 + for cid in list(self.trace.calltrace.keys()): + if dst_addr == self.trace.calltrace[cid].dst_addr: + if self.start_cid > cid: + counter += 1 + else: + counter += 2 + return counter + + def ret_interesting_locations(self): + start_cid = None + start_tid = None + end_cid = 9e999 + end_tid = -1 + + node_out = [] + + with open(self.trace_pn, 'rb') as f: + lines = f.readlines() + for line in lines: + + if b" DC " in line or b" IC " in line or b" IJ " in line or b" FR " in line: + src_addr = int(line.split(b"->")[0].split(b"(")[0].split(b"0x")[1], 16) + # dst_addr = int(line.split(b"->")[1].split(b"(")[0], 16) + if src_addr not in node_out: + node_out.append(src_addr) + + if self.start_func in line and b"0x0" in line: + cid = int(line.split(b"CALLID[")[1].split(b"]")[0]) + tid = int(line.split(b"TID[")[1].split(b"]")[0]) + + if start_cid == None: + start_cid = cid + start_tid = tid + + if self.end_func in line and b"0x0" in line: + cid = int(line.split(b"CALLID[")[1].split(b"]")[0]) + tid = int(line.split(b"TID[")[1].split(b"]")[0]) + + end_cid = cid + end_tid = tid + + return start_cid, end_cid, start_tid, node_out + + +class DominatorTrace(Trace): + # TODO: parse all traces, now we are tracing specified threadID with starting point + def __init__(self, trace_pn, start_cid, end_cid, interesting_tid): + super().__init__(trace_pn, DUMPDIR, interesting_tid, start_cid, build=False) + self.func_boundary = {} + self.src_to_dst = {} + self.dst_from_src = {} + self.start_cid = start_cid + self.end_cid = end_cid + self.all_callers = {} + self.node_list = [] + self.callgraph = nx.DiGraph() + + self.possible_return = {} + self.build() + + # test + # print(self.calltrace[0]) + # print(self.rettrace[1]) + + self.func_boundary = self.sanitize_func_boundary() # dict{start:dst} + self.generate_digraph() # store to self.callgraph + # self.show_graph() + + # print self.func_boundary + + """ TEST function boundary (start ==> dst) + for key in self.func_boundary.keys(): + print(hex(key), hex(self.func_boundary[key])) # src->dst + """ + + """ TEST whether same source has multiple targets (i.e., indirect call) + for key in self.src_to_dst.keys(): + print(hex(key), self.src_to_dst[key]) + """ + + """ TEST Xref functions (i.e., who are calling this function?) + for key in self.dst_from_src.keys(): + print(hex(key), self.dst_from_src[key]) + """ + + def generate_digraph(self): + for cid in list(self.calltrace.keys()): + te = self.calltrace[cid] + + start = te.src_addr + end = te.dst_addr + func_start = self.get_func_start(start, merge_boundary=MERGE_BOUNDARY) + # if func_start == 0 and src_addr == 4303892: + # print "here", src_addr, dst_addr + if func_start not in self.node_list: + self.node_list.append(func_start) + # print src_addr, func_start + self.callgraph.add_edge(func_start, end) + # print(f"Call {hex(start)}: {hex(func_start)} -> {hex(end)}") + + def show_graph(self): + nx.draw(self.callgraph, node_size=200) + plt.show() + + def get_func_start(self, addr, merge_boundary=False): + # for start_addr in self.func_boundary.keys(): + # if start_addr <= addr <= self.func_boundary[start_addr]: + # return start_addr + + # # one more try to find and change the function boundary + # """ + # Before> + # A ---- B C-----D + # E + # After> + # A ---------E C-----D + + # option1: just return A + # option1: return A and modify the function boundary + # """ + # if merge_boundary: + # # find cloest boundary + # diff = 0x100000 + # candidate = 0 + # for start_addr in self.func_boundary.keys(): + # end_addr = self.func_boundary[start_addr] + # if end_addr < addr < end_addr + diff: + # candidate = start_addr + # diff = addr - end_addr + + # if candidate != 0: + # return candidate + + # Addition: use function type database + mod, mod_base, _ = self.find_module(addr) + if mod: + ft = self.functype_manager.get(mod) + fi = ft.by_addr_near(addr - mod_base) + return fi.addr + mod_base + + # finally we give up + return 0 + + def sanitize_func_boundary(self): + out = {} + + for key in list(self.func_boundary.keys()): + start_addr = key + end_addr = self.func_boundary[key] + + # print start_addr, end_addr + if end_addr == -1: + continue + + if end_addr - start_addr > 0x10000: + continue + + out[start_addr] = end_addr + return out + + def insert_relationship(self, src, dst): + + if src not in list(self.src_to_dst.keys()): + self.src_to_dst[src] = [] + self.src_to_dst[src].append(dst) + else: + if dst not in self.src_to_dst[src]: + self.src_to_dst[src].append(dst) + + if dst not in list(self.dst_from_src.keys()): + self.dst_from_src[dst] = [] + self.dst_from_src[dst].append(src) + else: + if src not in self.dst_from_src[dst]: + self.dst_from_src[dst].append(src) + + def windows_target(self, te: TraceElement): + dlls = [b"kernelbase.dll", b"kernel32.dll", b"ntdll.dll"] + + if te.dst_module and te.dst_module not in dlls: + return False + + return True + + def parse_call(self, chunk: bytes, tid, cid): + te = super().parse_call(chunk, tid, cid) + src_addr = te.src_addr + dst_addr = te.dst_addr + + # for finding fucntion boundary + self.possible_return[src_addr] = dst_addr + + if dst_addr not in self.func_boundary: + self.func_boundary[dst_addr] = -1 + + # for making src_to_dst and dst_from_src relationship + self.insert_relationship(src_addr, dst_addr) + + if cid >= self.start_cid and cid <= self.end_cid: + if te.has_symbols and src_addr not in self.all_callers and not self.windows_target(te): + self.all_callers[src_addr] = te.src_symbol + + return te + + def parse_ret(self, chunk: bytes): + te = super().parse_ret(chunk) + + # update for building function boundary + self.update_ret_addr(te.dst_addr, te.src_addr) + + return te + + # dst_addr is returning address (returining target) + def update_ret_addr(self, dst_addr, current_addr): + for addr in self.possible_return.keys(): + # found return place + if addr <= dst_addr <= addr + 7: + start_addr = self.possible_return[addr] + self.func_boundary[start_addr] = max( + self.func_boundary[start_addr], current_addr) + + +def main(): + global DUMPDIR + # DEFINE PARSER + parser = argparse.ArgumentParser() + parser.add_argument("-t", "--trace", dest="trace_file", type=str, + default=None, help="Trace file collected from DynamoRIO", + required=True) + parser.add_argument("-d", "--memory-dump", dest="dump_dir", type=str, + default=None, help="memory dump file directory (pre/post)", + required=True) + parser.add_argument("-s", "--start", dest="start_func", type=str, + default=None, help="name of the starting function to process", + required=False) + parser.add_argument("-e", "--end", dest="end_func", type=str, + default=None, help="name of the ending function to process", + required=False) + parser.add_argument("-sample", "--sample-name", dest="sample_name", type=str, + default=None, help="name of the original sample name", + required=False) + args = parser.parse_args() + # END PARSER + + # Ctrl-c handler + signal.signal(signal.SIGINT, exit_gracefully( + signal.getsignal(signal.SIGINT))) + + # Start harness dominator + DUMPDIR = args.dump_dir + domi = Dominator(args.trace_file, args.dump_dir, + args.start_func, args.end_func, args.sample_name) + + +if __name__ == '__main__': + main() diff --git a/harnessgen/harconf.py b/harnessgen/harconf.py new file mode 100644 index 0000000..b0bb056 --- /dev/null +++ b/harnessgen/harconf.py @@ -0,0 +1,35 @@ +import glob +import os + +# for functype manager; it invokes an IDAPython script +# TODO: edit this if it doesn't work +ROOT = os.path.dirname(__file__) +IDA_PATH = glob.glob(r'C:\Program Files\IDA *\idat.exe')[0] +IDA_SCRIPT = os.path.join(ROOT, 'util/ida_func_type.py') +FUNCTYPE_CACHE_PATH = os.path.join(ROOT, 'cache') + +if not os.path.isdir(FUNCTYPE_CACHE_PATH): + os.mkdir(FUNCTYPE_CACHE_PATH) + +if not os.path.exists(IDA_SCRIPT): + print("Check if the following file exists:", IDA_SCRIPT) + exit(1) + +# for harnesor +TRACE_PN = "trace" +FUNCTYPE = "functype" +NORMAL_POSTFIX = "normal" +TRACE_PREFIX = ["DC", "RET", "IC", "IJ"] + +THRESHOLD = 0x20 +TRACE_MAX = 20 +INPUT = "QQQQ" + +NAMEDIC = {} +NAMEDIC['CP'] = "Code Pointer" +NAMEDIC['DP'] = "Data Pointer" + +# for synthesizer +BINREAD = 0x1000 +POINTER_SEARCH_LIMIT = 100 +INPUT_COMP_LENGTH = 10 \ No newline at end of file diff --git a/harnessgen/harnessor.py b/harnessgen/harnessor.py new file mode 100644 index 0000000..49bcddd --- /dev/null +++ b/harnessgen/harnessor.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +import os +import glob + +from harconf import * + + +class Tracer(object): + def __init__(self, project): + self.project = project + self.dir = os.path.join(TRACE_PN, self.project) + self.filelist = glob.glob(self.dir+"/*.log") + self.normalfile = self.find_normal_file() # store normal (no interesting points) + self.filelist.remove(self.normalfile) + + self.trace_normal = None + self.trace_interesting = {} + self.trace_interesting_line = {} + self.trace_interesting_line2index = {} + self.list_diff = {} + self.funcline_diff = {} + self.modulename = None + + def find_normal_file(self): + for filename in self.filelist: + if "_"+NORMAL_POSTFIX in filename: + return filename + raise Exception("No normal file which contains _normal") + + @property + def files(self): + return self.filelist + + def collect_trace(self): + self.unique_functions_normal, self.modulename, self.unique_functions_line, _ \ + = parse_trace_unique_callee(self.normalfile) + + for filename in self.filelist: + #print (filename) + basename = os.path.basename(filename) + self.trace_interesting[basename], _, self.trace_interesting_line[basename], \ + self.trace_interesting_line2index[basename] = parse_trace_unique_callee(filename) + + def extract_unique_callsite(self): + for filename in self.filelist: + #print (list_diff(self.trace_interesting[os.path.basename(filename)], self.unique_functions_normal)) + basename = os.path.basename(filename) + self.list_diff[basename] = list_diff(self.trace_interesting[basename], self.unique_functions_normal) + self.funcline_diff[basename] = list_diff(self.trace_interesting_line[basename], self.unique_functions_line) + + def print_unique_trace(self): + # print-out the first line of unique call (indirect/direct/ind-jmp) + + for key in self.list_diff.keys(): + print('\n' + key) + for x in range(len(self.funcline_diff[key])): + print(self.funcline_diff[key][x], self.trace_interesting_line2index[key][self.funcline_diff[key][x]]) + + def extract_interesting_trace(self, outdir): + + for key in self.list_diff.keys(): + #print ('\n' + key) + idx_list = [] + for x in range(len(self.funcline_diff[key])): + idx = self.trace_interesting_line2index[key][self.funcline_diff[key][x]] + idx_list.append(idx) + + min_val, max_val = extract_minmax(idx_list) + out_pn = os.path.join(outdir, key) + dump_extracted_trace(min_val, max_val, out_pn, self.project) + + +def extract_minmax(idx_list): + min_val = min(idx_list) + max_val = max(idx_list) + + while True: + if max_val - min_val > TRACE_MAX: + idx_list.remove(max_val) + max_val = max(idx_list) + else: + break + + return min_val, max_val + + +def list_diff(li1, li2): + return (list(set(li1) - set(li2))) + + +def get_module_name(chunk): + lines = chunk.split("\n") + if "T2M" in chunk: + return lines[1].split(".dll")[0].split(" ")[-1]+".dll" + + elif "M2T" in chunk: + return lines[0].split(".dll")[0].split(" ")[-1]+".dll" + + +def get_baseaddr(chunk, modulename): + lines = chunk.split("\n") + for line in lines: + if modulename in line: + return int(line.split(',')[2], 16) + raise Exception("No modulename in the entry?") + + +def sanitize_fcall_line(line, baseaddr): + addr = line.split("0x")[1].split(" ")[0] + newaddr = hex(int(addr, 16) - baseaddr)[2:] + + if " @ " in line: + line = line.split(" @ ")[1].strip() + if "=> " in line: + line = line.split("=> ")[1].strip() + + line = line.replace(" ? ??:0", "") + line = line.replace(" ??:0", "") + line = line.replace(" to ", "") + line = line.replace(addr, newaddr) + + return line.strip() + + +def extract_call_addr(chunk, baseaddr): + lines = chunk.split("\n") + + if lines[0].split(" ")[0].strip() in TRACE_PREFIX: + if lines[0].split(" ")[1].strip() == "T2M": + return int(lines[1].split("0x")[1].split(" ")[0], 16) - baseaddr, sanitize_fcall_line(lines[1], baseaddr) + elif lines[0].split(" ")[1].strip() == "M2T": + return int(lines[0].split("0x")[1].split(" ")[0], 16) - baseaddr, sanitize_fcall_line(lines[0], baseaddr) + + #print (chunk) + raise Exception("No keyword? T2M or M2T?") + else: + return None, None + + +def dump_extracted_trace(min_val, max_val, out_pn, project): + trace_pn = os.path.join(TRACE_PN, project, os.path.basename(out_pn)) + #print (trace_pn, filename) + fdata = None + out = "" + + with open(trace_pn, 'r') as f: + fdata = f.read().split("==\n") + + modulename = get_module_name(fdata[1]) + baseaddr = get_baseaddr(fdata[-1], modulename) + + for x in range(min_val, max_val+1): + chunk = fdata[x+1] + out += chunk + "==\n" + out += modulename + "|" + hex(baseaddr) + + with open(out_pn, 'w') as f: + f.write(out) + + +def parse_trace_unique_callee(filename): + unique_targets = [] + unique_funcline = [] + line2index = {} + fdata = None + + with open(filename, 'r') as f: + fdata = f.read().split("==\n") + + modulename = get_module_name(fdata[1]) + baseaddr = get_baseaddr(fdata[-1], modulename) + + for x in range(len(fdata)-1): + chunk = fdata[x+1] + fcall_addr, relavant_line = extract_call_addr(chunk, baseaddr) + if fcall_addr is not None: + line2index[relavant_line.strip()] = x + if fcall_addr not in unique_targets: + unique_targets.append(fcall_addr) + unique_funcline.append(relavant_line) + + """ + if x not in line2index.keys(): + line2index[relavant_line.strip()] = x + """ + + return unique_targets, modulename, unique_funcline, line2index + + +def main(): + tr = Tracer("notepad++") + #tr = Tracer("aimp") + tr.collect_trace() + tr.extract_unique_callsite() + # tr.print_unique_trace() + tr.extract_interesting_trace("trace/extracted") + + +if __name__ == '__main__': + main() diff --git a/harnessgen/lib/.gitignore b/harnessgen/lib/.gitignore new file mode 100644 index 0000000..e7b7851 --- /dev/null +++ b/harnessgen/lib/.gitignore @@ -0,0 +1,385 @@ +pin +pin2 + +ndk + +# Created by https://www.gitignore.io/api/c++,visualstudio +# Edit at https://www.gitignore.io/?templates=c++,visualstudio + +### C++ ### +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +### VisualStudio ### +## Ignore Visual Studio temporary files, build results, and +## files generated by popular Visual Studio add-ons. +## +## Get latest from https://github.com/github/gitignore/blob/master/VisualStudio.gitignore + +# User-specific files +*.rsuser +*.suo +*.user +*.userosscache +*.sln.docstates + +# User-specific files (MonoDevelop/Xamarin Studio) +*.userprefs + +# Build results +[Dd]ebug/ +[Dd]ebugPublic/ +[Rr]elease/ +[Rr]eleases/ +x64/ +x86/ +[Aa][Rr][Mm]/ +[Aa][Rr][Mm]64/ +bld/ +[Bb]in/ +[Oo]bj/ +[Ll]og/ + +# Visual Studio 2015/2017 cache/options directory +.vs/ +# Uncomment if you have tasks that create the project's static files in wwwroot +#wwwroot/ + +# Visual Studio 2017 auto generated files +Generated\ Files/ + +# MSTest test Results +[Tt]est[Rr]esult*/ +[Bb]uild[Ll]og.* + +# NUNIT +*.VisualState.xml +TestResult.xml + +# Build Results of an ATL Project +[Dd]ebugPS/ +[Rr]eleasePS/ +dlldata.c + +# Benchmark Results +BenchmarkDotNet.Artifacts/ + +# .NET Core +project.lock.json +project.fragment.lock.json +artifacts/ + +# StyleCop +StyleCopReport.xml + +# Files built by Visual Studio +*_i.c +*_p.c +*_h.h +*.ilk +*.meta +*.iobj +*.pdb +*.ipdb +*.pgc +*.pgd +*.rsp +*.sbr +*.tlb +*.tli +*.tlh +*.tmp +*.tmp_proj +*_wpftmp.csproj +*.log +*.vspscc +*.vssscc +.builds +*.pidb +*.svclog +*.scc + +# Chutzpah Test files +_Chutzpah* + +# Visual C++ cache files +ipch/ +*.aps +*.ncb +*.opendb +*.opensdf +*.sdf +*.cachefile +*.VC.db +*.VC.VC.opendb + +# Visual Studio profiler +*.psess +*.vsp +*.vspx +*.sap + +# Visual Studio Trace Files +*.e2e + +# TFS 2012 Local Workspace +$tf/ + +# Guidance Automation Toolkit +*.gpState + +# ReSharper is a .NET coding add-in +_ReSharper*/ +*.[Rr]e[Ss]harper +*.DotSettings.user + +# JustCode is a .NET coding add-in +.JustCode + +# TeamCity is a build add-in +_TeamCity* + +# DotCover is a Code Coverage Tool +*.dotCover + +# AxoCover is a Code Coverage Tool +.axoCover/* +!.axoCover/settings.json + +# Visual Studio code coverage results +*.coverage +*.coveragexml + +# NCrunch +_NCrunch_* +.*crunch*.local.xml +nCrunchTemp_* + +# MightyMoose +*.mm.* +AutoTest.Net/ + +# Web workbench (sass) +.sass-cache/ + +# Installshield output folder +[Ee]xpress/ + +# DocProject is a documentation generator add-in +DocProject/buildhelp/ +DocProject/Help/*.HxT +DocProject/Help/*.HxC +DocProject/Help/*.hhc +DocProject/Help/*.hhk +DocProject/Help/*.hhp +DocProject/Help/Html2 +DocProject/Help/html + +# Click-Once directory +publish/ + +# Publish Web Output +*.[Pp]ublish.xml +*.azurePubxml +# Note: Comment the next line if you want to checkin your web deploy settings, +# but database connection strings (with potential passwords) will be unencrypted +*.pubxml +*.publishproj + +# Microsoft Azure Web App publish settings. Comment the next line if you want to +# checkin your Azure Web App publish settings, but sensitive information contained +# in these scripts will be unencrypted +PublishScripts/ + +# NuGet Packages +*.nupkg +# The packages folder can be ignored because of Package Restore +**/[Pp]ackages/* +# except build/, which is used as an MSBuild target. +!**/[Pp]ackages/build/ +# Uncomment if necessary however generally it will be regenerated when needed +#!**/[Pp]ackages/repositories.config +# NuGet v3's project.json files produces more ignorable files +*.nuget.props +*.nuget.targets + +# Microsoft Azure Build Output +csx/ +*.build.csdef + +# Microsoft Azure Emulator +ecf/ +rcf/ + +# Windows Store app package directories and files +AppPackages/ +BundleArtifacts/ +Package.StoreAssociation.xml +_pkginfo.txt +*.appx + +# Visual Studio cache files +# files ending in .cache can be ignored +*.[Cc]ache +# but keep track of directories ending in .cache +!?*.[Cc]ache/ + +# Others +ClientBin/ +~$* +*~ +*.dbmdl +*.dbproj.schemaview +*.jfm +*.pfx +*.publishsettings +orleans.codegen.cs + +# Including strong name files can present a security risk +# (https://github.com/github/gitignore/pull/2483#issue-259490424) +#*.snk + +# Since there are multiple workflows, uncomment next line to ignore bower_components +# (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) +#bower_components/ +# ASP.NET Core default setup: bower directory is configured as wwwroot/lib/ and bower restore is true +**/wwwroot/lib/ + +# RIA/Silverlight projects +Generated_Code/ + +# Backup & report files from converting an old project file +# to a newer Visual Studio version. Backup files are not needed, +# because we have git ;-) +_UpgradeReport_Files/ +Backup*/ +UpgradeLog*.XML +UpgradeLog*.htm +ServiceFabricBackup/ +*.rptproj.bak + +# SQL Server files +*.mdf +*.ldf +*.ndf + +# Business Intelligence projects +*.rdl.data +*.bim.layout +*.bim_*.settings +*.rptproj.rsuser + +# Microsoft Fakes +FakesAssemblies/ + +# GhostDoc plugin setting file +*.GhostDoc.xml + +# Node.js Tools for Visual Studio +.ntvs_analysis.dat +node_modules/ + +# Visual Studio 6 build log +*.plg + +# Visual Studio 6 workspace options file +*.opt + +# Visual Studio 6 auto-generated workspace file (contains which files were open etc.) +*.vbw + +# Visual Studio LightSwitch build output +**/*.HTMLClient/GeneratedArtifacts +**/*.DesktopClient/GeneratedArtifacts +**/*.DesktopClient/ModelManifest.xml +**/*.Server/GeneratedArtifacts +**/*.Server/ModelManifest.xml +_Pvt_Extensions + +# Paket dependency manager +.paket/paket.exe +paket-files/ + +# FAKE - F# Make +.fake/ + +# JetBrains Rider +.idea/ +*.sln.iml + +# CodeRush personal settings +.cr/personal + +# Python Tools for Visual Studio (PTVS) +__pycache__/ +*.pyc + +# Cake - Uncomment if you are using it +# tools/** +# !tools/packages.config + +# Tabs Studio +*.tss + +# Telerik's JustMock configuration file +*.jmconfig + +# BizTalk build output +*.btp.cs +*.btm.cs +*.odx.cs +*.xsd.cs + +# OpenCover UI analysis results +OpenCover/ + +# Azure Stream Analytics local run output +ASALocalRun/ + +# MSBuild Binary and Structured Log +*.binlog + +# NVidia Nsight GPU debugger configuration file +*.nvuser + +# MFractors (Xamarin productivity tool) working folder +.mfractor/ + +# Local History for Visual Studio +.localhistory/ + +# BeatPulse healthcheck temp database +healthchecksdb + +# End of https://www.gitignore.io/api/c++,visualstudio + diff --git a/harnessgen/lib/Tracer/Tracer.sln b/harnessgen/lib/Tracer/Tracer.sln new file mode 100644 index 0000000..cb2d0c9 --- /dev/null +++ b/harnessgen/lib/Tracer/Tracer.sln @@ -0,0 +1,31 @@ + +Microsoft Visual Studio Solution File, Format Version 12.00 +# Visual Studio Version 16 +VisualStudioVersion = 16.0.29306.81 +MinimumVisualStudioVersion = 10.0.40219.1 +Project("{8BC9CEB8-8B4A-11D0-8D11-00A0C91BC942}") = "Tracer", "Tracer.vcxproj", "{639EF517-FCFC-408E-9500-71F0DC0458DB}" +EndProject +Global + GlobalSection(SolutionConfigurationPlatforms) = preSolution + Debug|x64 = Debug|x64 + Debug|x86 = Debug|x86 + Release|x64 = Release|x64 + Release|x86 = Release|x86 + EndGlobalSection + GlobalSection(ProjectConfigurationPlatforms) = postSolution + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Debug|x64.ActiveCfg = Debug|x64 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Debug|x64.Build.0 = Debug|x64 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Debug|x86.ActiveCfg = Debug|Win32 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Debug|x86.Build.0 = Debug|Win32 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Release|x64.ActiveCfg = Release|x64 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Release|x64.Build.0 = Release|x64 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Release|x86.ActiveCfg = Release|Win32 + {639EF517-FCFC-408E-9500-71F0DC0458DB}.Release|x86.Build.0 = Release|Win32 + EndGlobalSection + GlobalSection(SolutionProperties) = preSolution + HideSolutionNode = FALSE + EndGlobalSection + GlobalSection(ExtensibilityGlobals) = postSolution + SolutionGuid = {07CB9EAD-D3D1-4D13-96E3-68BBB515F7B7} + EndGlobalSection +EndGlobal diff --git a/harnessgen/lib/Tracer/Tracer.vcxproj b/harnessgen/lib/Tracer/Tracer.vcxproj new file mode 100644 index 0000000..30d019b --- /dev/null +++ b/harnessgen/lib/Tracer/Tracer.vcxproj @@ -0,0 +1,265 @@ + + + + + Debug + Win32 + + + Debug + x64 + + + Release + Win32 + + + Release + x64 + + + + {639EF517-FCFC-408E-9500-71F0DC0458DB} + MyPinTool + Win32Proj + + + + DynamicLibrary + MultiByte + true + v142 + + + DynamicLibrary + MultiByte + v142 + + + DynamicLibrary + MultiByte + true + v142 + + + DynamicLibrary + MultiByte + v142 + + + + + + + + + + + + + + + + + + + <_ProjectFileVersion>10.0.40219.1 + $(ProjectDir)$(Configuration)\ + $(Configuration)\ + false + false + $(ProjectDir)$(Platform)\$(Configuration)\ + $(Platform)\$(Configuration)\ + false + false + $(ProjectDir)$(Configuration)\ + $(Configuration)\ + false + false + $(ProjectDir)$(Platform)\$(Configuration)\ + $(Platform)\$(Configuration)\ + false + false + AllRules.ruleset + + + AllRules.ruleset + + + AllRules.ruleset + + + AllRules.ruleset + + + + + + /GR- /GS- /EHs- /EHa- /FP:strict /Oi- /FIinclude/msvc_compat.h %(AdditionalOptions) + Disabled + ..\..\include\pin;..\..\include\pin\gen;..\InstLib;..\..\..\extras\xed-ia32\include\xed;..\..\..\extras\components\include;..\..\..\extras\stlport\include;..\..\..\extras;..\..\..\extras\libstdc++\include;..\..\..\extras\crt\include;..\..\..\extras\crt;..\..\..\extras\crt\include\arch-x86;..\..\..\extras\crt\include\kernel\uapi;..\..\..\extras\crt\include\kernel\uapi\asm-x86;%(AdditionalIncludeDirectories) + TARGET_IA32;HOST_IA32;TARGET_WINDOWS;WIN32;__PIN__=1;PIN_CRT=1;__i386__ + false + + + Default + MultiThreaded + false + true + NotSet + false + + + Level3 + ProgramDatabase + 4530;%(DisableSpecificWarnings) + + + /export:main /ignore:4210 %(AdditionalOptions) + pin.lib;xed.lib;pinvm.lib;kernel32.lib;pincrt.lib;ntdll-32.lib;crtbeginS.obj + ..\..\..\ia32\lib;..\..\..\ia32\lib-ext;..\..\..\extras\xed-ia32\lib;..\..\..\ia32\runtime\pincrt;%(AdditionalLibraryDirectories) + true + %(IgnoreSpecificDefaultLibraries) + true + NotSet + false + Ptrace_DllMainCRTStartup%4012 + 0x55000000 + MachineX86 + true + false + + + + + X64 + + + /GR- /GS- /EHs- /EHa- /FP:strict /Oi- /FIinclude/msvc_compat.h %(AdditionalOptions) + Disabled + ..\..\include\pin;..\..\include\pin\gen;..\InstLib;..\..\..\extras\xed-intel64\include\xed;..\..\..\extras\components\include;..\..\..\extras\stlport\include;..\..\..\extras;..\..\..\extras\libstdc++\include;..\..\..\extras\crt\include;..\..\..\extras\crt;..\..\..\extras\crt\include\arch-x86_64;..\..\..\extras\crt\include\kernel\uapi;..\..\..\extras\crt\include\kernel\uapi\asm-x86;%(AdditionalIncludeDirectories) + TARGET_IA32E;HOST_IA32E;TARGET_WINDOWS;WIN32;__PIN__=1;PIN_CRT=1;__LP64__ + false + + + Default + MultiThreaded + false + true + false + + + Level3 + ProgramDatabase + 4530;%(DisableSpecificWarnings) + + + /export:main %(AdditionalOptions) + pin.lib;xed.lib;pinvm.lib;kernel32.lib;pincrt.lib;ntdll-64.lib;crtbeginS.obj + ..\..\..\intel64\lib;..\..\..\intel64\lib-ext;..\..\..\extras\xed-intel64\lib;..\..\..\intel64\runtime\pincrt;%(AdditionalLibraryDirectories) + true + %(IgnoreSpecificDefaultLibraries) + true + NotSet + false + Ptrace_DllMainCRTStartup + 0xC5000000 + MachineX64 + true + + + + + /GR- /GS- /EHs- /EHa- /FP:strict /Oi- /FIinclude/msvc_compat.h %(AdditionalOptions) + false + false + ..\..\include\pin;..\..\include\pin\gen;..\InstLib;..\..\..\extras\xed-ia32\include\xed;..\..\..\extras\components\include;..\..\..\extras\stlport\include;..\..\..\extras;..\..\..\extras\libstdc++\include;..\..\..\extras\crt\include;..\..\..\extras\crt;..\..\..\extras\crt\include\arch-x86;..\..\..\extras\crt\include\kernel\uapi;..\..\..\extras\crt\include\kernel\uapi\asm-x86;%(AdditionalIncludeDirectories) + TARGET_IA32;HOST_IA32;TARGET_WINDOWS;WIN32;__PIN__=1;PIN_CRT=1;__i386__ + false + + + Default + MultiThreaded + false + true + NotSet + false + + + Level3 + + + 4530;%(DisableSpecificWarnings) + + + /export:main %(AdditionalOptions) + pin.lib;xed.lib;pinvm.lib;kernel32.lib;pincrt.lib;ntdll-32.lib;crtbeginS.obj + ..\..\..\ia32\lib;..\..\..\ia32\lib-ext;..\..\..\extras\xed-ia32\lib;..\..\..\ia32\runtime\pincrt;%(AdditionalLibraryDirectories) + true + %(IgnoreSpecificDefaultLibraries) + true + NotSet + true + + + + + Ptrace_DllMainCRTStartup%4012 + 0x55000000 + MachineX86 + false + + + + + X64 + + + /GR- /GS- /EHs- /EHa- /FP:strict /Oi- /FIinclude/msvc_compat.h %(AdditionalOptions) + false + false + ..\..\include\pin;..\..\include\pin\gen;..\InstLib;..\..\..\extras\xed-intel64\include\xed;..\..\..\extras\components\include;..\..\..\extras\stlport\include;..\..\..\extras;..\..\..\extras\libstdc++\include;..\..\..\extras\crt\include;..\..\..\extras\crt;..\..\..\extras\crt\include\arch-x86_64;..\..\..\extras\crt\include\kernel\uapi;..\..\..\extras\crt\include\kernel\uapi\asm-x86;%(AdditionalIncludeDirectories) + TARGET_IA32E;HOST_IA32E;TARGET_WINDOWS;WIN32;__PIN__=1;PIN_CRT=1;__LP64__ + false + + + Default + MultiThreaded + false + true + false + + + Level3 + + + 4530;%(DisableSpecificWarnings) + + + /export:main %(AdditionalOptions) + pin.lib;xed.lib;pinvm.lib;kernel32.lib;pincrt.lib;ntdll-64.lib;crtbeginS.obj + ..\..\..\intel64\lib;..\..\..\intel64\lib-ext;..\..\..\extras\xed-intel64\lib;..\..\..\intel64\runtime\pincrt;%(AdditionalLibraryDirectories) + true + %(IgnoreSpecificDefaultLibraries) + true + NotSet + true + + + + + Ptrace_DllMainCRTStartup + 0xC5000000 + MachineX64 + + + + + + + + + + + + \ No newline at end of file diff --git a/harnessgen/lib/Tracer/library_trace.cpp b/harnessgen/lib/Tracer/library_trace.cpp new file mode 100644 index 0000000..3e6aa0e --- /dev/null +++ b/harnessgen/lib/Tracer/library_trace.cpp @@ -0,0 +1,1597 @@ +#include "pin.H" +#define _WINDOWS_H_PATH_ C:/Program Files (x86)/Windows Kits/10/Include/10.0.18362.0/um +namespace W { +#include +}; +W::HANDLE current_process; + +enum { + UNIQUE, RELATION, ALL, INDIRECT, DOMINATOR +} trace_mode; + +typedef enum { + DRLTRC_NONE_POINTER, + DRLTRC_CODE_POINTER, + DRLTRC_DATA_POINTER +} drltrc_pointer_type_t; + +#define BUFFER_SIZE_BYTES(buf) sizeof(buf) +#define BUFFER_SIZE_ELEMENTS(buf) (BUFFER_SIZE_BYTES(buf) / sizeof((buf)[0])) +#define BUFFER_LAST_ELEMENT(buf) (buf)[BUFFER_SIZE_ELEMENTS(buf) - 1] + +int tracemode; + +#define VNOTIFY(level, msg, ...) + +#define MAX_PTR_DEPTH 4 + +/* Frontend scope is defined here because if logdir is a forbidden path we have to change + * it and provide for our client manually. + */ +KNOB op_logdir +(KNOB_MODE_WRITEONCE, "pintool", "logdir", ".", "Log directory to print library call data\n" + "Specify log directory where library call data will be written, in a separate file per " + "process. The default value is \".\" (current dir). If set to \"-\", data for all " + "processes are printed to stderr (warning: this can be slow)."); + +std::string logdir; + +KNOB op_functype +(KNOB_MODE_WRITEONCE, "pintool", "functype", "none", "functype information\n" "Specify functype file"); + +KNOB op_only_from_app +(KNOB_MODE_WRITEONCE, "pintool", "only_from_app", "0", "Reports only library calls from the app\n" + "Only reports library calls from the application itself, as opposed to all calls even " + "from other libraries or within the same library."); + +KNOB op_follow_children +(KNOB_MODE_WRITEONCE, "pintool", "follow_children", "1", "Trace child processes\n" + "Trace child processes created by a target application. Specify -no_follow_children " + "to disable."); + +KNOB op_print_ret_addr +(KNOB_MODE_WRITEONCE, "pintool", "print_ret_addr", "0", "Print library call's return address\n" + "Print return addresses of library calls."); + +KNOB op_disable_dump +(KNOB_MODE_WRITEONCE, "pintool", "disable_dump", "0", "Disable Memory Dump\n" + "Disable Memory Dump"); + +KNOB op_ind_call_tracer +(KNOB_MODE_WRITEONCE, "pintool", "ind_call_tracer", "0", "Print all indirect-call (addr)\n" + "Print all indirect-call (addr)."); + +KNOB op_print_callback +(KNOB_MODE_WRITEONCE, "pintool", "print_callback", "0", "Print library callback functions\n" + "Print callback functions."); + +KNOB op_unknown_args +(KNOB_MODE_WRITEONCE, "pintool", "num_unknown_args", "10", "Number of unknown libcall args to print\n" + "Number of arguments to print for unknown library calls. Specify 0 to disable " + "unknown args printing."); + +KNOB op_max_args +(KNOB_MODE_WRITEONCE, "pintool", "num_max_args", "10", "Maximum number of arguments to print\n" + "Maximum number of arguments to print. This option allows to limit the number of " + "arguments to be printed. Specify 0 to disable args printing (including unknown)."); + +KNOB op_config_file_default +(KNOB_MODE_WRITEONCE, "pintool", "default_config", "1", "Use default config file.\n" + "Use config file that comes with drltrace and located in the same path. Specify " + "no_use_config and provide a path to custom config file using -config option."); + +KNOB op_config_file +(KNOB_MODE_WRITEONCE, "pintool", "config", "", "The path to custom config file.\n" + "Specify a custom path where config is located. The config file describes the prototype" + " of library functions for printing library call arguments. See drltrace documentation" + " for more details."); + +KNOB op_ignore_underscore +(KNOB_MODE_WRITEONCE, "pintool", "ignore_underscore", "0", "Ignores library routine names " + "starting with \"_\".\n" "Ignores library routine names starting with \"_\"."); + +KNOB op_only_to_lib +(KNOB_MODE_WRITEONCE, "pintool", "only_to_lib", "", "Only reports calls to the library . \n" + "Only reports calls to the library . Argument is case insensitive on Windows."); + +KNOB op_only_to_target +(KNOB_MODE_WRITEONCE, "pintool", "only_to_target", "XXXXXX", "Only reports calls to the target . \n" + "Only reports calls to the target . Argument is case insensitive on Windows."); + +KNOB op_trace_mode +(KNOB_MODE_WRITEONCE, "pintool", "trace_mode", "unique", "Setup level of trace\n" + "Setup the level of trace. (unique, relation, all, ind"); + +KNOB op_use_config +(KNOB_MODE_WRITEONCE, "pintool", "use_config", "1", "Use config file\n" + "Use config file for library call arguments printing. Specify no_use_config to disable."); + +KNOB op_ltracelib_ops +(KNOB_MODE_WRITEONCE, "pintool", "ltracelib_ops", "0", "(For internal use: sweeps up drltracelib options)\n" + "This is an internal option that sweeps up other options to pass to the drltracelib."); + +#include +#include +#include +#include +#include + +using namespace std; + +#define MAXIMUM_PATH 260 +#define MAX_SYM_RESULT 256 + +/* Where to write the trace */ +static FILE * outf; +static FILE* out_table; + +bool table_flag = true; +bool printed = false; +bool trace_ind_call = false; +bool found_module; +bool found_target; +int print_callback = 0; +int indent = 0; +int callnum = 0; +int retnum = 0; +int calltype; +int second_level_elements = 100; +const int pagesize = 0x1000; +PIN_MUTEX as_built_lock; +PIN_MUTEX as_built_lock2; + +typedef UINT8 byte; +typedef ADDRINT app_pc; +typedef ADDRINT ptr_uint_t; + +byte safe_buf[pagesize + 1] = { 0 }; +byte safe_buf2[pagesize + 1] = { 0 }; +byte safe_buf_pointer[(pagesize / 4) + 1] = { 0 }; +byte safe_buf_pointer2[(pagesize / 4) + 1] = { 0 }; +app_pc last_addr = 0x0; + +#ifdef X64 //windows x64 +int memory_size = 8; +#else +int memory_size = 4; +#endif + +void print_sp(ADDRINT addr); + +using namespace std; +vector> print_args(ADDRINT RSP, int count, int tid); +bool fast_safe_read(void* base, size_t size, void* out_buf, size_t* outsize); +void print_pointer_arrow(drltrc_pointer_type_t type); +bool print_if_printable(ADDRINT addr, int tid); +static VOID event_app_instruction(TRACE trace, VOID*); +static void print_ret_value(ADDRINT addr); +void dump_at_return(vector> dumped, int tid); + +std::set dlls; +std::vector file_related; +std::vector return_candidate; + +/* Avoid exe exports, as on Linux many apps have a ton of global symbols. */ +static app_pc exe_start; +static app_pc module_start; +static app_pc module_end; +static app_pc target_start; +static app_pc target_end; + +/* map to store function name and the number of argument */ +map func_arg_map; +map::iterator map_it; +map tid_callid; +map tid_retid; +typedef map> > CallListMap; + +CallListMap calls; + +TLS_KEY calllist_map_key; + +class MUTEX { +public: + MUTEX() { + PIN_LockClient(); + // PIN_MutexLock(&as_built_lock); + } + + ~MUTEX() { + PIN_UnlockClient(); + // PIN_MutexUnlock(&as_built_lock); + } +}; + +CallListMap& calllist_map_() { + return *(CallListMap*)(PIN_GetThreadData(calllist_map_key)); +} + +void insert_calllist(app_pc addr, vector> data) { + //FIXME: check x64 case + //fprintf(outf, "insert %x\n", addr); + //last_addr = (app_pc)addr+6; + calls.insert(make_pair(addr, data)); +} + +void increase_callnum(int tid) { + // new thread-id? + if (tid_callid.find(tid) == tid_callid.end()) { + tid_callid.insert(make_pair(tid, 0)); + } + else { + tid_callid[tid]++; + } +} + +int get_callnum(int tid) { + if (tid_callid.find(tid) == tid_callid.end()) { + return 0; + } + else { + return tid_callid[tid]; + } +} + +// some hardcoded routine to check file operations +bool is_file_related(char* symbol, bool check_offset) { + if (symbol == NULL) { + return false; + } + if ((strcasestr(symbol, "CreateFil") != NULL) || + (strcasestr(symbol, "ReadFil") != NULL) || + (strcasestr(symbol, "SetFilePoin") != NULL) || + (strcasestr(symbol, "WriteFil") != NULL)) { + + if (check_offset && strcasestr(symbol, "+0x0") != NULL) { + return true; + } + + else if (!check_offset) { + return true; + } + } + return false; +} + +void increase_retid(int tid) { + // new thread-id? + if (tid_retid.find(tid) == tid_retid.end()) { + tid_retid.insert(make_pair(tid, 0)); + } + else { + tid_retid[tid]++; + } +} + +int get_retid(int tid) { + if (tid_retid.find(tid) == tid_retid.end()) { + return 0; + } + else { + return tid_retid[tid]; + } +} + +/* return argument-map given the address */ +vector> check_calladdr(app_pc addr) { + //fprintf(outf, "finding %x\n", addr); + + auto calllist_map = calls; + if (calllist_map.find(addr) == calllist_map.end()) { + //fprintf(outf, "nothing in callist_map\n"); + vector> temp; + temp.push_back(make_pair(-1, -1)); + //vector> temp = calllist_map[last_addr]; + //calllist_map.erase(last_addr); + return temp; + + } + else { + //fprintf(outf, "something in callist_map\n"); + vector> data = calllist_map[addr]; + //calllist_map.erase(addr); + //fprintf(outf, "FOUND return patch with cid:%d\n", data); + return data; + } +} + +void insert_func_arg(string funcname, int args) { + func_arg_map.insert(make_pair(funcname, args)); +} + +int ret_func_arg(string funcname) { + + map_it = func_arg_map.find(funcname); + if (map_it != func_arg_map.end()) { + return func_arg_map[funcname]; + } + return -1; +} + +static std::string is_printable(ptr_uint_t arg_val) +{ + std::stringstream stream; + stream << std::hex << arg_val; + std::string result(stream.str()); + uint j; + std::string h; + std::string s; + if (result.length() != 8) + return s; + unsigned int x; + + for (j = 0; j < 4; j++) { + std::stringstream ss; + char c[5]; + h = ""; + h.push_back(result.at(j * 2)); + h.push_back(result.at(j * 2 + 1)); + ss << std::hex << h; + ss >> x; + if (static_cast(x) < 32 || static_cast(x) > 126) + return s; + sprintf(c, "%c", static_cast(x)); + s.push_back(c[0]); + } + reverse(s.begin(), s.end()); + return s; +} + +/* get the decoded string of a specific address (terminated by space or non-ascii char) */ +static std::string get_string(ptr_uint_t arg_val, ptr_uint_t ptr_val, size_t sz) +{ + std::string res; + std::string ret = is_printable(arg_val); + if (ret.length() < sz) + return ret; + res += ret; + /* keep access the subsequent value */ + if (arg_val == ptr_val) + return res; + while (true) { + ptr_uint_t deref = 0; + ptr_val += sz; + fast_safe_read((void*)(ptr_val), sz, &deref, NULL); + if (deref == 0) + break; + ret = is_printable(deref); + if (ret.length() < sz) { + res += ret; + break; + } + res += ret; + } + return res; +} + +//FIXME: I assume that all writable region as DATA pointer (regardless of execution) +/* determine whether a pointer is code pointer or data pointer + implement for the print structure feature for harness generation */ +static bool is_code_pointer(ptr_uint_t arg_val) +{ + //fprintf(outf, "\naddr:%x\n", arg_val); + VOID* pAddress = (void*)arg_val; + OS_MEMORY_AT_ADDR_INFORMATION info; + NATIVE_PID pid; + OS_GetPid(&pid); + if (OS_QueryMemory(pid, pAddress, &info).generic_err != OS_RETURN_CODE_QUERY_FAILED) { + if (info.Protection == (OS_PAGE_PROTECTION_TYPE_READ | OS_PAGE_PROTECTION_TYPE_EXECUTE)) + return true; + else + return false; + } + return false; +} + +/* determine whether an address is data or pointer + implement for the print structure feature for harness generation */ +static drltrc_pointer_type_t _is_pointer(ptr_uint_t arg_val, int sz) +{ + ptr_uint_t deref = 0; + bool ret = fast_safe_read((void*)arg_val, sz, &deref, NULL); + /* arg_val is a pointer */ + if (ret) { + if (is_code_pointer(arg_val)) + return DRLTRC_CODE_POINTER; + else + return DRLTRC_DATA_POINTER; + } + return DRLTRC_NONE_POINTER; +} + +/* try to print the module and the symbol of a specific address */ +static void print_mod_and_symbol(void* drcontext, ptr_uint_t addr) +{ + MUTEX lock; + IMG mod = IMG_FindByAddress(addr); + string sym, modname; + + /* get the module name and function name via the module table first */ + if (IMG_Valid(mod)) { + sym = RTN_FindNameByAddress(addr); + modname = IMG_Name(mod); + } + + /* if fail to get module name via module table => use GetModuleFileName API to get it */ + //if (modname.empty()) { + // MEMORY_BASIC_INFORMATION memInfo; + // TCHAR dlpath[MAX_PATH]; + // if (VirtualQuery((LPCVOID)addr, &memInfo, sizeof(memInfo)) != 0) { + // DWORD r = GetModuleFileName((HMODULE)memInfo.AllocationBase, dlpath, MAX_PATH); + // if (symres == DRSYM_SUCCESS || symres == DRSYM_ERROR_LINE_NOT_AVAILABLE) + // fprintf(outf, "[%s!%s]", dlpath, sym.name); + // else + // fprintf(outf, "[%s]", dlpath); + // } + //} + if (!modname.empty() && !sym.empty()) { + fprintf(outf, "[%s!%s]", modname.c_str(), sym.c_str()); + } + else if (modname.empty() && !sym.empty()) { + fprintf(outf, "[!%s]", sym.c_str()); + } + else if (!modname.empty() && sym.empty()) { + fprintf(outf, "[%s!unknown_symbol]", modname); + } + else { + fprintf(outf, ""); + } +} + +static void print_structure(void* drcontext, ptr_uint_t addr, int sz, int level, ptr_uint_t ptr_val) +{ + std::string ret; + drltrc_pointer_type_t type = _is_pointer(addr, sz); + /* access the value of the pointer, if it is not a pointer, try to decode it to be a string. */ + if (type == DRLTRC_NONE_POINTER) { + fprintf(outf, " (DATA) "); + ret = get_string(addr, ptr_val, sz); + if (ret.length() != 0) { + fprintf(outf, " ( string: %s )", ret.c_str()); + } + } + else if (type == DRLTRC_CODE_POINTER) { + fprintf(outf, " (CODE_POINTER) "); + print_mod_and_symbol(drcontext, addr); + } + else if (type == DRLTRC_DATA_POINTER) { + fprintf(outf, " (DATA_POINTER) "); + ptr_uint_t deref = 0; + bool flag = fast_safe_read((void*)addr, sz, &deref, NULL); + if (flag && level < MAX_PTR_DEPTH) { + fprintf(outf, " => " "%p", deref); + print_structure(drcontext, deref, sz, level + 1, addr); + } + } +} + +bool fast_safe_read(void* base, size_t size, void* out_buf, size_t* outsize) +{ + /* For all of our uses, a failure is rare, so we do not want + * to pay the cost of the syscall (DrMemi#265). + */ + bool res = true; + res = !!W::ReadProcessMemory(current_process, base, out_buf, size, (W::SIZE_T*)outsize); + return res; +} + +bool addr_belongs_module(app_pc addr) { + if (addr >= module_start && addr <= module_end) { + return true; + } + return false; +} + +bool addr_belongs_target(app_pc addr) { + if (addr >= target_start && addr <= target_end) { + return true; + } + return false; +} + +static void print_two_modules(app_pc inst_addr, app_pc target_addr) +{ + MUTEX lock; + IMG data1 = IMG_FindByAddress(inst_addr); + + if (!IMG_Valid(data1)) { + return; + } + auto modname1 = IMG_Name(data1); + + IMG data2 = IMG_FindByAddress(target_addr); + if (!IMG_Valid(data2)) { + return; + } + auto modname2 = IMG_Name(data2); + auto pair = modname1 + "--" + modname2; + + if (!dlls.count(pair)) { + char* remove_str = ":\\windows"; // use windows default path + if (strcasestr(modname1.c_str(), remove_str) == NULL && strcasestr(modname2.c_str(), remove_str) == NULL) { + fprintf(outf, "%s\n", modname2.c_str()); + } + else { + fprintf(outf, "Windows Library ===> %s\n", modname2.c_str()); + } + dlls.insert(pair); + } +} + +static void print_ind_address(app_pc src, app_pc dest, const char* prefix) { + fprintf(outf, "%s:" "%p" "," "%p" " \n", prefix, src, dest); +} + +int print_thread_id(int tid, bool is_call) { + if (is_call) { + increase_callnum(tid); + fprintf(outf, "CALLID[%d] TID[%d] ", get_callnum(tid), tid); + } + else { + increase_retid(tid); + fprintf(outf, "RETID[%d] TID[%d] ", get_retid(tid), tid); + } + + return tid; +} + +//// from instrcalls +static std::string print_address(app_pc addr, const char* prefix, bool newline) +{ + MUTEX lock; + IMG data = IMG_FindByAddress(addr); + string sym; + if (!IMG_Valid(data)) { + fprintf(outf, "%s " "%p" " ? ??:0\n", prefix, addr); + return sym; + } + + RTN rtn = RTN_FindByAddress(addr); + + if (RTN_Valid(rtn)) { + sym = RTN_Name(rtn); + ADDRINT rtn_addr = RTN_Address(rtn); + auto modname = IMG_Name(data); + if (modname.empty()) + modname = ""; + //fprintf(outf, "[LEV%d] %s " "%p" " %s!%s+" "%p", indent, prefix, addr, modname, sym.name, addr - data->start - sym.start_offs); + if (newline) + fprintf(outf, "%s""%p""(%s!%s+" "%p"")\n", prefix, addr, modname.c_str(), sym.c_str(), addr - rtn_addr); + else + fprintf(outf, "%s""%p""(%s!%s+" "%p"")", prefix, addr, modname.c_str(), sym.c_str(), addr - rtn_addr); + + } + else + + // apply when from? + if (newline) + fprintf(outf, "%s""%p""\n", prefix, addr); + else + fprintf(outf, "%s""%p""", prefix, addr); + return sym; +} + +static char* get_symbol_at_addr(app_pc addr) +{ + MUTEX lock; + static char ret_str[0x100] = { 0 }; + RTN rtn = RTN_FindByAddress(addr); + + if (RTN_Valid(rtn)) { + sprintf(ret_str, "%s+" "%p""", RTN_Name(rtn).c_str(), addr - RTN_Address(rtn)); + return ret_str; + } + + ret_str[0] = '?'; + ret_str[1] = '\0'; + return ret_str; +} + + +static void PIN_FAST_ANALYSIS_CALL at_call(app_pc instr_addr, app_pc target_addr, app_pc RSP, int tid, app_pc next_addr) +{ + if (tracemode == RELATION) { + if (addr_belongs_module(instr_addr) || addr_belongs_module(target_addr)) + print_two_modules(instr_addr, target_addr); + } + + else if (tracemode == ALL) { + if (found_module == true) { + // if call-from-module and call-to-outside + char name[0x30]; + bool print = false; + + if (print_callback == 1) { + if (addr_belongs_module(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "DC M2T @"); //in case of indirect jmp, this may be M2J + print = true; + } + } + + // if call-from-outside and call-to-module + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "DC T2M @"); + print = true; + } + + if (print) { + //PIN_MutexLock(&as_built_lock); + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + vector> dumped = print_args(RSP, 10, tid); + insert_calllist(next_addr, dumped); + //PIN_MutexUnlock(&as_built_lock); + } + } + } + + else if (tracemode == DOMINATOR) { + PIN_MutexLock(&as_built_lock); + // target to target (binary to binary) + char name[0x30]; + bool print = false; + bool dump = false; + + //We need to print address (simple information) + //print_address(target_addr, "==", true); + char* sym = get_symbol_at_addr(target_addr); + //fprintf(outf, "%s\n", sym); + if (is_file_related(sym, true)) { + //fprintf(outf, "sym:%s\n", sym); + sprintf(name, "FR "); // File Related + print = true; + dump = true; + } + + if (addr_belongs_target(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "DC T2T "); + print = true; + } + else if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "DC T2M "); + print = true; + } + + if (print) { + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); //print unique call_id and thread_id + //print_ret_addr(instr_addr, cur_instr_length); + //fprintf(outf, "here1\n"); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_sp(RSP); + + if (dump) { + print_args(RSP, 10, tid); + return_candidate.push_back(instr_addr); + } + else { + if (print_if_printable(RSP, tid)) { + return_candidate.push_back(instr_addr); + } + } + //fprintf(outf, "here2\n"); + } + PIN_MutexUnlock(&as_built_lock); + } +} + +static void PIN_FAST_ANALYSIS_CALL at_call_ind(app_pc instr_addr, app_pc target_addr, app_pc RSP, int tid, app_pc next_addr) +{ + app_pc ret_addr; + //fprintf(outf, "at_call_ind\n"); + + if (tracemode == RELATION) { + if (addr_belongs_module(instr_addr) || addr_belongs_module(target_addr)) + print_two_modules(instr_addr, target_addr); + } + + else if (tracemode == ALL) { + if (found_module == true) { + // if call-from-module and call-to-outside + char name[0x30]; + bool print = false; + + if (print_callback == 1) { + if (addr_belongs_module(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "IC M2T @"); //in case of indirect jmp, this may be M2J + print = true; + } + } + + // if call-from-outside and call-to-module + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "IC T2M @"); + print = true; + } + + if (print) { + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + + vector> dumped = print_args(RSP, 10, tid); + insert_calllist(next_addr, dumped); + } + } + } + + else if (tracemode == DOMINATOR) { + PIN_MutexLock(&as_built_lock); + // target to target (binary to binary) + char name[0x30]; + bool print = false; + bool dump = false; + + char* sym = get_symbol_at_addr(target_addr); + if (is_file_related(sym, true)) { + //fprintf(outf, "sym:%s\n", sym); + sprintf(name, "FR "); + print = true; + dump = true; + } + + //print_address(target_addr, "--", true); + + if (addr_belongs_target(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "IC T2T "); + print = true; + } + else if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "IC T2M "); + print = true; + } + + if (print) { + //fprintf(outf, "here3\n"); + //fprintf(outf, "ins_size:%d\n", cur_instr_length); + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); //print unique call_id and thread_id + //print_ret_addr(instr_addr, cur_instr_length); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_sp(RSP); + //print_args(10, tid); + if (dump) { + print_args(RSP, 10, tid); + return_candidate.push_back(instr_addr); + } + else { + if (print_if_printable(RSP, tid)) { + return_candidate.push_back(instr_addr); + } + } + } + PIN_MutexUnlock(&as_built_lock); + } +} + +static void PIN_FAST_ANALYSIS_CALL at_jmp_ind(app_pc instr_addr, app_pc target_addr, ADDRINT RSP, int tid) +{ + app_pc ret_addr; + //fprintf(outf, "at_jmp_ind\n"); + + if (tracemode == RELATION) { + if (addr_belongs_module(instr_addr) || addr_belongs_module(target_addr)) + print_two_modules(instr_addr, target_addr); + } + + else if (tracemode == ALL) { + if (found_module == true) { + // if call-from-module and call-to-outside + char name[0x30]; + bool print = false; + + if (print_callback == 1) { + if (addr_belongs_module(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "IJ M2T @"); //in case of indirect jmp, this may be M2J + + print = true; + } + } + + // if call-from-outside and call-to-module + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "IJ T2M @"); + print = true; + } + + if (print) { + //PIN_MutexLock(&as_built_lock); + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + vector> dumped = print_args(10, tid, ret_addr); + insert_calllist(ret_addr, dumped); + + //PIN_MutexUnlock(&as_built_lock); + } + } + } + + else if (tracemode == DOMINATOR) { + PIN_MutexLock(&as_built_lock); + // target to target (binary to binary) + char name[0x30]; + bool print = false; + bool dump = false; + + char* sym = get_symbol_at_addr(target_addr); + if (is_file_related(sym, true)) { + //fprintf(outf, "sym:%s\n", sym); + sprintf(name, "FR "); + print = true; + dump = true; + } + //print_address(target_addr, "~~", true); + + if (addr_belongs_target(instr_addr) && addr_belongs_target(target_addr)) { + sprintf(name, "IC T2T "); + print = true; + } + else if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "IC T2M "); + print = true; + } + + if (print) { + MUTEX mutex; + //fprintf(outf, "here5\n"); + fprintf(outf, "==\n"); + tid = print_thread_id(tid, true); //print unique call_id and thread_id + //print_ret_addr(instr_addr, cur_instr_length); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_sp(RSP); + //print_args(10, tid); + if (dump) { + print_args(RSP, 10, tid); + return_candidate.push_back(instr_addr); + } + else { + if (print_if_printable(RSP, tid)) { + return_candidate.push_back(instr_addr); + } + } + } + PIN_MutexUnlock(&as_built_lock); + } +} + +bool return_to_candidate(app_pc target_addr) { + app_pc del_item = 0; + size_t index = 0; + for (size_t i = 0; i < return_candidate.size(); i++) { + if (return_candidate[i] < target_addr + 7 && return_candidate[i] > target_addr - 7) { + del_item = return_candidate[i]; + index = i; + } + } + + + if (del_item == 0) + return false; + else { + return_candidate.erase(return_candidate.begin() + index); + return true; + } +} + +static void PIN_FAST_ANALYSIS_CALL at_return(app_pc instr_addr, app_pc target_addr, ADDRINT RSP, ADDRINT RAX, int tid) +{ + if (tracemode == RELATION) { + if (addr_belongs_module(instr_addr) || addr_belongs_module(target_addr)) + print_two_modules(instr_addr, target_addr); + } + + else if (tracemode == ALL) { + char* name; + bool print = false; + + if (print_callback == 1) { + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + name = "RET2M "; + print = true; + } + } + + if (found_module == true) { + if (addr_belongs_module(instr_addr) && addr_belongs_target(target_addr)) { + name = "RET2T "; + print = true; + } + } + + if (print) { + MUTEX mutex; + fprintf(outf, "==\n"); + tid = print_thread_id(tid, false); + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_ret_value(RAX); + + // dump the arguments of function + vector> dumped = check_calladdr(target_addr); + dump_at_return(dumped, tid); + } + } + + else if (tracemode == DOMINATOR) { + PIN_MutexLock(&as_built_lock); + + char name[0x30]; + bool print = false; + bool dump = false; + + /* + char* sym = get_symbol_at_addr(instr_addr); + if (is_file_related(sym, false)){ + //fprintf(outf, "sym:%s\n", sym); + sprintf(name, "RET_FR "); + print = true; + dump = true; + //fprintf(outf, "here2-2\n"); + } + //fprintf(outf, "here3\n"); + */ + + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + sprintf(name, "RET2M "); + //name = "RET2M "; + print = true; + } + + else if ((addr_belongs_module(instr_addr) && addr_belongs_target(target_addr)) || + (addr_belongs_target(instr_addr) && addr_belongs_target(target_addr))) { + sprintf(name, "RET2T "); + //name = "RET2T "; + print = true; + } + + if (print) { + MUTEX mutex; + //fprintf(outf, "here3-1\n"); + if (print == false) { + sprintf(name, "RETFR "); + } + + fprintf(outf, "==\n"); + tid = print_thread_id(tid, false); //print unique call_id and thread_id + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_sp(RSP); + print_ret_value(RAX); + } + + if (return_to_candidate(target_addr)) { + MUTEX mutex; + if (print == false) { + sprintf(name, "RETFR "); + } + + fprintf(outf, "==\n"); + tid = print_thread_id(tid, false); //print unique call_id and thread_id + print_address(instr_addr, name, false); + print_address(target_addr, "->", true); + print_sp(RSP); + print_ret_value(RAX); + } + PIN_MutexUnlock(&as_built_lock); + } + + else if (tracemode == INDIRECT) { + if (addr_belongs_target(instr_addr) && addr_belongs_module(target_addr)) { + print_ind_address(instr_addr, target_addr, "RETURN"); + } + } +} + + +void print_pointer_arrow(drltrc_pointer_type_t type) { + if (type == DRLTRC_NONE_POINTER) { + fprintf(outf, "[D]"); + } + else if (type == DRLTRC_CODE_POINTER) { + fprintf(outf, "[CP]"); + //fprintf(outf, " -> " "%p", deref); + } + else if (type == DRLTRC_DATA_POINTER) { + fprintf(outf, "[DP]"); + } +} + +void check_pointer_wholepage(int arg_index, int cid, int tid, char* fix) { + + void* deref = 0; + bool result = false; + drltrc_pointer_type_t _type; + ptr_uint_t* addr; + + char memdump_pn[0x100]; + char out_pn[0x100]; + + //init the array + for (int i = 0; i < (pagesize / 4) + 1; i++) + safe_buf_pointer[i] = 0; + + // read data and check whether it is pointer or not + for (int i = 0; i < pagesize + 1; i = i + 4) { + + addr = ((ptr_uint_t*)(safe_buf + i)); + //fprintf(outf, "%x ", *test); + result = fast_safe_read((void*)* addr, 4, (void*)& deref, NULL); + if (result) { + _type = _is_pointer(*addr, memory_size); + if (_type == DRLTRC_CODE_POINTER) { + safe_buf_pointer[i / 4] = DRLTRC_CODE_POINTER; + } + else { + safe_buf_pointer[i / 4] = DRLTRC_DATA_POINTER; + + // dump the second level pointer + if (i < second_level_elements && cid >= 0) { + sprintf(memdump_pn, "%s\\memdump\\t%d-c%d-a%d-%s", logdir.c_str(), tid, cid, arg_index, fix); + OS_MkDir(memdump_pn, 0777); + fast_safe_read((void*)* addr, pagesize, safe_buf2, NULL); + + sprintf(out_pn, "%s\\%d", memdump_pn, i); + FILE* out_fp = fopen(out_pn, "wb"); + fwrite(safe_buf2, 1, pagesize, out_fp); + fclose(out_fp); + } + } + } + else { + //fprintf(outf, "N"); + safe_buf_pointer[i / 4] = DRLTRC_NONE_POINTER; + } + } +} + +void dump_address(ptr_uint_t addr, int arg_index, int cid, int tid) { + //void *data = NULL; + int current_callid; + char* fix; + + // at-call || at-call-ind || at-jmp-ind + if (cid == -1) { + current_callid = get_callnum(tid); + fix = "pre"; + } + // at-return + else { + current_callid = get_retid(tid); + fix = "post"; + } + bool result = false; + size_t bytes_read; + char out_pn[0x100]; + + byte safe_buf[pagesize + 1] = { 0 }; + memset(safe_buf, 0x0, sizeof(safe_buf)); + + // safe_read target address + result = fast_safe_read((void*)addr, pagesize, safe_buf, &bytes_read); + check_pointer_wholepage(arg_index, cid, tid, fix); + + // store to file (e.g., memdump\\c1-a1.bin) + sprintf(out_pn, "memdump\\t%d-c%d-a%d.%s", tid, current_callid, arg_index, fix); + FILE* out_fp = fopen((logdir + "\\" + out_pn).c_str(), "wb"); + fwrite(safe_buf, 1, pagesize, out_fp); + fwrite(safe_buf_pointer, 1, pagesize / 4, out_fp); + fclose(out_fp); +} + +void dump_at_return(vector> dumped, int tid) { + + if (dumped.size() == 1) { + return; + } + int cid = dumped.at(0).second; + int len = dumped.size(); + + int _arg; + ptr_uint_t _addr; + + for (int i = 1; i < len; i++) { + _arg = dumped.at(i).first; + _addr = dumped.at(i).second; + + if (!op_disable_dump) + dump_address(_addr, _arg, cid, tid); + } +} + +/* +fast_safe_read(addr, sizeof(arg), &arg); +result = access_and_print((ptr_uint_t)arg, i, tid, addr); +*/ +bool is_printable_ascii(char* input) { + + for (int i = 0; i < 4; i++) { + if (!isprint(input[i])) + return false; + } + return true; +} + +char* should_access_and_print(ptr_uint_t arg, int index, int tid, ADDRINT* addr) { + bool rst = false; + char s_buf[0x101]; + static char s_out[0x101]; + static char s_out2[0x101]; + bzero(s_buf, 0x101); + bzero(s_out, 0x101); + bzero(s_out2, 0x101); + + rst = fast_safe_read((void*)arg, 0x100, (void*)& s_buf, NULL); + + if (index == 0 && rst == true) { + sprintf(s_out, "%s", s_buf); + sprintf(s_out2, "%s", s_buf); + if (is_printable_ascii(s_out)) { + //fprintf(outf, " -STR: %s\n", s_out); + return s_out; + } + + else if (is_printable_ascii(s_out2)) { + //fprintf(outf, " -STR: %s\n", s_out2); + return s_out2; + } + } + return NULL; +} + +bool access_and_print(ptr_uint_t arg, int index, int tid, ADDRINT* addr) { + bool rst = false; + char s_buf[0x101]; + char s_out[0x101]; + char s_out2[0x101]; + bzero(s_buf, 0x101); + rst = fast_safe_read((void*)arg, 0x100, &s_buf, NULL); + if (index == 0 && rst == true) { + sprintf(s_out, "%s", s_buf); + sprintf(s_out2, "%s", s_buf); + if (is_printable_ascii(s_out)) + fprintf(outf, " -STR: %s\n", s_out); + else if (is_printable_ascii(s_out2)) + fprintf(outf, " -STR: %s\n", s_out2); + } + + std::string ret = is_printable(arg); + void* deref = 0; + void* next_arg = NULL; + bool result = false; + bool has_datapointer = false; + + //FIXME: make it as a function + result = fast_safe_read((void*)arg, 4, &deref, NULL); + fprintf(outf, " -A%d: ""%p""", index, arg); + if (result) { + if (ret.length() == 4) + fprintf(outf, " (str:%s)", ret.c_str()); + } + + // 1st dereference + result = fast_safe_read((void*)arg, 4, &deref, NULL); + if (result) { + drltrc_pointer_type_t _type = _is_pointer(arg, memory_size); + print_pointer_arrow(_type); + + // test dump here + // TODO: now we only consider dump for the first layer + if (_type == DRLTRC_DATA_POINTER) { + has_datapointer = true; + //fprintf(outf, "dump_address()\n"); + if (!op_disable_dump) + dump_address(arg, index, -1, tid); + } + + fprintf(outf, " > " "%p", deref); + ret = is_printable((ptr_uint_t)deref); + if (ret.length() == 4) { + if (deref > 0) + fprintf(outf, " (str:%s)", ret.c_str()); + } + } + + // 2nd dereference + next_arg = deref; + deref = 0; + result = fast_safe_read((void*)next_arg, 4, &deref, NULL); + if (result) { + drltrc_pointer_type_t _type = _is_pointer((ptr_uint_t)next_arg, memory_size); + print_pointer_arrow(_type); + fprintf(outf, " > " "%p", deref); + ret = is_printable((ptr_uint_t)deref); + if (ret.length() == 4) { + if (deref > 0) + fprintf(outf, " (str:%s)", ret.c_str()); + } + } + + // 3rd dereference + next_arg = deref; + deref = 0; + result = fast_safe_read((void*)next_arg, 4, &deref, NULL); + if (result) { + drltrc_pointer_type_t _type = _is_pointer((ptr_uint_t)next_arg, memory_size); + print_pointer_arrow(_type); + fprintf(outf, " > " "%p", deref); + ret = is_printable((ptr_uint_t)deref); + if (ret.length() == 4) { + if (deref > 0) + fprintf(outf, " (str:%s)", ret.c_str()); + } + } + + return has_datapointer; +} + +bool print_if_printable(ADDRINT addr, int tid) { + ADDRINT arg; + char* rst; + + fast_safe_read((void*)addr, sizeof(arg), &arg, nullptr); + rst = should_access_and_print((ptr_uint_t)arg, 0, tid, (ADDRINT*)addr); + if (rst != NULL) { + fprintf(outf, " -STR: %s\n", rst); + return true; + } + return false; +} + +void print_sp(ADDRINT addr) { + ADDRINT arg = 0x3f3f3f3f; + fast_safe_read((void*)addr, sizeof(arg), &arg, nullptr); + + fprintf(outf, "SP: ""%p"" -> ""%p""\n", addr, arg); +} + +vector> print_args(app_pc RSP, int count, int tid) { + //void print_args (int count){ + ADDRINT* addr; + void* arg; + bool result = false; + vector> dumped_args; + +#ifdef X64 //windows x64 + fprintf(outf, " - ARG%d: ""%p""\n", 0, mc.rcx); + fprintf(outf, " - ARG%d: ""%p""\n", 1, mc.rdx); + fprintf(outf, " - ARG%d: ""%p""\n", 2, mc.r8); + fprintf(outf, " - ARG%d: ""%p""\n", 3, mc.r9); + + for (int i = 4; i < count; i++) { + addr = (reg_t*)(mc.xsp + (i + 0) * sizeof(reg_t)); + fast_safe_read(addr, sizeof(arg), &arg); + fprintf(outf, " - ARG%d: ""%p""\n", i, arg); + } +#else // windows x86 + dumped_args.push_back(make_pair(-1, get_callnum(tid))); + for (int i = 0; i < count; i++) { + arg = nullptr; + addr = (ADDRINT*)(RSP + (i + 0) * sizeof(ADDRINT)); + fast_safe_read(addr, sizeof(arg), &arg, nullptr); + result = access_and_print((ptr_uint_t)arg, i, tid, addr); + if (result == true) { + dumped_args.push_back(make_pair(i, (ptr_uint_t)arg)); + } + fprintf(outf, "\n"); + } +#endif + return dumped_args; +} + +static void print_ret_addr(app_pc instr_addr, int instr_size) { + /* + dr_mcontext_t mc = {sizeof(mc), DR_MC_ALL}; + dr_get_mcontext(dr_get_current_drcontext(), &mc); + + reg_t *addr; + void *arg; + addr = (reg_t *) (mc.xsp); + fast_safe_read(addr, sizeof(arg), &arg); + fprintf(outf, "(ret:""%p"", xsp:""%p"")", arg, addr); + */ + + fprintf(outf, "(ret:""%p"") ", instr_addr + instr_size); +} + +static void print_ret_value(ADDRINT addr) { + fprintf(outf, "RETVAL: ""%p""\n", addr); +} + +static bool library_matches_filter(IMG info) +{ + if (!op_only_to_lib.Value().empty()) { + string libname = IMG_Name(info); + return (!libname.empty() && strcasestr(libname.c_str(), op_only_to_lib.Value().c_str()) != NULL); + } + return true; +} + +static bool target_matches_filter(IMG info) +{ + if (!op_only_to_target.Value().empty()) { + string libname = IMG_Name(info); + return (!libname.empty() && strcasestr(libname.c_str(), op_only_to_target.Value().c_str()) != NULL); + } + return true; +} + +struct ImageTracker { + ADDRINT start, end; + ADDRINT entry; + ADDRINT r1, r2, r3; + std::string path; +}; + +static std::vector ss; +void loaded(IMG info) { + static int i; + + ss.push_back({ + IMG_LowAddress(info), + IMG_HighAddress(info) + 1, + IMG_LowAddress(info) + IMG_EntryAddress(info), + 0, 0, 0, + IMG_Name(info) + }); + return; +} + +static void event_exit(INT32 code, VOID* v); + +static void event_module_load(IMG info, VOID*) +{ + loaded(info); + + //fprintf(outf, "!!! %s\n", IMG_Name(info).c_str()); + char* remove_str = ":\\windows"; // use windows default path + if (tracemode == UNIQUE && strcasestr(IMG_Name(info).c_str(), remove_str) == NULL) { + dlls.insert(IMG_Name(info)); + } + + if (!IMG_IsMainExecutable(info) && library_matches_filter(info)) { + module_start = IMG_LowAddress(info); + module_end = IMG_HighAddress(info) + 1; + + if (tracemode == ALL) { + print_address(module_start, "LIBRARY MODULE START ADDR:", true); + print_address(module_end, "LIBRARY MODULE END ADDR:", true); + } + + found_module = true; + } + + if (IMG_IsMainExecutable(info) && tracemode == ALL) { + fprintf(outf, "CHECKING MODULE...\n"); + exe_start = IMG_LowAddress(info); + target_start = exe_start; + target_end = IMG_HighAddress(info) + 1; + + print_address(target_start, "TARGET MODULE START ADDR:", true); + print_address(target_end, "TARGET MODULE END ADDR:", true); + found_target = true; + } + + if (target_matches_filter(info)) { + target_start = IMG_LowAddress(info); + target_end = IMG_HighAddress(info) + 1; + + if (tracemode == ALL) { + print_address(target_start, "TARGET MODULE START ADDR:", true); + print_address(target_end, "TARGET MODULE END ADDR:", true); + } + + found_target = true; + } + + auto rtn = RTN_FindByName(info, "NtTerminateProcess"); + if (RTN_Valid(rtn)) { + RTN_Open(rtn); + RTN_InsertCall(rtn, IPOINT_BEFORE, (AFUNPTR)event_exit, IARG_UINT32, 0, IARG_ADDRINT, 0, IARG_END); + RTN_Close(rtn); + } +} + +static void init_file_related(void) { + file_related.push_back("CreateFile"); + file_related.push_back("ReadFile"); + file_related.push_back("SetFilePointer"); + file_related.push_back("WriteFile"); +} + +static void open_log_file(void) +{ + if (op_logdir.Value().compare("-") == 0) { + outf = stderr; + } + else { + NATIVE_PID pid; OS_GetPid(&pid); + char filename[100]; + sprintf(filename, "drltrace.%d.log", pid); + outf = fopen((logdir + "\\" + filename).c_str(), "wb"); + ASSERT(outf, "failed to open log file"); + //setvbuf(outf, nullptr, _IONBF, 0); + VNOTIFY(0, "drltrace log file is %s""\n", buf); + } +} + +void reset_memdumnp_storage(void) { + const char memdump_pn[] = "memdump"; + { + std::stringstream s; + s << "cd \"" << logdir << "\" && rmdir /s /q " << memdump_pn; + system(s.str().c_str()); + } + OS_MkDir((logdir + "\\" + memdump_pn).c_str(), 0777); +} + +static void open_functype_file(void) +{ + VNOTIFY(0, "Opened functype is %s""\n", op_functype.Value().c_str()); + FILE* fp; + + char line[255]; + fp = fopen(op_functype.Value().c_str(), "r"); + + while (fgets(line, sizeof(line), fp) != NULL) { + + char* val1 = strtok(line, "|"); + char* val2 = strtok(NULL, "|"); + int val3 = atoi(strtok(NULL, "|")); + //fprintf("insert %s %d\n", val1, val3); + std::string val1_str(val1); + insert_func_arg(val1_str, val3); + } + +} + +static void +event_thread_init(void* drcontext) +{ + PIN_SetThreadData(calllist_map_key, new CallListMap); +} + +static void +event_thread_exit(void* drcontext) +{ +} + +void drmodtrack_dump(FILE* out) { + int index = 0; + fprintf(outf, "Module Table: version 4, count %d\n", ss.size()); + for (auto&& item : ss) { + fprintf(out, "%-3d, %-3d, 0x%08x, 0x%08x, 0x%08x, %016x, 0x%08x, 0x%08x, %s\n", + index, index, item.start, item.end, item.entry, item.r1, item.r2, item.r3, item.path.c_str()); + index++; + } +} + +static void event_exit(INT32 code, VOID* v) +{ + //drmgr_unregister_thread_init_event(event_thread_init); + //drmgr_unregister_thread_exit_event(event_thread_exit); + + //if (op_use_config) + // libcalls_hashtable_delete(); + + if (outf != stderr && tracemode != UNIQUE && tracemode != RELATION) { + + MUTEX mutex; + fprintf(outf, "\n\n==\n"); + drmodtrack_dump(outf); + fclose(outf); + } +} + +INT32 print_usage() +{ + cerr << "This tool counts the number of dynamic instructions executed" << endl; + cerr << endl << KNOB_BASE::StringKnobSummary() << endl; + return -1; +} + +int main(int argc, CHAR** argv) +{ + if (PIN_Init(argc, argv)) { + return print_usage(); + } + +#define MAXPATH 0x1000 + char logdir_buf[MAXPATH] = {}; + if (!W::GetFullPathNameA(op_logdir.Value().c_str(), MAXPATH, logdir_buf, NULL)) { + return print_usage(); + } + + logdir = logdir_buf; + + PIN_InitSymbols(); + PIN_AddFiniFunction(event_exit, nullptr); + + calllist_map_key = PIN_CreateThreadDataKey(NULL); + + //FOR BB + TRACE_AddInstrumentFunction(event_app_instruction, nullptr); + + //FOR MODULE + IMG_AddInstrumentFunction(event_module_load, nullptr); + + trace_ind_call = op_ind_call_tracer.Value(); + + PIN_MutexInit(&as_built_lock); + PIN_MutexInit(&as_built_lock2); + open_log_file(); + init_file_related(); + + if (!strcasestr("none", op_functype.Value().c_str()) != NULL) + open_functype_file(); + print_callback = op_print_callback; + + //take care of memdump storage + reset_memdumnp_storage(); + + // read option for tracing + // 1) unique: print out unique dlls loaded + // 2) relation: print out all dlls related with specified dll + // 3) all: print all traces between two dlls + if (strcasestr("unique", op_trace_mode.Value().c_str()) != NULL) { + tracemode = UNIQUE; + } + else if (strcasestr("relation", op_trace_mode.Value().c_str()) != NULL) { + tracemode = RELATION; + } + else if (strcasestr("all", op_trace_mode.Value().c_str()) != NULL) { + tracemode = ALL; + } + else if (strcasestr("indirect", op_trace_mode.Value().c_str()) != NULL) { + tracemode = INDIRECT; + } + else if (strcasestr("dominator", op_trace_mode.Value().c_str()) != NULL) { + tracemode = DOMINATOR; + } + + current_process = W::GetCurrentProcess(); + PIN_StartProgram(); +} + +static VOID event_app_instruction(TRACE trace, VOID*) +{ + for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl)) { + + for (INS instr = BBL_InsHead(bbl); INS_Valid(instr); instr = INS_Next(instr)) { + auto next_addr = INS_Address(instr) + INS_Size(instr); + if (INS_IsDirectCall(instr)) { + INS_InsertCall(instr, IPOINT_BEFORE, (AFUNPTR)at_call, IARG_FAST_ANALYSIS_CALL + , IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_REG_VALUE, REG_ESP, IARG_THREAD_ID, IARG_ADDRINT, next_addr, IARG_END); + } + + else if (INS_IsRet(instr)) { + INS_InsertCall(instr, IPOINT_BEFORE, (AFUNPTR)at_return, IARG_FAST_ANALYSIS_CALL + , IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_REG_VALUE, REG_ESP, IARG_REG_VALUE, REG_EAX, IARG_THREAD_ID, IARG_END); + } + + else if (INS_IsIndirectControlFlow(instr)) { + if (INS_Opcode(instr) != XED_ICLASS_JMP) + INS_InsertCall(instr, IPOINT_BEFORE, (AFUNPTR)at_call_ind, IARG_FAST_ANALYSIS_CALL + , IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_REG_VALUE, REG_ESP, IARG_THREAD_ID, IARG_ADDRINT, next_addr, IARG_END); + else + INS_InsertCall(instr, IPOINT_BEFORE, (AFUNPTR)at_jmp_ind, IARG_FAST_ANALYSIS_CALL + , IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_REG_VALUE, REG_ESP, IARG_THREAD_ID, IARG_END); + } + + } + } + return; +} + diff --git a/harnessgen/logger.py b/harnessgen/logger.py new file mode 100644 index 0000000..9261161 --- /dev/null +++ b/harnessgen/logger.py @@ -0,0 +1,25 @@ +import logging + +class CuteHandler(logging.StreamHandler): + def emit(self, record): + color = hash(record.name) % 7 + 31 + try: + record.name = ("\x1b[%dm" % color) + record.name + "\x1b[0m" + except Exception: + pass + + try: + record.msg = ("\x1b[%dm" % color) + record.msg + "\x1b[0m" + except Exception: + pass + + super(CuteHandler, self).emit(record) + +def getlogger(name): + logger = logging.getLogger(name) + logger.setLevel(logging.INFO) + stream_handler = CuteHandler() + stream_handler.setFormatter(logging.Formatter('%(levelname)-7s | %(asctime)-23s | %(name)-8s | %(message)s')) + logger.addHandler(stream_handler) + + return logger diff --git a/harnessgen/requirements.txt b/harnessgen/requirements.txt new file mode 100644 index 0000000..7959aca --- /dev/null +++ b/harnessgen/requirements.txt @@ -0,0 +1,3 @@ +tqdm +networkx +matplotlib \ No newline at end of file diff --git a/harnessgen/syn-multi.py b/harnessgen/syn-multi.py new file mode 100644 index 0000000..660e854 --- /dev/null +++ b/harnessgen/syn-multi.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python2 +import os +import glob +import argparse +import signal +import typing +from util import exit_gracefully +from logger import * +from harconf import * +from template import * + +from common import Synthesizer + +""" +./syn-multi.py harness -t temptrace_math +./syn-multi.py harness -t temptrace_math --start test +./syn-multi.py harness -t temptrace_math --start test --sample-name input +""" + +""" +[TARCE_DIR] + - cor1_1 : 1st trace using the 1st correct input + - cor1_2 : 2nd trace using the 1st correct input + - cor2_1 : 3rd trace using the 2nd correct input + - functype : function type information from IDA ==> fixed with dynamic information + - decompile: decompiled results where the library call started + +[NOTE] + - cor1_1 and cor1_2 is used to infer the const value of argument + - cor1_1 and cor2_1 is used to infer the possible impact of file-related operation +""" + +logger = getlogger("Synthesizer") + + +class Identifier: + def __init__(self, traces: typing.Dict[str, Synthesizer]): + self.cor_trace1 = traces["cor1"] + self.cor_trace2 = traces["cor2"] + self.diff_trace = traces["diff"] + + # comparison result using same input should have result + self.comp_cor = self.compare_cortrace() + + # comparison result using different inputs may not have result (i.e., null dict) + self.comp_diff = self.compare_difftrace() + + self.report = self.make_report(self.comp_cor, self.comp_diff) + + def make_report(self, cor_dict, diff_dict): + # NOTE: assume that diff_dict may not have any result + # - also, assume that the number of CIDs are same + + out = {} + for cid in cor_dict.keys(): + arg = cor_dict[cid] + + if cid in diff_dict: + arg += "\n" + diff_dict[cid] + + out[cid] = arg + + return out + + def analyze_args(self, args1, args2, cor): + # args e.g., [[(2291628, 'DP'), (11, 'D')], [(9376112, 'CP'), (2347535189, 'D')]] + # args e.g., [[(210, 'D')]] + + # NOTE: we assume that the length of args are same + + out = "" + + if cor == True: + tag = " // [DIFF] (Multi-runs using same input)" + else: + tag = " // [DIFF] (Result with different inputs)" + + for x in range(len(args1)): + rst = "" + + # pointer + if args1[x][0][1] == 'DP' or args1[x][0][1] == 'CP': + if args1[x][1][0] == args2[x][1][0]: + rst = "SAME" + else: + rst = "DIFFERENT" + out += "\n // - Arg[%d]: %s (referenced value is %s) |" % ( + x, NAMEDIC[args1[x][0][1]], rst) + + # data + else: + if args1[x][0][0] == args2[x][0][0]: + rst = "SAME" + else: + rst = "DIFFERENT" + out += "\n // - Arg[%d]: DATA (value is %s) |" % (x, rst) + + out = tag + out + + return out + + def compare_cortrace(self): + # compare cor_trace1 and cor_trace2 + """ report + - arg1: data_pointer (referenced value SAME/DIFF), arg2: data (SAME) + - data_pointer / code_pointer / data, SAME or DIFF + """ + + result = {} + + for cid in self.cor_trace1.trace.cid_sequence: + + if cid in self.cor_trace1.trace.calltrace and \ + cid in self.cor_trace2.trace.calltrace: + + calltrace1 = self.cor_trace1.trace.calltrace[cid] + calltrace2 = self.cor_trace2.trace.calltrace[cid] + + rettrace1 = self.cor_trace1.trace.rettrace[cid] + rettrace2 = self.cor_trace2.trace.rettrace[cid] + + funcname1 = calltrace1.dst_symbol + funcname2 = calltrace2.dst_symbol + + analyzed_result = self.analyze_args( + calltrace1.args, calltrace2.args, cor=True) + result[cid] = analyzed_result + # print analyzed_result + + else: + result[cid] = None + + return result + + def compare_difftrace(self): + # compare cor_Trace1 and diff_trace + + # NOTE: we not only compare the argument if the length of cids are same + # - we don't provide alignment for this problem + + result = {} + + # check all function names for each sequence is same + for cid in self.cor_trace1.trace.cid_sequence: + if cid in self.cor_trace1.trace.calltrace and \ + cid in self.diff_trace.trace.calltrace: + + calltrace1 = self.cor_trace1.trace.calltrace[cid] + calltrace2 = self.diff_trace.trace.calltrace[cid] + funcname1 = calltrace1.dst_symbol + funcname2 = calltrace2.dst_symbol + + if funcname1 != funcname2: + return result + + for cid in self.cor_trace1.trace.cid_sequence: + + if cid in self.cor_trace1.trace.calltrace and \ + cid in self.diff_trace.trace.calltrace: + + calltrace1 = self.cor_trace1.trace.calltrace[cid] + calltrace2 = self.diff_trace.trace.calltrace[cid] + + rettrace1 = self.cor_trace1.trace.rettrace[cid] + rettrace2 = self.diff_trace.trace.rettrace[cid] + + funcname1 = calltrace1.dst_symbol + funcname2 = calltrace2.dst_symbol + + analyzed_result = self.analyze_args( + calltrace1.args, calltrace2.args, cor=False) + result[cid] = analyzed_result + # print analyzed_result + + else: + result[cid] = None + + return result + + +class MultiSynthesizer(Synthesizer): + # one of the core logics for discovering pointer relationship + def build_body(self, report): + # we should consider both call_tarce and ret_trace + # - report: dictionary from the diff analysis + + pending_flag = True + pending_count = 0 + prev_src_addr = 0 + prev_dst_addr = 0 + + # 1) we need to handle used argument (raw value and pointer, ret-chain) + for cid in self.trace.cid_sequence: + variables = [] # defined variables: e.g., int a=0 + arguments = [] # used arguments: func(&a) + calltrace = self.trace.calltrace[cid] + rettrace = self.trace.rettrace[cid] + funcname = calltrace.dst_symbol.decode() + fi = self.trace.find_function(calltrace.dst_addr) + ret_type = fi.ret_type + args_type = fi.args + + msg = "" + args = calltrace.args + args_dump = calltrace.args_dump + args_ptr = calltrace.pointer + dst_addr = calltrace.dst_addr + src_addr = calltrace.src_addr + + if src_addr == prev_src_addr and dst_addr == prev_dst_addr: + pending_count = pending_count + 1 + pending_flag = True + else: + pending_count = 0 + pending_flag = False + + # one of core functions (we should do pointer analysis here) + need_to_define, arguments = self.ret_arg_code( + cid, args, args_dump, args_type, args_ptr) + input_digging_result = self.dig_userinput( + cid, args, args_dump, args_ptr) + need_to_define_str = ' '.join(need_to_define).strip() + self.history[cid] = (need_to_define, arguments) + + # build function snippet + # print funcname, arguments + if funcname not in self.defined_func: + self.defined_func.append(funcname) + func_snippet = FUNC.replace("{funcname}", funcname) + else: + func_snippet = FUNC_WO.replace("{funcname}", funcname) + + if need_to_define_str == '': + func_snippet = func_snippet.replace( + "{print_cid}", "// Harness function #%d " % cid) + else: + func_snippet = func_snippet.replace( + "{print_cid}", "// Harness function #%d \n %s" % (cid, ' '.join(need_to_define))) + + func_snippet = func_snippet.replace( + "{arguments}", ', '.join(arguments)) + + if ret_type == '': + func_snippet = func_snippet.replace("{ret_statement}", "") + func_snippet = func_snippet.replace("{dbg_printf}", + 'dbg_printf("%s\\n");' % (funcname)) + else: + func_snippet = func_snippet.replace( + "{ret_statement}", "%s %s_ret_%d = " % (ret_type, funcname, cid)) + func_snippet = func_snippet.replace("{dbg_printf}", + 'dbg_printf("%s, ret = %%d\\n", %s_ret_%d);' % (funcname, funcname, cid)) + + msg += "\n%s" % report[cid] + + # print "dist:", prev_src_addr, hex(abs (prev_src_addr - src_addr)) + if prev_src_addr != 0 and abs(prev_src_addr - src_addr) > THRESHOLD: + msg += "\n // Distance between prev/current libcall:%s, we recomment you to check the condition between them " % ( + hex(abs(prev_src_addr - src_addr))) + + if pending_count > 0: + msg += "\n // [LOOP] This is %dth execution of %s() " % ( + pending_count+1, funcname) + self.body.append(msg+func_snippet) + else: + self.body.append(msg+func_snippet) + + prev_src_addr = src_addr + prev_dst_addr = dst_addr + + +def main(): + # DEFINE PARSER + parser = argparse.ArgumentParser() + subparser = parser.add_subparsers(help="subparser") + + # subparser: harness generation + har_parser = subparser.add_parser('harness') + har_parser.add_argument("-t", "--trace", dest="trace_dir", type=str, + default=None, help="Trace dir collected from DynamoRIO", + required=False) + har_parser.add_argument("-s", "--start", dest="start_func", type=str, + default=None, help="name of the starting function to process", + required=False) + har_parser.add_argument("-sample", "--sample-name", dest="sample_name", type=str, + default=None, help="name of the original sample name", + required=False) + har_parser.set_defaults(action='harness') + + args = parser.parse_args() + # END PARSER + + # Ctrl-c handler + signal.signal(signal.SIGINT, exit_gracefully( + signal.getsignal(signal.SIGINT))) + + # Start harness synthesizer + if args.action == 'harness': + cor_trace_1 = glob.glob(os.path.join( + args.trace_dir, MAIN_TRACE) + "/*.log")[0] + cor_trace_2 = glob.glob(os.path.join( + args.trace_dir, SECOND_TRACE) + "/*.log")[0] + diff_trace = glob.glob(os.path.join( + args.trace_dir, DIFF_TRACE) + "/*.log")[0] + + dumpdir = os.path.join(args.trace_dir, MAIN_TRACE, "memdump") + dumpdir2 = os.path.join(args.trace_dir, SECOND_TRACE, "memdump") + dumpdir_diff = os.path.join(args.trace_dir, DIFF_TRACE, "memdump") + functype_pn = os.path.join(args.trace_dir, "functype_") + + syn_cor1 = MultiSynthesizer(cor_trace_1, dumpdir, + functype_pn, args.start_func, args.sample_name) + syn_cor2 = MultiSynthesizer(cor_trace_2, dumpdir2, + functype_pn, args.start_func, args.sample_name) + syn_diff = MultiSynthesizer(diff_trace, dumpdir_diff, + functype_pn, args.start_func, args.sample_name) + + traces = {} + traces["cor1"] = syn_cor1 + traces["cor2"] = syn_cor2 + traces["diff"] = syn_diff + + identifier = Identifier(traces) + report = identifier.report + + syn_cor1.build_body(report) + syn_cor1.emit_code() + # syn.search_pointer(0x1f75ac0) + # syn.search_pointer(0x65aa89a0) + + +if __name__ == '__main__': + main() diff --git a/harnessgen/synthesizer.py b/harnessgen/synthesizer.py new file mode 100644 index 0000000..c513f43 --- /dev/null +++ b/harnessgen/synthesizer.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python2 +import argparse +import operator +import signal +from util import exit_gracefully + +from logger import * +from harconf import * +from template import * + +from common import SimpleTrace, Synthesizer + +""" +./synthesizer.py harness -t temptrace/trace.log -d temptrace +./synthesizer.py harness -t temptrace/trace.log -d temptrace --start avformat_open_input +./synthesizer.py harness -t temptrace_gom/trace.log -d temptrace_gom -f temptrace_gom/functype --start avformat_open_input --sample-name small.mp4 + +./synthesizer.py diff --input-dummy temptrace_aimp/aimp_without_playlist.log --input-parse temptrace_aimp/aimp_with_playlist.log --output temptrace_aimp/trace.log +""" + +logger = getlogger("Synthesizer") + + +class Differ: + def __init__(self, in_dummy_pn, in_parse_pn, output_pn): + self.in_dummy_pn = in_dummy_pn + self.in_parse_pn = in_parse_pn + self.output_pn = output_pn + + self.dummy_trace = SimpleTrace(self.in_dummy_pn) + self.parse_trace = SimpleTrace(self.in_parse_pn) + + dummy_trace_sort = sorted( + list(self.dummy_trace.unique_call.items()), key=operator.itemgetter(1)) + parse_trace_sort = sorted( + list(self.parse_trace.unique_call.items()), key=operator.itemgetter(1)) + + dummy_calls = [] + for k, v in dummy_trace_sort: + dummy_calls.append(k) + + for k, v in parse_trace_sort: + if k not in dummy_calls: + print("CALLID[{0:05d}]: {1}".format(v, k)) + + # print len(self.dummy_trace.unique_call) + # print len(self.parse_trace.unique_call) + + +class SingleSynthesizer(Synthesizer): + def build_body(self): + # we should consider both call_tarce and ret_trace + + # 1) we need to handle used argument (raw value and pointer, ret-chain) + for cid in self.trace.cid_sequence: + variables = [] # defined variables: e.g., int a=0 + arguments = [] # used arguments: func(&a) + calltrace = self.trace.calltrace[cid] + rettrace = self.trace.rettrace[cid] + funcname = calltrace.dst_symbol.decode() + + fi = self.trace.find_function(calltrace.dst_addr) + ret_type = fi.ret_type + args_type = fi.args + + args = calltrace.args + args_dump = calltrace.args_dump + args_ptr = calltrace.pointer + + # one of core functions (we should do pointer analysis here) + need_to_define, arguments = self.ret_arg_code(cid, args, args_dump, args_type, args_ptr) + need_to_define_str = ' '.join(need_to_define).strip() + self.history[cid] = (need_to_define, arguments) + + func_snippet = FUNC.replace("{funcname}", funcname) + + if need_to_define_str == '': + func_snippet = func_snippet.replace("{print_cid}", "/* Harness function #%d */" % cid) + else: + func_snippet = func_snippet.replace("{print_cid}", "/* Harness function #%d */\n %s" % (cid, ' '.join(need_to_define))) + + func_snippet = func_snippet.replace("{arguments}", ', '.join(arguments)) + + if ret_type == '': + func_snippet = func_snippet.replace("{ret_statement}", "") + func_snippet = func_snippet.replace("{dbg_printf}", + 'dbg_printf("%s\\n");' % (funcname)) + else: + func_snippet = func_snippet.replace("{ret_statement}", "%s %s_ret = " % (ret_type, funcname)) + func_snippet = func_snippet.replace("{dbg_printf}", + 'dbg_printf("%s, ret = %%d\\n", %s_ret);' % (funcname, funcname)) + + self.body.append(func_snippet) + + +def main(): + # DEFINE PARSER + parser = argparse.ArgumentParser() + subparser = parser.add_subparsers(help="subparser") + + # subparser: harness generation + har_parser = subparser.add_parser('harness') + har_parser.add_argument("-t", "--trace", dest="trace_file", type=str, + default=None, help="Trace file collected from DynamoRIO", + required=False) + har_parser.add_argument("-d", "--memory-dump", dest="dump_dir", type=str, + default=None, help="memory dump file directory (pre/post)", + required=False) + har_parser.add_argument("-f", "--function-type", dest="functype", type=str, + default=None, help="function type information from IDA (and fixed pointer by trace)", + required=False) + har_parser.add_argument("-s", "--start", dest="start_func", type=str, + default=None, help="name of the starting function to process", + required=False) + har_parser.add_argument("-sample", "--sample-name", dest="sample_name", type=str, + default=None, help="name of the original sample name", + required=False) + har_parser.set_defaults(action='harness') + + # subparser: diff trace + diff_parser = subparser.add_parser('diff') + diff_parser.add_argument('--input-dummy', dest="input_dummy", type=str) + diff_parser.add_argument('--input-parse', dest="input_parse", type=str) + diff_parser.add_argument('--output', dest="output", type=str) + diff_parser.set_defaults(action='diff') + + args = parser.parse_args() + # END PARSER + + # Ctrl-c handler + signal.signal(signal.SIGINT, exit_gracefully(signal.getsignal(signal.SIGINT))) + + # Start harness synthesizer + if args.action == 'harness': + syn = SingleSynthesizer(args.trace_file, args.dump_dir, + args.functype, args.start_func, args.sample_name) + syn.build_body() + syn.emit_code() + # syn.search_pointer(0x1f75ac0) + # syn.search_pointer(0x65aa89a0) + + elif args.action == 'diff': + diff = Differ(args.input_dummy, args.input_parse, args.output) + + +if __name__ == '__main__': + main() diff --git a/harnessgen/template.py b/harnessgen/template.py new file mode 100644 index 0000000..0d6a71d --- /dev/null +++ b/harnessgen/template.py @@ -0,0 +1,90 @@ + +ASCII_BYTE = ' !"#\\$%&\'\\(\\)\\*\\+,-\\./0123456789:;<=>\\?@ABCDEFGHIJKLMNOPQRSTUVWXYZ\\[\\]\\^_`abcdefghijklmnopqrstuvwxyz\\{\\|\\}\\\\~\t' + +# Directory structure +MAIN_TRACE = 'cor1_1' +SECOND_TRACE = 'cor1_2' +DIFF_TRACE = 'cor2_1' +INPUT1 = 'input1' +INPUT2 = 'input2' +FUNCTYPE = 'functype' + +HEADER = """ +#include +#include +#include +#include +#include +#include +#include + +#define dbg_printf (void)printf + +// Macro to help to loading functions +#define LOAD_FUNC(h, n) \\ + n##_func = (n##_func_t)GetProcAddress(h, #n); \\ + if (!n##_func) { \\ + dbg_printf("failed to load function " #n "\\n"); \\ + exit(1); \\ + } + +// Macro help creating unique nop functions +#define NOP(x) \\ + int nop##x() { \\ + dbg_printf("==> nop%d called, %p\\n", ##x, _ReturnAddress());\\ + return (DWORD)x; \\ + } + +HMODULE dlllib; + +{typedef} + +""" + +FUZZME = """ +void fuzz_me(char* filename){ + +{funcdef} + +{harness} + +} +""" + +MAIN = """ +int main(int argc, char ** argv) +{ + if (argc < 2) { + printf("Usage %s: \\n", argv[0]); + printf(" e.g., harness.exe input\\n"); + exit(1); + } + + dlllib = LoadLibraryA("%s"); + if (dlllib == NULL){ + dbg_printf("failed to load library, gle = %d\\n", GetLastError()); + exit(1); + } + + char * filename = argv[1]; + fuzz_me(filename); + return 0; +} +""" + +""" + LOAD_FUNC(dlllib, avformat_open_input); + int ret_avformat_open_input = avformat_open_input_func(&ctx_org, filename, 0, &avformat_open_input_arg3); // zeros: if pointer one page ==> copy original page to that page ==> if error + dbg_printf("avformat_open_input: ret = %d\n", ret_avformat_open_input); // @jinho: check crash / progress +""" + +FUNC = """ + {print_cid} + LOAD_FUNC(dlllib, {funcname}); + {ret_statement}{funcname}_func({arguments}); + {dbg_printf} """ + +FUNC_WO = """ + {print_cid} + {ret_statement}{funcname}_func({arguments}); + {dbg_printf} """ \ No newline at end of file diff --git a/harnessgen/util.py b/harnessgen/util.py new file mode 100644 index 0000000..19a2d33 --- /dev/null +++ b/harnessgen/util.py @@ -0,0 +1,49 @@ +import sys +import struct +import string +import signal + +printable = set(string.printable.encode()) + + +def strings(data: bytes): + found_str = b"" + while True: + if not data: + break + for char in data: + if char in printable: + found_str += bytes([char]) + elif len(found_str) >= 4: + yield found_str + break + break + yield b'' + + +def u32(b, off=0): + return struct.unpack(" ").lower().startswith('y'): + sys.exit(1) + + except KeyboardInterrupt: + print("Ok ok, quitting") + sys.exit(1) + + # restore the exit gracefully handler here + signal.signal(signal.SIGINT, _exit_gracefully) + return _exit_gracefully diff --git a/harnessgen/util/ida_func_type.py b/harnessgen/util/ida_func_type.py new file mode 100644 index 0000000..4f84f25 --- /dev/null +++ b/harnessgen/util/ida_func_type.py @@ -0,0 +1,54 @@ +import json +import os +import typing + +from idautils import Segments, Functions +from idc import get_segm_start, get_segm_end, get_func_name +import idaapi + + +def serialize(tif: idaapi.tinfo_t) -> typing.Union[dict, None]: + fi = idaapi.func_type_data_t() + if not tif.get_func_details(fi): + return None + + args = [{"type": str(arg.type), "name": arg.name} for arg in fi] + + return ({ + 'args': args, + 'ret_type': str(fi.rettype), + 'cc': '__stdcall' if fi.cc == idaapi.CM_CC_STDCALL else '__cdecl' + }) + + +def main(): + idaapi.auto_wait() + base = idaapi.get_imagebase() + tif = idaapi.tinfo_t() + + f = open(os.environ.get("DESTPATH", "functype_"), 'w') + + for ea in Segments(): + # only code segment + if idaapi.segtype(ea) != idaapi.SEG_CODE: + continue + + for fva in Functions(get_segm_start(ea), get_segm_end(ea)): + func_name = get_func_name(fva) + has_type = idaapi.get_tinfo(tif, fva) or idaapi.guess_tinfo(tif, fva) + + if not has_type: + continue + + info = serialize(tif) + if info is None: + continue + + print(hex(fva-base)[:-1], "|", func_name, "|", tif, "|", len(info['args'])) + f.write("0x%x|%s|%s\n" % (fva-base, func_name, json.dumps(info))) + + f.close() + idaapi.qexit(0) + + +main()