1#!/usr/local/bin/python3.8
2import sys
3import os
4import subprocess
5import tempfile
6import select
7import shutil
8HELP="""OVERVIEW: ZESTI like wrapper of KLEE
9
10USAGE:  klee-zesti [klee-options] <input bytecode> <concrete program arguments>
11
12WARNING this script is not equivalent to ZESTI in ICSE 2012. It just provides a similar interface to KLEE. Namely it first explores the path of <concrete program arguments> and then continues symbolic execution from that point. Most importantly it does not implement the ZESTI searcher.
13"""
14
15
16KLEE="klee"
17GEN_BOUT="gen-bout"
18
19def find_klee_bin_dir():
20  global KLEE
21  global GEN_BOUT
22  bin_dir = os.path.dirname(os.path.realpath(__file__))
23  KLEE = bin_dir + "/klee"
24  GEN_BOUT = bin_dir + "/gen-bout"
25  if not os.path.isfile(KLEE):
26      print("WARNING can't find klee at " + KLEE)
27      KLEE= shutil.which("klee")
28      print("Using klee in PATH", KLEE)
29  if not os.path.isfile(GEN_BOUT):
30      print("WARNING can't find gen-bout at " + GEN_BOUT)
31      GEN_BOUT= shutil.which("gen-bout")
32      print("Using gen-bout in PATH", GEN_BOUT)
33  if GEN_BOUT is None or KLEE is None:
34      print("Failed to find KLEE at this script location or in PATH. Quitting ...")
35      sys.exit(1)
36  print("Using", KLEE)
37
38
39
40def split_args():
41  prog = None
42  prog_args = []
43  klee_args = []
44  is_progargs = False
45  for a in sys.argv[1:]:
46      if is_progargs:
47          prog_args += [a]
48      elif a.startswith("-"):
49          klee_args += [a]
50      else:
51          prog = a
52          is_progargs = True
53  return klee_args, prog, prog_args
54
55def maybe_file_size(name):
56  try:
57    return os.path.getsize(name)
58  except:
59    return None
60
61def get_stdin_file(tmpdir):
62  stdin = ""
63  stdin_size = 0
64  if sys.stdin in select.select([sys.stdin], [], [], 0)[0]:
65    stdin += sys.stdin.readline()
66  if stdin == "":
67      return None, stdin_size
68  stdin_file_name = tmpdir.name + "/stdin.file"
69  with open(stdin_file_name, 'w') as f:
70    stdin_size = f.write(stdin)
71  return stdin_file_name, stdin_size
72
73
74
75def prog_args_to_posix(prog_args):
76  posix_args = []
77  sym_file = 'A'
78  sym_file_sizes = []
79  gen_out_args = []
80  for parg in prog_args:
81      file_size = maybe_file_size(parg)
82      if file_size is None:
83          posix_args += ['--sym-arg', str(len(parg))]
84          gen_out_args += [parg]
85      else:
86          sym_file_sizes += [file_size]
87          posix_args += [sym_file]
88          sym_file = chr(ord(sym_file) + 1)
89          gen_out_args += ['--sym-file', parg]
90
91  if ord(sym_file) - ord('A') > 0:
92      posix_args += ['--sym-files', str(ord(sym_file) - ord('A')), str(max(sym_file_sizes))]
93  return posix_args, gen_out_args
94
95def create_ktest_file(gen_out_args, tmpdir):
96  out_file=tmpdir + "/test.ktest"
97  subprocess.run([GEN_BOUT, "--bout-file", out_file] + gen_out_args, check=True)
98  return out_file
99
100
101
102def main():
103  klee_args, prog, prog_args = split_args()
104  if len(sys.argv) == 1 or prog is None:
105      print(HELP)
106      return
107  find_klee_bin_dir()
108  tmpdir = tempfile.TemporaryDirectory()
109  stdin_file, stdin_size = get_stdin_file(tmpdir)
110  posix_args, gen_out_args = prog_args_to_posix(prog_args)
111  if stdin_file is not None:
112      gen_out_args += ["--sym-stdin", stdin_file]
113      posix_args += ["--sym-stdin", str(stdin_size)]
114  ktest_file = create_ktest_file(gen_out_args,tmpdir.name)
115  klee_args += ["-seed-file=" + ktest_file]
116
117  proc = subprocess.Popen([KLEE] + klee_args + [prog] + posix_args, stdout=sys.stdout, stderr=sys.stderr)
118  while proc.returncode is None:
119      try:
120        proc.wait()
121      except KeyboardInterrupt:
122        pass # This is expected when stopping KLEE, so we wait for KLEE to finish
123  sys.exit(proc.returncode)
124
125
126if __name__ == "__main__":
127  main()
128
129
130