/*
 *
 * Portions of code from PspLink:
 * Copyright (c) 2005 James F <tyranid@gmail.com>
 *
 */
#include <pspkernel.h>
#include <pspdebug.h>
#include <pspsysmem_kernel.h>
#include <psputilsforkernel.h>
#include <pspmoduleexport.h>
#include <psploadcore.h>
#include <pspthreadman_kernel.h>
#include <pspsdk.h>
#include <stdio.h>
#include <string.h>
#include "apihook.h"

#define DEBUG_APIHOOK 0

void sioPrintWord(u32 value);
void sioPrint(const char* str);
void sioPrintHex(char ch);

typedef struct HookRecord
{
 char modname[32];
 char libname[32];
 u32 nid;
 void *func;
 void *oldfunc;
 void *syscall;
} HookRecord;

#define MAX_HOOK_RECORDS 30

HookRecord hookRecords[30];
int        numHookRecords = 0;
SceUID     hookSema = 0;

struct SyscallHeader
{
	void *unk;
	unsigned int basenum;
	unsigned int topnum;
	unsigned int size;
};

void initApiHooks(void)
{
	hookSema = sceKernelCreateSema("pikey_hooks", 0, 1, 1, NULL);
	numHookRecords = 0;
}

void *find_syscall_addr(u32 addr)
{
	struct SyscallHeader *head;
	u32 *syscalls;
	void **ptr;
	int size;
	int i;

	asm(
			"cfc0 %0, $12\n"
			: "=r"(ptr)
	   );

	if(!ptr)
	{
		sioPrint("Null ptr from cfc\n");
		return NULL;
	}

	head = (struct SyscallHeader *) *ptr;
	syscalls = (u32*) (*ptr + 0x10);
	size = (head->size - 0x10) / sizeof(u32);

	for(i = 0; i < size; i++)
	{
		if(syscalls[i] == addr)
		{
			return &syscalls[i];
		}
	}

	sioPrint("Failed to find syscall for addr\n");

	return NULL;
}

static void *apiHookAddr(u32 *addr, void *func)
{
	int intc;

	if(!addr)
	{
		return NULL;
	}

	intc = pspSdkDisableInterrupts();
	*addr = (u32) func;
	sceKernelDcacheWritebackInvalidateRange(addr, sizeof(addr));
	sceKernelIcacheInvalidateRange(addr, sizeof(addr));
	pspSdkEnableInterrupts(intc);

	return addr;
}

static struct SceLibraryEntryTable *_libsFindLibrary(SceUID uid, const char *library)
{
	struct SceLibraryEntryTable *entry;
	SceModule *pMod;
	void *entTab;
	int entLen;

	pMod = sceKernelFindModuleByUID(uid);
	if(pMod != NULL)
	{
		int i = 0;

		entTab = pMod->ent_top;
		entLen = pMod->ent_size;
		while(i < entLen)
		{
			entry = (struct SceLibraryEntryTable *) (entTab + i);

			if((entry->libname) && (strcmp(entry->libname, library) == 0))
			{
				return entry;
			}
			else if(!entry->libname && !library)
			{
				return entry;
			}

			i += (entry->len * 4);
		}

		sioPrint("Can't find libname in module\n");
	}
	else
	{
		sioPrint("Can't find module UID\n");
	}

	return NULL;
}

void* libsFindExportAddrByNid(SceUID uid, const char *library, u32 nid)
{
	struct SceLibraryEntryTable *entry;
	u32 *addr = NULL;

	entry = _libsFindLibrary(uid, library);
	if(entry)
	{
		int count;
		int total;
		unsigned int *vars;

		total = entry->stubcount + entry->vstubcount;
		vars = entry->entrytable;

		if(entry->stubcount > 0)
		{
			for(count = 0; count < entry->stubcount; count++)
			{
				if(vars[count] == nid)
				{
					return &vars[count+total];
				}
			}
		}
		else
		{
			sioPrint("No stubs for entry\n");
		}
	}																	

	return addr;
}

u32 libsFindExportByNid(SceUID uid, const char *library, u32 nid)
{
	u32 *addr;

	addr = libsFindExportAddrByNid(uid, library, nid);
	if(!addr)
	{
		return 0;
	}

	return *addr;
}

