#include <stdio.h>
#include <sys/ptrace.h>
#include <sys/user.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>

#ifndef __WALL
#define __WALL 0
#endif

void enable_core_dump(void);

int inject(pid_t pid, const char *shellcode, int size)
{
        long ptr;
        int i;
        struct user_regs_struct data;

        if (ptrace(PTRACE_ATTACH, pid, NULL, NULL) == -1) {
                perror("Attach");
                return -1;
        }
        /* wait for the stopping of the target process */
        if (waitpid(pid, NULL, __WALL) == -1) {
                perror("waitpid");
                return -1;
        }

        if (ptrace(PTRACE_GETREGS, pid, NULL, &data) == -1) {
                perror("Getregs");
                return -1;
        }

        /* save the return address, since we use jmp to call the
         * function instead of call instruction */
        data.esp -= sizeof(long);
        if (ptrace(PTRACE_POKETEXT, pid, data.esp, data.eip) == -1) {
                perror("Poketext");
                return -1;
        }

        /* transfer the shellcode to the target process */
        if (size < 0)
                size = strlen(shellcode);
        ptr = data.eip = data.esp - size - 1024;
        for (i = 0; i < size; i += sizeof(long)) {
                if (ptrace(PTRACE_POKETEXT, pid, ptr, *((long*)(shellcode + i))) == -1) {
                        perror("Poktext");
                        return -1;
                }
                ptr += sizeof(long);
        }

        /* set the instruction counter */
        data.eip += 2; /* skip the two instructions: nop */
        if (ptrace(PTRACE_SETREGS, pid, NULL, &data) == -1) {
                perror("Setregs");
                return -1;
        }

        /* detach the target process and let it run... */
        if (ptrace(PTRACE_DETACH, pid, NULL, NULL) == -1) {
                perror("Detach");
                return -1;
        }

        return 0;
}

int main(int argc, char *argv[])
{
        pid_t pid;

        if (argc < 2) {
                printf("Usage: %s pid [-k]\n", argv[0]);
                return -1;
        }
        pid = atoi(argv[1]);

        printf("Start injecting(%d)...", pid);
        if (inject(pid, (const char *)enable_core_dump, 0x2c) != 0) {
                printf("Failed\n");
                return -1;
        }
        printf("OK\n");

        /* sleep for a moment. When waken up, the core dump of the target
         * process should be enabled. */
        if (argc > 2 && strcmp(argv[2], "-k") == 0) {
                sleep(3);
                kill(pid, SIGSEGV);
        }

        return 0;
}

/*
 * call the function: setrlimit(RLIMIT_CORE, {-1, -1});
 * after calling that, the process is allowed to save
 * the core dump file if exception, such as SIGSEGV, occurs.
 */

/* XXX:
 * Something important to do before the exploitation is to put two nops bytes
 * before the shellcode. Reason is simple : if ptrace has interrupted a syscall
 * being executed, the kernel will subtract two bytes from eip after the
 * PTRACE_DETACH to restart the syscall.
 */

.globl enable_core_dump

enable_core_dump:
        nop
        nop
        pusha
        push %ebp
        mov %esp, %ebp
        sub $8, %esp
        movl $-1, -8(%ebp)
        movl $-1, -4(%ebp)
        lea -8(%ebp), %eax
        mov %eax, %ecx
        xor %ebx, %ebx
        mov $4, %bl
        xor %eax, %eax
        mov $75, %al
        int $0x80
        leave
        popa
        ret
        nop /* padding byte */
        nop
        nop