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