#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/ioctl.h>
#include <sys/ptrace.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/user.h>
#include <sys/wait.h>

#include <linux/kdev_t.h>
#include <linux/major.h>

#include "neercs.h"

static int memcpy_from_target(pid_t pid, void* dest, void* src, size_t n)
{
    unsigned int i;
    long *d, *s;
    d = (long*) dest;
    s = (long*) src;
    n /= sizeof(long);
    for (i = 0; i < n; i++) {
	d[i] = ptrace(PTRACE_PEEKTEXT, pid, s+i, 0);
	if (errno) {
	    perror("ptrace(PTRACE_PEEKTEXT)");
	    return -1;
	}
    }
    return 0;
}

static int memcpy_into_target(pid_t pid, void* dest, void* src, size_t n)
{
    unsigned int i;
    long *d, *s;
    d = (long*) dest;
    s = (long*) src;
    for (i = 0; i < n / sizeof(long); i++) {
	if (ptrace(PTRACE_POKETEXT, pid, d+i, s[i]) == -1) {
	    perror("ptrace(PTRACE_POKETEXT)");
	    return -1;
	}
    }
    return 0;
}

static int do_syscall(pid_t pid, struct user_regs_struct *regs)
{

    struct user_regs_struct oregs;
    long oinst;

    if(ptrace(PTRACE_GETREGS, pid, NULL, &oregs) < 0)
    {
	return errno;
    }

    regs->esp = oregs.esp-4;

    oinst = ptrace(PTRACE_PEEKTEXT, pid, regs->esp, 0);
    if (errno) {
        fprintf(stderr, "PTRACE_PEEKTEXT failed\n");
	return errno;
    }

    if (ptrace(PTRACE_POKETEXT, pid, regs->esp, 0x80cd) < 0) { /* int 0x80 */
        fprintf(stderr, "PTRACE_POKETEXT failed\n");
	return errno;
    }

    regs->eip = regs->esp;
    if(ptrace(PTRACE_SETREGS, pid, NULL, regs) < 0)
    {
	return errno;
    }

    if(ptrace(PTRACE_SINGLESTEP, pid, NULL, NULL) < 0)
    {
	return errno;
    }
    waitpid(pid, NULL, 0);

    if(ptrace(PTRACE_GETREGS, pid, NULL, regs) < 0)
    {
	return errno;
    }

    if(ptrace(PTRACE_SETREGS, pid, NULL, &oregs) < 0)
    {
        fprintf(stderr, "PTRACE_SETREGS failed\n");
	return errno;
    }

    if(ptrace(PTRACE_POKETEXT, pid, oregs.esp-4 , oinst) < 0)
    {
        fprintf(stderr, "PTRACE_POKETEXT failed\n");
	return errno;
    }

    return 0;
}

static int do_close(pid_t pid, int fd)
{
    struct user_regs_struct regs;
    if(ptrace(PTRACE_GETREGS, pid, NULL, &regs) < 0)
    {
        fprintf(stderr, "PTRACE_GETREGS failed\n");
        return errno;
    }

    regs.eax = SYS_close;
    regs.ebx = fd;

    return do_syscall(pid, &regs);
}

static int do_dup2(pid_t pid, int oldfd, int newfd)
{
    struct user_regs_struct regs;
    if(ptrace(PTRACE_GETREGS, pid, NULL, &regs) < 0)
    {
        fprintf(stderr, "PTRACE_GETREGS failed\n");
        return errno;
    }

    regs.eax = SYS_dup2;
    regs.ebx = oldfd;
    regs.ecx = newfd;

    return do_syscall(pid, &regs);
}

