def trace_payload(self, data, info):
# Legacy implementation of -trace (now -trace_cb) using libxdc_edge_callback hook.
# This is generally slower and produces different bitmaps so we execute it in
# a different phase as part of calibration stage.
# Optionally pickup pt_trace_dump* files as well in case both methods are enabled.
trace_edge_in = self.config.work_dir + "/redqueen_workdir_%d/pt_trace_results.txt" % self.pid
trace_dump_in = self.config.work_dir + "/pt_trace_dump_%d" % self.pid
trace_edge_out = self.config.work_dir + "/traces/fuzz_cb_%05d.lst" % info[
'id']
trace_dump_out = self.config.work_dir + "/traces/fuzz_cb_%05d.bin" % info[
'id']
logger.info("%s Tracing payload_%05d.." % (self, info['id']))
if len(data) > self.payload_limit:
data = data[:self.payload_limit]
try:
self.q.set_payload(data)
old_timeout = self.q.get_timeout()
self.q.set_timeout(0)
self.q.set_trace_mode(True)
exec_res = self.q.send_payload()
self.q.set_trace_mode(False)
self.q.set_timeout(old_timeout)
if os.path.exists(trace_edge_in):
with open(trace_edge_in, 'rb') as f_in:
with lz4.LZ4FrameFile(trace_edge_out + ".lz4",
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
if os.path.exists(trace_dump_in):
with open(trace_dump_in, 'rb') as f_in:
with lz4.LZ4FrameFile(trace_dump_out + ".lz4",
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
if not exec_res.is_regular():
self.statistics.event_reload(exec_res.exit_reason)
self.q.reload()
except Exception as e:
logger.info("%s Failed to produce trace %s: %s (skipping..)" %
(self, trace_edge_out, e))
return None
return exec_res
def trace_payload(self, data, info):
trace_file_in = self.config.argument_values[
'work_dir'] + "/redqueen_workdir_%d/pt_trace_results.txt" % self.slave_id
trace_folder = self.config.argument_values['work_dir'] + "/traces/"
trace_file_out = trace_folder + "payload_%05d" % info['id']
log_slave("Tracing payload_%05d.." % info['id'], self.slave_id)
try:
self.q.set_payload(data)
exec_res = self.q.execute_in_trace_mode(timeout_detection=False)
with open(trace_file_in, 'rb') as f_in:
with lz4.LZ4FrameFile(
trace_file_out + ".lz4",
'wb',
compression_level=lz4.COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
if not exec_res.is_regular():
self.statistics.event_reload()
self.q.reload()
except Exception as e:
log_slave(
"Failed to produce trace %s: %s (skipping..)" %
(trace_file_out, e), self.slave_id)
return None
return exec_res
def read(self, path: str, mode: str):
with lz4.LZ4FrameFile(path) as lz4d:
archive = tarfile_open(mode=mode, fileobj=lz4d, **self.lz4_kwargs)
try:
yield archive
finally:
archive.close()
def parse_trace_file(self, trace_file, trace_id):
if not os.path.isfile(trace_file):
print_note("Could not find trace file %s, skipping.." % trace_file)
return None
gaps = set()
bbs = set()
edges = set()
with lz4.LZ4FrameFile(trace_file, 'rb') as f:
#for line in f.readlines():
# info = (json.loads(line.decode()))
# if 'trace_enable' in info:
# gaps.add(info['trace_enable'])
# if 'edge' in info:
# edges.add("%s_%s" % (info['edge'][0], info['edge'][1]))
# bbs.add(info['edge'][0])
# bbs.add(info['edge'][1])
# slightly faster than above line-wise json parsing
for m in re.finditer("\{.(\w+).: \[?(\d+),?(\d+)?\]? \}",
f.read().decode()):
if m.group(1) == "trace_enable":
gaps.add(m.group(2))
if m.group(1) == "edge":
edges.add("%s_%s" % (m.group(2), m.group(3)))
bbs.add(m.group(2))
bbs.add(m.group(3))
return {'bbs': bbs, 'edges': edges, 'gaps': gaps}
def write(self, path: str, mode: str):
with lz4.LZ4FrameFile(path, mode=mode[0]) as lz4c:
archive = tarfile_open(mode=mode, fileobj=lz4c, **self.lz4_kwargs)
try:
yield archive
finally:
archive.close()
def store_trace(self, node, tmp_trace):
if tmp_trace and os.path.exists(tmp_trace):
trace_dump_out = "%s/traces/fuzz_%05d.bin" % (self.config.work_dir,
node.get_id())
with open(tmp_trace, 'rb') as f_in:
with lz4.LZ4FrameFile(
trace_dump_out + ".lz4",
'wb',
compression_level=lz4.COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
os.remove(tmp_trace)
def generate_traces_worker(config, pid, input_list):
def sigterm_handler(signal, frame):
if q:
q.async_exit()
sys.exit(0)
# override config - workdir root should be tempdir!
pname = mp.current_process().name
config.argument_values['work_dir'] += "_%s" % pname
config.argument_values['purge'] = True
prepare_working_dir(config)
work_dir = config.argument_values['work_dir']
trace_dir = config.argument_values["input"] + "/traces/"
signal.signal(signal.SIGTERM, sigterm_handler)
os.setpgrp()
q = qemu(1337, config, debug_mode=False)
if not q.start():
print_fail("%s: Could not start Qemu. Exit." % pname)
return None
pbar = tqdm(total=len(input_list),
desc=pname,
dynamic_ncols=True,
smoothing=0.1,
position=pid + 1)
try:
for input_path in input_list:
trace_file = trace_dir + os.path.basename(input_path) + ".lz4"
if os.path.exists(trace_file):
#printf("Skipping %s.." % os.path.basename(input_path))
pbar.update()
continue
#print("Processing %s.." % os.path.basename(input_path))
if funky_trace_run(q, input_path):
with open(
work_dir +
"/redqueen_workdir_1337/pt_trace_results.txt",
'rb') as f_in:
with lz4.LZ4FrameFile(trace_file,
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
pbar.update()
except:
q.async_exit()
raise
q.shutdown()
def parse_trace_file(trace_file):
if not os.path.isfile(trace_file):
logger.warn("Could not find trace file %s, skipping.." %
trace_file)
return None
bbs = set()
edges = dict()
with lz4.LZ4FrameFile(trace_file, 'rb') as f:
for m in re.finditer("([\da-f]+),([\da-f]+)", f.read().decode()):
edges["%s,%s" % (m.group(1), m.group(2))] = 1
bbs.add(m.group(1))
bbs.add(m.group(2))
return {'bbs': bbs, 'edges': edges}
def parse_trace_file(trace_file):
if not os.path.isfile(trace_file):
print_note("Could not find trace file %s, skipping.." % trace_file)
return None
gaps = set()
bbs = set()
edges = set()
with lz4.LZ4FrameFile(trace_file, 'rb') as f:
for m in re.finditer("\{.(\w+).: \[?(\d+),?(\d+)?\]? \}",
f.read().decode()):
if m.group(1) == "trace_enable":
gaps.add(m.group(2))
if m.group(1) == "edge":
edges.add("%s,%s" % (m.group(2), m.group(3)))
bbs.add(m.group(2))
bbs.add(m.group(3))
return {'bbs': bbs, 'edges': edges, 'gaps': gaps}
def compress_dir(entry):
dst_file = Path(entry.path + ".tar.lz4")
if not dst_file.is_file():
with open(entry.path + ".txt", 'w') as txt:
for path, _subdirs, files in walk(entry.path):
for name in files:
txt.write(path + name)
txt.write("\n")
print("Compress: " + entry.name)
lz4_file = frame.LZ4FrameFile(
dst_file, mode='w',
compression_level=frame.COMPRESSIONLEVEL_MIN,
block_checksum=True
)
with tarfile.open(mode='w', fileobj=lz4_file) as tar_xz_file:
tar_xz_file.add(entry.path)
lz4_file.close()
else:
print(dst_file.__str__() + " already exist")
print("Remove Directory: " + entry.path)
rmtree(entry.path)
def main(args):
blacklist = Utils.init_blacklist(BLACKLIST)
col_stats = ColumnStatsD.fromFile(COLSTATS)
sys.stdout.flush()
dataset = MalDictionary.fromJsonFile(
args.input_file,
blacklist,
col_stats
)
tags = sorted(dataset.query_tags)
if args.limit:
tags = tags[:args.limit]
tag_map_file = open(args.tag_map, 'a') if args.tag_map else None
counter = 0
for tag in tags:
out_name = args.output_files.replace('XXX', '%03d' % counter)
short_name = os.path.basename(out_name)
if '.' in short_name:
short_name = short_name[:short_name.index('.')]
if tag_map_file:
tag_map_file.write('{}:{}\n'.format(tag, short_name))
counter += 1
contents = dataset.filter(lambda x: x.tag == tag)
with open(out_name, 'wb') as f:
if out_name.endswith('.lz4'):
f = lz4frame.LZ4FrameFile(f, mode='wb')
sys.stderr.write('\r[{}/{}] tag {}'.format(counter, len(tags), tag))
pickle.dump(contents, f)
f.close()
sys.stderr.write('\rDone \n')
if tag_map_file:
tag_map_file.close()
def generate_traces_worker(config, pid, work_queue):
dump_mode = True
def sigterm_handler(signal, frame):
if q:
q.async_exit()
sys.exit(0)
pname = mp.current_process().name
pnum = mp.current_process()._identity[0]
ptdump_path = config.ptdump_path
if config.resume:
# spawn worker in same workdir, picking up snapshot + page_cache
config.purge = False # not needed?
qemu_id = int(pnum) # get unique qemu ID != {0,1337}
else:
# spawn worker in separate workdir, booting a new VM state
config.work_dir += "_%s" % pname
config.purge = True # not needed?
qemu_id = 1337 # debug instance
prepare_working_dir(config)
work_dir = config.work_dir
trace_dir = config.input + "/traces/"
signal.signal(signal.SIGTERM, sigterm_handler)
os.setpgrp()
# FIXME: really ugly switch between -trace and -dump_pt
if dump_mode:
print("Tracing in '-trace' mode..")
# new dump_pt mode - translate to edge trace in separate step
config.trace = True
config.trace_cb = False
else:
# traditional -trace mode - more noisy and no bitmap to check
print("Tracing in legacy '-trace_cb' mode..")
config.trace = False
config.trace_cb = True
q = qemu(qemu_id, config, debug_mode=False)
if not q.start():
logger.error("%s: Could not start Qemu. Exit." % pname)
return None
pbar = tqdm(total=len(work_queue),
desc=pname,
dynamic_ncols=True,
smoothing=0.1,
position=pid + 1)
f = tempfile.NamedTemporaryFile(delete=False)
tmpfile = f.name
f.close()
try:
for input_path, dump_file, trace_file in work_queue:
print("\nProcessing %s.." % os.path.basename(input_path))
if dump_mode:
# -trace mode (pt dump)
if not os.path.exists(dump_file):
qemu_file = work_dir + "/pt_trace_dump_%d" % qemu_id
if simple_trace_run(q, read_binary_file(input_path),
q.send_payload):
with open(qemu_file, 'rb') as f_in:
with lz4.LZ4FrameFile(
dump_file,
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
if not os.path.exists(trace_file):
with tempfile.NamedTemporaryFile(delete=False) as pt_tmp:
with lz4.LZ4FrameFile(dump_file, 'rb') as pt_dump_lz4:
shutil.copyfileobj(pt_dump_lz4, pt_tmp)
pt_tmp.close()
cmd = [
ptdump_path, work_dir + "/page_cache", pt_tmp.name,
tmpfile
]
for i in range(2):
key = "ip" + str(i)
if getattr(config, key, None):
range_a = hex(getattr(config, key)[0]).replace(
"L", "")
range_b = hex(getattr(config, key)[1]).replace(
"L", "")
cmd += [range_a, range_b]
try:
subprocess.run(cmd, timeout=180)
os.unlink(pt_tmp.name)
except subprocess.TimeoutExpired as e:
print(e)
os.unlink(pt_tmp.name)
continue
with open(tmpfile, 'rb') as f_in:
with lz4.LZ4FrameFile(
trace_file,
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
else:
# -trace_cb mode (libxdc callback)
if not os.path.exists(trace_file):
qemu_file = work_dir + "/redqueen_workdir_%d/pt_trace_results.txt" % qemu_id
if simple_trace_run(q, read_binary_file(input_path),
q.send_payload):
with open(qemu_file, 'rb') as f_in:
with lz4.LZ4FrameFile(
trace_file,
'wb',
compression_level=lz4.
COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
pbar.update()
except Exception:
q.async_exit()
raise
finally:
os.unlink(tmpfile)
q.shutdown()
def generate_traces(config, input_list):
work_dir = config.argument_values['work_dir']
data_dir = config.argument_values["input"]
trace_dir = data_dir + "/traces/"
if data_dir == work_dir:
print_note("Workdir must be separate from input/data dir. Aborting.")
return None
prepare_working_dir(config)
if os.path.exists(trace_dir):
print_note(
"Input data_dir already has a traces/ subdir. Skipping trace generation..\n"
)
return trace_dir
# real deal. delete trace dir if it exists and (re-)create traces
shutil.rmtree(trace_dir, ignore_errors=True)
os.makedirs(trace_dir)
# TODO What is the effect of not defining a trace region? will it trace?
if not config.argument_values['ip0']:
print_warning("No trace region configured!")
if os.path.exists(work_dir + "redqueen_workdir_1337"):
print_fail(
"Leftover files from 1337 instance. This should not happen.")
return None
q = qemu(1337, config, debug_mode=False)
if not q.start():
print_fail("Could not start Qemu. Exit.")
return None
start = time.time()
try:
for input_path, nid, timestamp in input_list:
print("Processing: %s" % input_path)
q.set_payload(read_binary_file(input_path))
exec_res = q.execute_in_trace_mode(timeout_detection=False)
if not exec_res:
print_note("Failed to execute input %s. Continuing anyway..." %
input_path)
assert (q.restart())
continue
# TODO: reboot by default, persistent by option
if exec_res.is_crash():
q.reload()
with open(work_dir + "/redqueen_workdir_1337/pt_trace_results.txt",
'rb') as f_in:
with lz4.LZ4FrameFile(
trace_dir + os.path.basename(input_path) + ".lz4",
'wb',
compression_level=lz4.COMPRESSIONLEVEL_MINHC) as f_out:
shutil.copyfileobj(f_in, f_out)
except:
raise
finally:
q.async_exit()
end = time.time()
print("Time taken: %.2fs" % (end - start))
return trace_dir