u32 apiHookByNid(const char *modname, const char *library, u32 nid, void *func)
{
	sceKernelWaitSema(hookSema, 1, 0);

	if (numHookRecords == MAX_HOOK_RECORDS)
	{
		sceKernelSignalSema(hookSema, 1);
		return -1;
	}

	u32 lhookid = -1;

  /***************************************************************************/
	/* First, let's see if we already have a hook for this function.           */
	/***************************************************************************/
	int ii;
	for (ii = 0; ii < numHookRecords; ii++)
	{
		if ((hookRecords[ii].nid == nid) &&
				(strcmp(hookRecords[ii].libname, library) == 0) &&
				(strcmp(hookRecords[ii].modname, modname) == 0))
		{
			lhookid = ii;
		}
	}

	if (lhookid == -1)
	{
		// no existing hook, so this is simple.
		SceModule *pMod = sceKernelFindModuleByName(modname);
		if (pMod)
		{
			u32 addr = libsFindExportByNid(pMod->modid, library, nid);
			if(addr)
			{
				void * syscalladdr = find_syscall_addr(addr);
				if(!apiHookAddr(syscalladdr, func))
				{
					sioPrint("Failed to apiHookAddr\n");
				}
				else
				{
					// success hooking.  Record the details.
					lhookid = numHookRecords++;
					hookRecords[lhookid].oldfunc = (void*)addr;
					strcpy(hookRecords[lhookid].libname, library);
					strcpy(hookRecords[lhookid].modname, modname);
					hookRecords[lhookid].nid     = nid;
					hookRecords[lhookid].func    = func;
					hookRecords[lhookid].syscall = syscalladdr;
				}
			}
			else
			{
				sioPrint("Failed to find exportbynid\n");
			}
		}
	}
	else
	{
		// We have an existing hook for this function.
		// So, we:
		//   - copy the existing hook details
		//   - set 'oldfunc' of the existing hook to the new 'func'
		// so that the call tree next time the SCE func is called will be:
		//   - first hook
		//     - this hook
		//       - SCE func
		// THIS ASSUMES THAT ALL HOOK FUNCS WILL CALL THROUGH TO THE ORIG FUNC

#if DEBUG_APIHOOK==1
		sioPrint("Copy hook record from number ");
		sioPrintHex(lhookid);
		sioPrint("\r\n");
#endif

		int lnewhookid = numHookRecords++;
		memcpy(&(hookRecords[lnewhookid]), &(hookRecords[lhookid]), sizeof(HookRecord));
		hookRecords[lhookid].oldfunc = func;
		hookRecords[lnewhookid].syscall = 0;

		lhookid = lnewhookid;
	}

	sceKernelSignalSema(hookSema, 1);

#if DEBUG_APIHOOK==1
	sioPrint("Return hookid ");
	sioPrintHex(lhookid);
	sioPrint("\r\n");
#endif

	return lhookid;
}

typedef u32 (*HOOKFUNC)(u32, u32, u32, u32, u32, u32, u32);

u32 callApiHook(u32 hookid, ...)
{
	int lreturn;
	u32 args[7];
	int ii;
	va_list val;

	// save the argument list
	va_start(val, hookid);
	for (ii = 0; ii < 7; ii++)
	{
		args[ii] = va_arg(val, u32);
	}
	va_end(val);

#if DEBUG_APIHOOK==1
	sioPrint("call api hook #");
	sioPrintHex(hookid);
	sioPrint("\r\n");
#endif

	sceKernelWaitSema(hookSema, 1, 0);

	// find the hook
	if ((hookid >= 0) && (hookid < numHookRecords))
	{
		HOOKFUNC lfunc = (HOOKFUNC) (hookRecords[hookid].oldfunc);

		sceKernelSignalSema(hookSema, 1);

		if (lfunc)
		{
#if DEBUG_APIHOOK==1
			sioPrint("Calling with args: \r\n");
			for (ii = 0; ii < 2; ii++)
			{
				sioPrintWord(args[ii]);
				sioPrint("\r\n");
			}
#endif

			lreturn = lfunc(args[0], args[1], args[2], args[3], args[4], args[5], args[6]);

#if DEBUG_APIHOOK==1
			sioPrint("Return: ");
			sioPrintWord(lreturn);
			sioPrint("\r\n");
#endif
		}
		else
		{
			lreturn = -2;
		}
	}
	else
	{
		sceKernelSignalSema(hookSema, 1);

		lreturn = -1;
	}

	return lreturn;
}

void unHookApi(u32 hookid)
{
	sceKernelWaitSema(hookSema, 1, 0);

	// find the hook
	if ((hookid >= 0) && (hookid < numHookRecords))
	{

		/* Restore original function */
		if(hookRecords[hookid].oldfunc)
		{
			u32 *lsyscall = hookRecords[hookid].syscall;

			if (lsyscall)
			{
				*lsyscall = (u32) hookRecords[hookid].oldfunc;
				hookRecords[hookid].oldfunc = NULL;
				sceKernelDcacheWritebackInvalidateRange(lsyscall, sizeof(void *));
				sceKernelIcacheInvalidateRange(lsyscall, sizeof(void *));
			}
		}
	}

	sceKernelSignalSema(hookSema, 1);
}
