Adds pytest tests

Author: iosifache

pyproject.toml

@@ -8,7 +8,7 @@ readme = "README.md"
 packages = [{include = "zeratool_lib"}]
 
 [tool.poetry.dependencies]
-python = "^3.11"
+python = "^3.10"
 angr = "^9.2.52"
 claripy = "^9.2.52"
 ipython = "^8.13.2"
@@ -17,6 +17,9 @@ r2pipe = "^1.8.0"
 timeout-decorator = "^0.5.0"
 tqdm = "^4.65.0"
 
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.3.1"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

tests/test_zeratool_binaries.py

@@ -0,0 +1,67 @@
+from zeratool_lib import ZeratoolExploit, ZeratoolInputStreams, exploit
+
+
+def test_bof() -> None:
+    exploit_data = exploit(
+        "zeratool_binaries/bin/bof_32.elf",
+        input_stream=ZeratoolInputStreams.STDIN,
+        overflow_only=False,
+        force_shellcode=True,
+        skip_check=True,
+    )
+
+    assert exploit_data.payload is not None
+    assert exploit_data.outcome == ZeratoolExploit.Outcomes.SHELL
+
+
+def test_bof_with_nx() -> None:
+    exploit_data = exploit(
+        "zeratool_binaries/bin/bof_nx_32.elf",
+        input_stream=ZeratoolInputStreams.STDIN,
+        overflow_only=False,
+        leak_format=b"{.*}",
+        skip_check=True,
+    )
+
+    assert exploit_data.payload is not None
+    assert exploit_data.outcome == ZeratoolExploit.Outcomes.SHELL
+
+
+def test_bof_with_win_function() -> None:
+    exploit_data = exploit(
+        "zeratool_binaries/bin/bof_win_32.elf",
+        input_stream=ZeratoolInputStreams.STDIN,
+        overflow_only=True,
+        win_funcs=["get_secret"],
+        leak_format=b"{.*}",
+        skip_check=True,
+    )
+
+    assert exploit_data.payload is not None
+    assert exploit_data.outcome == ZeratoolExploit.Outcomes.CALL_TO_WIN
+
+
+def test_format_string_attack_for_leak() -> None:
+    exploit(
+        "zeratool_binaries/bin/read_stack_32.elf",
+        input_stream=ZeratoolInputStreams.STDIN,
+        format_only=True,
+        leak_format=b"(.*)BEGIN PRIVATE KEY(.*)",
+        skip_check=False,
+    )
+
+    # TODO: Check why this is not generating an exploit despite the replication of the
+    #       test from original README.md and tests/format_test.py.
+
+
+def test_format_string_attack_for_win_call() -> None:
+    exploit(
+        "zeratool_binaries/bin/format_pc_write_32.elf",
+        input_stream=ZeratoolInputStreams.STDIN,
+        format_only=True,
+        win_funcs=["secret_function"],
+        leak_format=b"(.*)BEGIN PRIVATE KEY(.*)",
+    )
+
+    # TODO: Check why this is not generating an exploit despite the replication of the
+    #       test from original README.md and tests/format_test.py.

tests/zeratool_binaries/Makefile

