#include "linklist.h"

struct api_s *add_api(struct api_s *ap, DWORD api_addr, DWORD rdata_addr, DWORD txt_addr, DWORD is_jmp_stx)
{
	struct api_s *new_ap = NULL;
	struct api_s *cur = NULL;

	if (ap)
	{
		if ((cur = find_api(ap, api_addr)) != NULL)
		{
			cur->rdata = add_rdata(cur->rdata, rdata_addr, txt_addr, is_jmp_stx);
			return ap;
		}
	}
	new_ap = (struct api_s*)malloc(sizeof (struct api_s));
	new_ap->next = ap;
	new_ap->api_addr = api_addr;
	new_ap->rdata = add_rdata(new_ap->rdata, rdata_addr, txt_addr, is_jmp_stx);
	return new_ap;
}

struct rdata_s *add_rdata(struct rdata_s *rd, DWORD rdata_addr, DWORD txt_addr, DWORD is_jmp_stx)
{
	struct rdata_s *new_rdata = NULL;

	new_rdata = (struct rdata_s*)malloc(sizeof (struct rdata_s));
	new_rdata->next = rd;
	new_rdata->rdata_addr = rdata_addr;
	new_rdata->txt_addr = txt_addr;
	new_rdata->is_jmp_stx = is_jmp_stx;
	return new_rdata;
}

void fix_api_rdata(struct api_s *ap)
{
	DWORD OldProtect;
	struct rdata_s *rdata_actual = NULL;

	while (ap)
	{
		rdata_actual = ap->rdata;
		VirtualProtect((LPVOID)rdata_actual->rdata_addr, 4, PAGE_EXECUTE_READWRITE, &OldProtect);
		*(DWORD*)rdata_actual->rdata_addr = ap->api_addr;
		VirtualProtect((LPVOID)rdata_actual->rdata_addr, 4, OldProtect, &OldProtect);
		while (rdata_actual)
		{
			if (rdata_actual->is_jmp_stx)
			{
				VirtualProtect((LPVOID)rdata_actual->txt_addr, 6, PAGE_EXECUTE_READWRITE, &OldProtect);
				*(BYTE*)(rdata_actual->txt_addr) = 0xFF;
				*(BYTE*)(rdata_actual->txt_addr + 1) = 0x15;
				*(DWORD*)(rdata_actual->txt_addr + 2) = rdata_actual->rdata_addr;
				VirtualProtect((LPVOID)rdata_actual->txt_addr, 6, OldProtect, &OldProtect);
			}
			else
			{
				VirtualProtect((LPVOID)rdata_actual->txt_addr, 4, PAGE_EXECUTE_READWRITE, &OldProtect);
				*(DWORD*)rdata_actual->txt_addr = rdata_actual->rdata_addr;
				VirtualProtect((LPVOID)rdata_actual->txt_addr, 4, OldProtect, &OldProtect);
			}
			rdata_actual = rdata_actual->next;
		}
		ap = ap->next;
	}
}

void reorder_api_rdata(struct api_s *ap)
{
	struct rdata_s *rdata_free = NULL;
	struct rdata_s *rdata_used = NULL;
	struct rdata_s *rdata_actual = NULL;
	DWORD actual_rdata_addr = 0;

	while (ap)
	{
		if (ap->rdata == NULL)
		{
			MessageBoxA(NULL, "There is something wrong!", "!?", 0);
			exit(0);
		}
		actual_rdata_addr = ap->rdata->rdata_addr;
		if (find_rdata(rdata_used, actual_rdata_addr))
		{
			if (rdata_free == NULL)
			{
				MessageBoxA(NULL, "FIX IT", "!?", 0);
				exit(0);
			}
			rdata_free = get_rdata(rdata_free, &actual_rdata_addr);
		}
		rdata_used = add_rdata(rdata_used, actual_rdata_addr, 0, 0);
		rdata_actual = ap->rdata;
		while (rdata_actual)
		{
			if (rdata_actual->rdata_addr != actual_rdata_addr)
			{
				if (find_rdata(rdata_used, actual_rdata_addr) == NULL)
					rdata_free = add_rdata(rdata_free, rdata_actual->rdata_addr, 0, 0);
				rdata_actual->rdata_addr = actual_rdata_addr;
			}
			rdata_actual = rdata_actual->next;
		}
		ap = ap->next;
	}
}

struct rdata_s *get_rdata(struct rdata_s *rd, DWORD *rdata_addr)
{
	struct rdata_s *rd_actual;

	if (rd)
	{
		*rdata_addr = rd->rdata_addr;
		rd_actual = rd;
		rd = rd->next;
		free(rd_actual);
		return rd;
	}
	return NULL;
}

struct api_s *find_api(struct api_s *ap, DWORD api_addr)
{
	while (ap)
	{
		if (ap->api_addr == api_addr)
			return ap;
		ap = ap->next;
	}
	return NULL;
}

struct rdata_s *find_rdata(struct rdata_s *rd, DWORD rdata_addr)
{
	while (rd)
	{
		if (rd->rdata_addr == rdata_addr)
			return rd;
		rd = rd->next;
	}
	return NULL;
}

struct rdata_s *find_txt(struct rdata_s *rd, DWORD txt_addr)
{
	while (rd)
	{
		if (rd->txt_addr == txt_addr)
			return rd;
		rd = rd->next;
	}
	return NULL;
}

void print_api(FILE *fp, struct api_s *ap)
{
	while (ap)
	{
		fprintf(fp, "API (0x%X) has ", ap->api_addr);
		print_rdata(fp, ap->rdata);
		fprintf(fp, "\n");
		ap = ap->next;
	}
}

void print_rdata(FILE *fp, struct rdata_s *rd)
{
	while (rd)
	{
		fprintf(fp, "rdata.0x%X (txt.0x%X)  ", rd->rdata_addr, rd->txt_addr);
		rd = rd->next;
	}
}