xref: /aosp_15_r20/external/AFLplusplus/afl-wine-trace (revision 08b48e0b10e97b33e7b60c5b6e2243bd915777f2) 
1#!/usr/bin/env python3
2
3import os
4import sys
5import pefile
6import shutil
7import subprocess
8
9if len(sys.argv) < 2:
10    print("[afl-wine-trace] usage: ./afl-wine-trace binary [args...]\n")
11    exit(1)
12
13if os.getenv("AFL_PATH"):
14    my_dir = os.getenv("AFL_PATH")
15else:
16    my_dir = os.path.dirname(os.path.abspath(__file__))
17
18os.environ["WINELOADERNOEXEC"] = "1"
19
20pe = pefile.PE(sys.argv[1])
21
22if "AFL_ENTRYPOINT" not in os.environ:
23    os.environ["AFL_ENTRYPOINT"] = "0x%x" % (pe.OPTIONAL_HEADER.ImageBase + pe.OPTIONAL_HEADER.AddressOfEntryPoint)
24if not os.getenv("AFL_INST_LIBS"):
25    if "AFL_CODE_START" not in os.environ:
26        os.environ["AFL_CODE_START"] = "0x%x" % (pe.OPTIONAL_HEADER.ImageBase + pe.OPTIONAL_HEADER.BaseOfCode)
27    if "AFL_CODE_END" not in os.environ:
28        os.environ["AFL_CODE_END"] = "0x%x" % (pe.OPTIONAL_HEADER.ImageBase + pe.OPTIONAL_HEADER.BaseOfCode + pe.OPTIONAL_HEADER.SizeOfCode)
29
30if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"] or pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_IA64"]:
31    os.environ["QEMU_SET_ENV"] = "LD_PRELOAD=" + os.path.join(my_dir, "qemu_mode/unsigaction/unsigaction64.so") + ",WINEARCH=win64"
32else:
33    os.environ["QEMU_SET_ENV"] = "LD_PRELOAD=" + os.path.join(my_dir, "qemu_mode/unsigaction/unsigaction32.so") + ",WINEARCH=win32"
34
35if os.getenv("WINECOV_QEMU_PATH"):
36    qemu_path = os.getenv("WINECOV_QEMU_PATH")
37elif os.path.exists(os.path.join(my_dir, "afl-qemu-trace")):
38    qemu_path = os.path.join(my_dir, "afl-qemu-trace")
39else:
40    qemu_path = "qemu-"
41    if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"] or pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_IA64"]:
42        qemu_path += "x86_64"
43    elif pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]:
44        qemu_path += "i386"
45    else:
46        print ("[afl-wine-trace] unsuppoted architecture\n")
47        exit(1)
48    qemu_path = shutil.which(qemu_path)
49
50wine_path = None
51if os.getenv("AFL_WINE_PATH"):
52    wine_path = os.getenv("AFL_WINE_PATH")
53else:
54    if not wine_path and shutil.which("wine"):
55        wine_path = shutil.which("wine")
56    if not wine_path and os.path.exists("/usr/bin/wine"):
57        wine_path = "/usr/bin/wine"
58    if not wine_path and os.path.exists("/usr/lib/wine/wine"):
59        wine_path = "/usr/lib/wine/wine"
60    if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"] or pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_IA64"]:
61        wine_path += "64"
62    elif pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]:
63        pass
64    else:
65        print ("[afl-wine-trace] unsopported architecture\n")
66        exit(1)
67
68argv = sys.argv[1:]
69for i in range(len(argv)):
70    if ".cur_input" in argv[i]:
71        # Get the Wine translated path using the winepath tool
72        arg_translated = subprocess.run([os.path.join(os.path.dirname(wine_path), "winepath"), "--windows", argv[i]], universal_newlines=True, stdout=subprocess.PIPE).stdout
73        # Remove the spurious LF at the end of the path
74        if len(arg_translated) > 0 and arg_translated[-1] == '\n':
75            arg_translated = arg_translated[:-1]
76        argv[i] = arg_translated
77        break
78
79print("[afl-wine-trace] exec:", " ".join([qemu_path, wine_path] + argv))
80os.execve(qemu_path, [qemu_path, wine_path] + argv, os.environ)
81