@@ -0,0 +1,28 @@
+# Adapted from Zeratool's tests/Makefile
+
+build_flags := -fno-stack-protector -z execstack \
+	-Wno-implicit-function-declaration -no-pie \
+	-Wno-format-security -fcf-protection=none -mno-shstk
+build_NX_flags := -fno-stack-protector \
+	-Wno-implicit-function-declaration -no-pie \
+	-Wno-format-security -z relro -fcf-protection=none -mno-shstk
+
+CC := gcc
+
+all: build_bof build_format
+
+build_bof:
+	$(CC) -m32 buffer_overflow.c -o bin/bof_32.elf $(build_flags)
+	$(CC) -m32 buffer_overflow.c -o bin/bof_nx_32.elf $(build_NX_flags)
+	$(CC) -m32 buffer_overflow.c -o bin/bof_win_32.elf -Dwin_func $(build_flags)
+
+build_format:
+	$(CC) -O0 -m32 -fno-stack-protector -o bin/read_stack_32.elf \
+		format_string.c -DEASY $(build_flags)
+	$(CC) -O0 -m32 -fno-stack-protector -o bin/format_pc_write_32.elf \
+		format_string.c -DMEDIUM $(build_flags) -z relro
+	$(CC) -O0 -m32 -fno-stack-protector -o bin/format_write_and_constrain_32.elf \
+		format_string.c -DHARD $(build_flags)
+
+clean:
+	rm -rf bin/*

zeratool_lib/zeratool/__init__.py tests/zeratool_binaries/bin/.gitkeepzeratool_lib/zeratool/__init__.py renamed to tests/zeratool_binaries/bin/.gitkeep

@@ -0,0 +1,75 @@
+// Zeratool's tests/buffer_overflow.c
+
+#include <stdio.h>
+
+#ifdef win_func
+void get_secret()
+{
+    puts("flag{you_did_it}");
+}
+#endif
+
+#ifdef srop_func
+__attribute__((naked)) void syscall_gad()
+{
+    __asm__("syscall");
+    __asm__("ret");
+}
+__attribute__((naked)) void pop_rax()
+{
+    __asm__("pop %rax");
+    __asm__("ret");
+}
+__attribute__((naked)) void pop_rdi()
+{
+    __asm__("pop %rdi");
+    __asm__("ret");
+}
+#endif
+
+#ifdef dlresolve_read_func
+
+void give_gadgets()
+{
+    __asm__("pop %rdx");
+    __asm__("ret");
+}
+
+int pwn_me()
+{
+    char my_buf[20] = {'\x00'};
+    printf("Your buffer is at %p\n", my_buf);
+    read(0, my_buf, 230);
+    return 0;
+}
+
+#else
+
+int pwn_me()
+{
+    char my_buf[20] = {'\x00'};
+    printf("Your buffer is at %p\n", my_buf);
+    gets(my_buf);
+    return 0;
+}
+
+#endif
+
+void does_nothing()
+{
+    puts("/bin/sh");
+    execve(NULL, NULL, NULL);
+}
+
+__attribute__((constructor)) void ignore_me()
+{
+    setbuf(stdout, NULL);
+    setbuf(stdin, NULL);
+}
+
+void main()
+{
+
+    puts("pwn_me:");
+    pwn_me();
+}

tests/zeratool_binaries/buffer_overflow.c

@@ -0,0 +1,64 @@
+// Zeratool's tests/format_string.c
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+/*
+ * compiled with:
+ * gcc -O0 -fno-stack-protector -o hard -z execstack -z norelro hard.c
+ * run with:
+ * socat TCP4-LISTEN:7803,tcpwrap=script,reuseaddr,fork EXEC:./hard
+ */
+
+#define SECRET "---BEGIN PRIVATE KEY---"
+
+#ifdef MEDIUM
+/*
+ * Test for point to win
+ */
+void secret_function(void)
+{
+    puts(SECRET);
+}
+#endif
+
+int main(int argc, char *argv[])
+{
+    int i = 0;
+
+#ifdef EASY
+    char buf[1024];
+    fgets(buf, 1024, stdin);
+    /*
+     * Test for stack reading
+     */
+    char key[64] = {};
+    strcpy(key, SECRET);
+#endif
+
+#ifdef MEDIUM
+    char buf[256];
+    read(0, buf, 256);
+#endif
+
+#ifdef HARD
+    char buf[1024];
+    /* read user input securely */
+    fgets(buf, 1024, stdin);
+    /*
+     * Test for point to shellcode AND
+     * satisfy constraints
+     */
+    /* convert string to lowercase */
+    for (i = 0; i < strlen(buf); i++)
+        if (buf[i] >= 'A' && buf[i] <= 'Z')
+            buf[i] = buf[i] ^ 0x20;
+#endif
+
+    /* print out our nice and new lowercase string */
+    printf(buf);
+
+    exit(EXIT_SUCCESS);
+    return EXIT_FAILURE;
+}

tests/zeratool_binaries/format_string.c

@@ -8,13 +8,11 @@
 import claripy
 import timeout_decorator
 from angr import sim_options as so
+from overflowExploiter import findShellcode, getRegValues
 from pwn import *
-
+from simgr_helper import getShellcode
 from zeratool import printf_model
 
-from .overflowExploiter import findShellcode, getRegValues
-from .simgr_helper import getShellcode
-
 log = logging.getLogger(__name__)
 
 
@@ -93,7 +91,9 @@ def exploitFormat(binary_name, properties, leak_format):
         log.info("[+] Binary does not have NX")
         log.info("[+] Overwriting GOT entry to point to shellcode")
 
