#include <spslr.h>
#include <stdio.h>
#include <stdlib.h>

#include <spslr_program.h>

// TODO
static int spslr_target(uint32_t uid, uint32_t size, uint32_t fieldcnt) { return 0; }
static int spslr_field(uint32_t offset, uint32_t size, uint32_t flags) { return 0; }
static int spslr_randomize(uint32_t target) { return 0; }
static int spslr_mprot(void* base, uint32_t pagecnt, uint8_t perm) { return 0; }
static int spslr_ipatch(void* ptr, uint32_t target, uint32_t field) { return 0; }
static int spslr_dpatch(void* ptr, uint32_t target) { return 0; }

/*
Postprocessing tool patches the value of __spslr_program to point to the SPSLR program section.
With ASLR, there are 2 options to make it function correctly:
	1. Make sure __spslr_program is relocated with program image shift (preferred)
	2. __spslr_program = spslr_ptr_absolute((uint64_t)__spslr_program) and postprocessor inserts relative value (kinda weird)
*/
static const uint8_t* __spslr_program = NULL;

static void* spslr_ptr_absolute(uint64_t relative) {
	// To allow ASLR, all patcher addresses are relative to &__spslr_program
	return ((uint8_t*)&__spslr_program) + relative;
}

static int spslr_do(const struct SPSLR_INST* inst) {
	if (!inst)
		return -1;

	static uint32_t pending_fields = 0;

	if (pending_fields) {
		if (inst->opcode != SPSLR_FIELD) {
			fprintf(stderr, "spslr_do encountered non field instruction where a field instruction was expected\n");
			return -1;
		}

		pending_fields--;
	}

	switch (inst->opcode) {
	case SPSLR_TARGET:
		pending_fields = inst->op2.target_fieldcnt;
		return spslr_target(inst->op0.target_uid, inst->op1.target_size, inst->op2.target_fieldcnt);
	case SPSLR_FIELD:
		return spslr_field(inst->op0.field_offset, inst->op1.field_size, inst->op2.field_flags);
	case SPSLR_RANDOMIZE:
		return spslr_randomize(inst->op0.randomize_target);
	case SPSLR_MPROT:
		return spslr_mprot(spslr_ptr_absolute(inst->op0.mprot_ptr), inst->op1.mprot_pagecnt, inst->op2.mprot_perm);
	case SPSLR_IPATCH:
		return spslr_ipatch(spslr_ptr_absolute(inst->op0.ipatch_ptr), inst->op1.ipatch_target, inst->op2.ipatch_field);
	case SPSLR_DPATCH:
		return spslr_dpatch(spslr_ptr_absolute(inst->op0.dpatch_ptr), inst->op1.dpatch_target);
	default:
		return -1;
	}
}

void spslr_selfpatch() {
	if (!__spslr_program) {
		fprintf(stderr, "spslr_selfpatch has no patcher program (finalize this binary to use SPSLR)\n");
		return;
	}

	int sz;
	struct SPSLR_INST inst;
	while ((sz = spslr_inst_load(&inst, __spslr_program)) > 0) {
		if (inst.opcode == SPSLR_EXIT)
			return;

		if (spslr_do(&inst) < 0) {
			fprintf(stderr, "spslr_selfpatch failed to execute an instruction\n");
			exit(1);
		}

		__spslr_program += sz;
	}

	fprintf(stderr, "spslr_selfpatch encountered invalid instruction\n");
	exit(1);
}