static int do_open(pid_t pid, char *path, int mode, int *fd)
{
    struct user_regs_struct regs;
    void *backup_page;
    void *target_page = (void*)(0x08048000);
    size_t size = (1<<12); /* 4K */
    int ret;

    /* Backup the page that we will use */
    backup_page = malloc(size);
    if (memcpy_from_target(pid, backup_page, target_page, size) < 0)
        return -1;

    if(ptrace(PTRACE_GETREGS, pid, NULL, &regs) < 0)
    {
        fprintf(stderr, "PTRACE_GETREGS failed\n");
        return errno;
    }

    /* +4 because it's trcuncated on a multiple of 4 and we need 1 */
    memcpy_into_target(pid, target_page, path, strlen(path)+4);

    regs.eax = SYS_open;
    regs.ebx = (long)target_page;
    regs.ecx = O_RDWR;
    regs.edx = 0755;

    if((ret = do_syscall(pid, &regs)) != 0)
    {
        return ret;
    }

    /* Restore the pages */
    memcpy_into_target(pid, target_page, backup_page, size);

    *fd  = regs.eax;

    return 0;
}

static int do_setsid(pid_t pid) {
    struct user_regs_struct regs;
    int ret;

    if(ptrace(PTRACE_GETREGS, pid, NULL, &regs) < 0)
    {
        fprintf(stderr, "PTRACE_GETREGS failed\n");
        return errno;
    }

    debug("Running setsid on process %d (sid=%d)", pid, getsid(pid));

    regs.eax = SYS_setpgid;
    regs.ebx = 0;
    regs.ecx = getsid(pid);

    if((ret = do_syscall(pid, &regs)) != 0)
    {
        return ret;
    }

    if (regs.eax != 0) {
        fprintf(stderr, "setpgid failed\n");
        return -regs.eax;
    }

    regs.eax = SYS_setsid;

    if((ret = do_syscall(pid, &regs)) != 0)
    {
        return ret;
    }

    debug("pid %d has now sid %d", pid, getsid(pid));

    if (regs.eax == -1) {
        fprintf(stderr, "getsid failed\n");
        return -regs.eax;
    }

    return 0;
}

int grab_process(pid_t pid, char *ptyname, int ptyfd)
{
    int i, fd = 0, mode, ret;
    char fdstr[1024];
    char to_open[3];
    struct stat stat_buf;

    debug("pty is %s", ptyname);

    kill(pid, SIGSTOP);
    if(ptrace(PTRACE_ATTACH, pid, 0, 0) < 0)
    {
        fprintf(stderr, "Cannot access process %d\n", pid);
        return errno;
    }
    kill(pid, SIGCONT);
    waitpid(pid, NULL, 0);

    for(i=0; i<=2; i++)
    {
	snprintf(fdstr, sizeof(fdstr), "/proc/%d/fd/%d", pid, i);
        to_open[i]=0;
	lstat(fdstr, &stat_buf);
	if ((stat_buf.st_mode & S_IRUSR) && (stat_buf.st_mode & S_IWUSR))
	    mode = O_RDWR;
	else if (stat_buf.st_mode & S_IWUSR)
	    mode = O_WRONLY;
	else
	    mode = O_RDONLY;

	if (stat(fdstr, &stat_buf) < 0)
	    continue;

	if (!S_ISCHR(stat_buf.st_mode) || MAJOR(stat_buf.st_rdev) != UNIX98_PTY_SLAVE_MAJOR)
            continue;

        debug("fd %d is a pty", i);
        if((ret = do_close(pid, i)))
        {
            fprintf(stderr, "do_close %s\n", strerror(ret));
        }
        to_open[i]=1;
    }

    if((ret = do_setsid(pid)))
    {
        fprintf(stderr, "do_setsid %s\n", strerror(ret));
    }

    for(i=0; i<=2; i++)
    {
        if(!to_open[i])
            continue;
        if((ret = do_open(pid, ptyname, mode, &fd)))
        {
            fprintf(stderr, "do_open %s\n", strerror(ret));
        }
        if((ret = do_dup2(pid, fd, i)))
        {
            fprintf(stderr, "do_dup2 %s\n", strerror(ret));
        }
    }

    kill(pid, SIGWINCH);
    ptrace(PTRACE_DETACH, pid, 0, 0);

    return 0;
}