-        return rediscoverAndExploit(binary_name, properties, stack_position, leak_format)
+        return rediscoverAndExploit(
+            binary_name, properties, stack_position, leak_format
+        )
 
 
 """
@@ -107,7 +107,7 @@ def rediscoverAndExploit(binary_name, properties, stack_position, leak_format):
 
     properties["shellcode"] = getShellcode(properties)
     properties["stack_position"] = stack_position
-    properties["leak_format" = leak_format]
+    properties["leak_format"] = leak_format
     inputType = properties["input_type"]
 
     extras = {so.REVERSE_MEMORY_NAME_MAP, so.TRACK_ACTION_HISTORY, so.TRACK_CONSTRAINTS}
@@ -176,12 +176,12 @@ def exploreBinary(simgr):
         stdin_str = str(end_state.posix.dumps(0))
         log.info("[+] Triggerable with STDIN : {}".format(stdin_str))
 
-        return stdin_str, end_state["fmt_outcome"]
+        return stdin_str, end_state.globals["fmt_outcome"]
     elif inputType == "ARG" and end_state is not None:
         arg_str = str(end_state.solver.eval(arg, cast_to=str))
         run_environ["input"] = arg_str
 
-        return arg_str, end_state["fmt_outcome"]
+        return arg_str, end_state.globals["fmt_outcome"]
 
 
 def get_num_constraints(chop_byte, state):
@@ -317,6 +317,8 @@ def checkExploitable(self, fmt):
         if greatest_count > 0:
             shellcode = properties["shellcode"]
             stack_pos = properties["stack_position"]
+            state_copy = None
+            results_n = None
 
             for got_name, got_addr in list(properties["protections"]["got"].items()):
                 #                for got_name,got_addr in [(x,y) for (x,y) in properties['protections']['got'].items() if x in " exit"]: #debug for hard_format
@@ -376,8 +378,8 @@ def checkExploitable(self, fmt):
                 exploit_results = {}
                 if results["success"] == True:
                     exploit_results["success"] = results["success"]
-                    exploit_results["input"] = format_input
-                    self.state["fmt_outcome"] = results["fmt_outcome"]
+                    exploit_results["input"] = user_input
+                    self.state.globals["fmt_outcome"] = results["fmt_outcome"]
                 else:  # Maybe angr still messed up the pointer
                     log.info("[-] Payload launch failed. Fixing angr stack pointer")
 
@@ -464,10 +466,7 @@ def checkExploitable(self, fmt):
 
                         leak_format = properties["leak_format"]
                         results_n = sendExploit(
-                            binary_name,
-                            properties,
-                            user_input,
-                            leak_format
+                            binary_name, properties, user_input, leak_format
                         )
                         if results_n["success"]:
                             log.info(
@@ -476,14 +475,19 @@ def checkExploitable(self, fmt):
                             self.state.globals["type"] = "Format"
                             self.state.globals["position"] = position
                             self.state.globals["length"] = greatest_count
-                            self.state["fmt_outcome"] = results_n["fmt_outcome"]
+                            self.state.globals["fmt_outcome"] = results_n["fmt_outcome"]
                             return True
 
                             # exploit_results['success'] = results_n['success']
                             # exploit_results['input'] = format_input
 
                 # Verify solution
-                if state_copy.globals["inputType"] == "STDIN" and results_n["success"]:
+                if (
+                    state_copy
+                    and state_copy.globals["inputType"] == "STDIN"
+                    and results_n
+                    and results_n["success"]
+                ):
                     stdin_str = str(state_copy.posix.dumps(0))
                     if format_payload in stdin_str or results["success"]:
                         var_value = self.state.memory.load(var_loc, var_len)
@@ -499,10 +503,10 @@ def checkExploitable(self, fmt):
                         self.state.globals["type"] = "Format"
                         self.state.globals["position"] = position
                         self.state.globals["length"] = greatest_count
-                        self.state["fmt_outcome"] = results_n["fmt_outcome"]
+                        self.state.globals["fmt_outcome"] = results_n["fmt_outcome"]
 
                         return True
-                if state_copy.globals["inputType"] == "ARG":
+                if state_copy and state_copy.globals["inputType"] == "ARG":
                     arg = state.globals["arg"]
                     arg_str = str(state_copy.solver.eval(arg, cast_to=str))
                     if format_payload in arg_str:
@@ -519,7 +523,7 @@ def checkExploitable(self, fmt):
                         self.state.globals["type"] = "Format"
                         self.state.globals["position"] = position
                         self.state.globals["length"] = greatest_count
-                        self.state["fmt_outcome"] = results_n["fmt_outcome"]
+                        self.state.globals["fmt_outcome"] = results_n["fmt_outcome"]
                         return True
                 state_copy = backup_state.copy()
 
@@ -651,9 +655,7 @@ def sendExploit(binary_name, properties, input_string, leak_format):
     # Flag not in stdout, we have a shell
     else:
 
-        if (
-            properties["input_type"] == "STDIN"
-        ):
+        if properties["input_type"] == "STDIN":
             proc = process(binary_name)
             proc.sendline(input_string)
         else: