#include "precomp.h"
#include "./../sys/session.h"
#include "./../sys/process/ps.h"

typedef struct _DESKTOPLIST
{
	DWORD dwSessionId;
	char * cpStationName;
	Packet * response;
} DESKTOPLIST, *LPDESKTOPLIST;

/*
 * Callback function for EnumDesktops when listing desktops on a station during desktop_list().
 */
BOOL CALLBACK desktop_enumdesktops_callback(LPTSTR cpDesktopName, LPARAM lpParam)
{
	DESKTOPLIST * dl = NULL;
	Tlv entry[3] = { 0 };
	DWORD dwSessionId = 0;

	do
	{
		dl = (DESKTOPLIST *)lpParam;
		if (!dl)
		{
			break;
		}

		if (!dl->cpStationName || !dl->response || !cpDesktopName)
		{
			break;
		}

		dwSessionId = htonl(dl->dwSessionId);

		entry[0].header.type = TLV_TYPE_DESKTOP_SESSION;
		entry[0].header.length = sizeof(DWORD);
		entry[0].buffer = (PUCHAR)&dwSessionId;

		entry[1].header.type = TLV_TYPE_DESKTOP_STATION;
		entry[1].header.length = (DWORD)(strlen(dl->cpStationName) + 1);
		entry[1].buffer = (PUCHAR)dl->cpStationName;

		entry[2].header.type = TLV_TYPE_DESKTOP_NAME;
		entry[2].header.length = (DWORD)(strlen(cpDesktopName) + 1);
		entry[2].buffer = (PUCHAR)cpDesktopName;

		packet_add_tlv_group(dl->response, TLV_TYPE_DESKTOP, entry, 3);

	} while (0);

	return TRUE;
}

/*
 * Callback function for EnumWindowStations when listing stations during request_ui_desktop_enum().
 */
BOOL CALLBACK desktop_enumstations_callback(LPTSTR cpStationName, LPARAM param)
{
	HWINSTA hWindowStation = NULL;
	DESKTOPLIST dl = { 0 };

	do
	{
		hWindowStation = OpenWindowStation(cpStationName, FALSE, MAXIMUM_ALLOWED); // WINSTA_ALL_ACCESS
		if (!hWindowStation)
		{
			break;
		}

		dl.dwSessionId = session_id(GetCurrentProcessId());
		dl.response = (Packet *)param;
		dl.cpStationName = cpStationName;

		EnumDesktops(hWindowStation, desktop_enumdesktops_callback, (LPARAM)&dl);

	} while (0);

	if (hWindowStation)
	{
		CloseWindowStation(hWindowStation);
	}

	return TRUE;
}

/*
 * Enumerate all accessible desktops on all stations.
 */
DWORD request_ui_desktop_enum(Remote * remote, Packet * request)
{
	DWORD dwResult = ERROR_SUCCESS;
	Packet * response = NULL;

	do
	{
		response = packet_create_response(request);
		if (!response)
		{
			BREAK_WITH_ERROR("[UI] desktop_enum. packet_create_response failed", ERROR_INVALID_HANDLE);
		}

		EnumWindowStations(desktop_enumstations_callback, (LPARAM)response);

	} while (0);

	if (response)
	{
		packet_transmit_response(dwResult, remote, response);
	}

	return ERROR_SUCCESS;
}

/*
 * Get the session/windows station/desktop we are currently using.
 */
DWORD request_ui_desktop_get(Remote * remote, Packet * request)
{
	DWORD dwResult = ERROR_SUCCESS;
	Packet * response = NULL;

	do
	{
		response = packet_create_response(request);
		if (!response)
		{
			BREAK_WITH_ERROR("[UI] desktop_get. packet_create_response failed", ERROR_INVALID_HANDLE);
		}

		lock_acquire(remote->lock);

		packet_add_tlv_uint(response, TLV_TYPE_DESKTOP_SESSION, remote->curr_sess_id);
		packet_add_tlv_string(response, TLV_TYPE_DESKTOP_STATION, remote->curr_station_name);
		packet_add_tlv_string(response, TLV_TYPE_DESKTOP_NAME, remote->curr_desktop_name);

		lock_release(remote->lock);

	} while (0);

	if (response)
	{
		packet_transmit_response(dwResult, remote, response);
	}

	return ERROR_SUCCESS;
}

/*
 * Set this process to use a specified window station and this thread to use
 * a specified desktop.
 */
DWORD request_ui_desktop_set(Remote * remote, Packet * request)
{
	DWORD dwResult = ERROR_SUCCESS;
	Packet * response = NULL;
	char * cpDesktopName = NULL;
	char * cpStationName = NULL;
	HWINSTA hWindowStation = NULL;
	HWINSTA hOrigWindowStation = NULL;
	HDESK hDesktop = NULL;
	BOOL bSwitch = FALSE;
	DWORD dwSessionId = 0;

	do
	{
		response = packet_create_response(request);
		if (!response)
		{
			BREAK_WITH_ERROR("[UI] desktop_set. packet_create_response failed", ERROR_INVALID_HANDLE);
		}

		dwSessionId = packet_get_tlv_value_uint(request, TLV_TYPE_DESKTOP_SESSION);
		if (!dwSessionId)
		{
			BREAK_WITH_ERROR("[UI] desktop_set. no TLV_TYPE_DESKTOP_SESSION provided", ERROR_INVALID_PARAMETER);
		}

		if (dwSessionId == -1)
		{
			dwSessionId = session_id(GetCurrentProcessId());
		}

		cpStationName = packet_get_tlv_value_string(request, TLV_TYPE_DESKTOP_STATION);
		if (!cpStationName)
		{
			BREAK_WITH_ERROR("[UI] desktop_set. no TLV_TYPE_DESKTOP_STATION provided", ERROR_INVALID_PARAMETER);
		}

		cpDesktopName = packet_get_tlv_value_string(request, TLV_TYPE_DESKTOP_NAME);
		if (!cpDesktopName)
		{
			BREAK_WITH_ERROR("[UI] desktop_set. no TLV_TYPE_DESKTOP_NAME provided", ERROR_INVALID_PARAMETER);
		}

		bSwitch = packet_get_tlv_value_bool(request, TLV_TYPE_DESKTOP_SWITCH);

		dprintf("[UI] desktop_set: Session %d\\%s\\%s (bSwitch=%d)", dwSessionId, cpStationName, cpDesktopName, bSwitch);

		// If we are switching desktop in our own session we proceed with the normal Windows API
		if (dwSessionId == session_id(GetCurrentProcessId()))
		{
			hWindowStation = OpenWindowStation(cpStationName, FALSE, WINSTA_ALL_ACCESS); // WINSTA_ALL_ACCESS MAXIMUM_ALLOWED
			if (!hWindowStation)
			{
				if (RevertToSelf())
				{
					hWindowStation = OpenWindowStation(cpStationName, FALSE, WINSTA_ALL_ACCESS);
				}
			}

			if (!hWindowStation)
			{
				BREAK_WITH_ERROR("[UI] desktop_set. Couldnt get the new Window Station", ERROR_INVALID_HANDLE);
			}

			hOrigWindowStation = GetProcessWindowStation();

			if (!SetProcessWindowStation(hWindowStation))
			{
				BREAK_ON_ERROR("[UI] desktop_set. SetProcessWindowStation failed");
			}

			hDesktop = OpenDesktop(cpDesktopName, 0, FALSE, GENERIC_ALL);
			if (!hDesktop)
			{
				BREAK_ON_ERROR("[UI] desktop_set. OpenDesktop failed");
			}

			if (!SetThreadDesktop(hDesktop))
			{
				BREAK_ON_ERROR("[UI] desktop_set. SetThreadDesktop failed");
			}

			if (bSwitch)
			{
				if (!SwitchDesktop(hDesktop))
				{
					BREAK_ON_ERROR("[UI] desktop_set. SwitchDesktop failed");
				}
			}

			core_update_desktop(remote, dwSessionId, cpStationName, cpDesktopName);
		}
		else
		{
			// if we are to use a desktop from a session which is not our own...
			BREAK_WITH_ERROR("[UI] desktop_set. Currently unable to set a desktop from an external session", ERROR_ACCESS_DENIED);
		}

	} while (0);

	if (response)
	{
		packet_transmit_response(dwResult, remote, response);
	}

	if (hDesktop)
	{
		CloseDesktop(hDesktop);
	}

	if (hWindowStation)
	{
		CloseWindowStation(hWindowStation);
	}

	if (hOrigWindowStation)
	{
		SetProcessWindowStation(hOrigWindowStation);
	}

	if (dwResult != ERROR_SUCCESS)
	{
		core_update_desktop(remote, -1, NULL, NULL);
	}

	return ERROR_SUCCESS;
}

/*
 * Worker thread for desktop screenshot. Creates a named pipe and reads in the 
 * screenshot for the first client which connects to it.
 */
DWORD THREADCALL desktop_screenshot_thread(THREAD * thread)
{
	DWORD dwResult = ERROR_ACCESS_DENIED;
	HANDLE hServerPipe = NULL;
	HANDLE hToken = NULL;
	char * cpNamedPipe = NULL;
	Packet * response = NULL;
	BYTE * pBuffer = NULL;
	DWORD dwRead = 0;
	DWORD dwLength = 0;
	DWORD dwTotal = 0;

	do
	{
		if (!thread)
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot_thread. invalid thread", ERROR_BAD_ARGUMENTS);
		}

		cpNamedPipe = (char *)thread->parameter1;
		response = (Packet *)thread->parameter2;

		if (!cpNamedPipe || !response)
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot_thread. invalid thread arguments", ERROR_BAD_ARGUMENTS);
		}

		dprintf("[UI] desktop_screenshot_thread. cpNamedPipe=%s", cpNamedPipe);

		// create the named pipe for the client service to connect to
		hServerPipe = CreateNamedPipe(cpNamedPipe, PIPE_ACCESS_DUPLEX, PIPE_TYPE_MESSAGE | PIPE_WAIT, 2, 0, 0, 0, NULL);
		if (!hServerPipe)
		{
			BREAK_ON_ERROR("[UI] desktop_screenshot_thread. CreateNamedPipe failed");
		}

		while (TRUE)
		{
			if (event_poll(thread->sigterm, 0))
			{
				BREAK_WITH_ERROR("[UI] desktop_screenshot_thread. thread->sigterm received", ERROR_DBG_TERMINATE_THREAD);
			}

			// wait for a client to connect to our named pipe...
			if (!ConnectNamedPipe(hServerPipe, NULL) && GetLastError() != ERROR_PIPE_CONNECTED)
			{
				continue;
			}

			dprintf("[UI] desktop_screenshot_thread. got client conn.");

			if (!ReadFile(hServerPipe, &dwLength, sizeof(DWORD), &dwRead, NULL))
			{
				BREAK_ON_ERROR("[UI] desktop_screenshot_thread. ReadFile 1 failed");
			}

			// a client can send a zero length to indicate that it cant get a screenshot.
			if (!dwLength)
			{
				BREAK_WITH_ERROR("[UI] desktop_screenshot_thread. dwLength == 0", ERROR_BAD_LENGTH);
			}

			pBuffer = (BYTE *)malloc(dwLength);
			if (!pBuffer)
			{
				BREAK_WITH_ERROR("[UI] desktop_screenshot_thread. pBuffer malloc failed", ERROR_NOT_ENOUGH_MEMORY);
			}

			while (dwTotal < dwLength)
			{
				DWORD dwAvailable = 0;

				if (!PeekNamedPipe(hServerPipe, NULL, 0, NULL, &dwAvailable, NULL))
				{
					break;
				}

				if (!dwAvailable)
				{
					Sleep(100);
					continue;
				}

				if (!ReadFile(hServerPipe, (LPVOID)(pBuffer + dwTotal), (dwLength - dwTotal), &dwRead, NULL))
				{
					break;
				}

				dwTotal += dwRead;
			}

			dwResult = packet_add_tlv_raw(response, TLV_TYPE_DESKTOP_SCREENSHOT, pBuffer, dwTotal);

			break;
		}

	} while (0);

	if (hServerPipe)
	{
		DisconnectNamedPipe(hServerPipe);
		CLOSE_HANDLE(hServerPipe);
	}

	SAFE_FREE(pBuffer);

	dprintf("[UI] desktop_screenshot_thread finishing, dwResult=%d", dwResult);

	return dwResult;
}

/*
 * Take a screenshot of the desktop and transmit the image (in JPEG format) back to the client.
 */
DWORD request_ui_desktop_screenshot(Remote * remote, Packet * request)
{
	DWORD dwResult = ERROR_SUCCESS;
	Packet * response = NULL;
	THREAD * pPipeThread = NULL;
	LPVOID lpDllBuffer = NULL;
	DLL_BUFFER DllBuffer = { 0 };
	char cNamedPipe[MAX_PATH] = { 0 };
	char cCommandLine[MAX_PATH] = { 0 };
	int quality = 0;
	DWORD dwDllLength = 0;
	DWORD dwPipeName = 0;
	DWORD dwCurrentSessionId = 0;
	DWORD dwActiveSessionId = 0;

	do
	{
		response = packet_create_response(request);
		if (!response)
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot. packet_create_response failed", ERROR_INVALID_HANDLE);
		}

		quality = packet_get_tlv_value_uint(request, TLV_TYPE_DESKTOP_SCREENSHOT_QUALITY);
		if (quality < 1 || quality > 100)
		{
			quality = 50;
		}

		// get the x86 and x64 screenshot dll's. we are not obliged to send both but we reduce the number of processes
		// we can inject into (wow64 and x64) if we only send one type on an x64 system. If we are on an x86 system
		// we dont need to send the x64 screenshot dll as there will be no x64 processes to inject it into.
		DllBuffer.dwPE32DllLenght = packet_get_tlv_value_uint(request, TLV_TYPE_DESKTOP_SCREENSHOT_PE32DLL_LENGTH);
		DllBuffer.lpPE32DllBuffer = packet_get_tlv_value_string(request, TLV_TYPE_DESKTOP_SCREENSHOT_PE32DLL_BUFFER);

		DllBuffer.dwPE64DllLenght = packet_get_tlv_value_uint(request, TLV_TYPE_DESKTOP_SCREENSHOT_PE64DLL_LENGTH);
		DllBuffer.lpPE64DllBuffer = packet_get_tlv_value_string(request, TLV_TYPE_DESKTOP_SCREENSHOT_PE64DLL_BUFFER);

		if (!DllBuffer.lpPE32DllBuffer && !DllBuffer.lpPE64DllBuffer)
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot. Invalid dll arguments, at least 1 dll must be supplied", ERROR_BAD_ARGUMENTS);
		}

		// get the session id that our host process belongs to
		dwCurrentSessionId = session_id(GetCurrentProcessId());

		// get the session id for the interactive session
		dwActiveSessionId = session_activeid();

		// create a uniuqe pipe name for our named pipe server
		dwPipeName = GetTickCount();

		_snprintf(cNamedPipe, MAX_PATH, "\\\\.\\pipe\\%08X", dwPipeName);

		// create the commandline to pass to the screenshot dll when we inject it
		_snprintf(cCommandLine, MAX_PATH, "/s /q:%d /p:0x%08X\x00", quality, dwPipeName);

		dprintf("[UI] desktop_screenshot. dwCurrentSessionId=%d, dwActiveSessionId=%d, cCommandLine=%s\n", dwCurrentSessionId, dwActiveSessionId, cCommandLine);

		// start a thread to create a named pipe server and wait for a client to connect an send back the JPEG screenshot.
		pPipeThread = thread_create(desktop_screenshot_thread, &cNamedPipe, response, NULL);
		if (!pPipeThread)
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot. thread_create failed", ERROR_INVALID_HANDLE);
		}

		if (!thread_run(pPipeThread))
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot. thread_run failed", ERROR_ACCESS_DENIED);
		}

		Sleep(500);

		// do the local process or session injection
		if (dwCurrentSessionId != dwActiveSessionId)
		{
			dprintf("[UI] desktop_screenshot. Injecting into active session %d...\n", dwActiveSessionId);
			if (session_inject(dwActiveSessionId, &DllBuffer, cCommandLine) != ERROR_SUCCESS)
			{
				BREAK_WITH_ERROR("[UI] desktop_screenshot. session_inject failed", ERROR_ACCESS_DENIED);
			}
		}
		else
		{
			dprintf("[UI] desktop_screenshot. Allready in the active session %d.\n", dwActiveSessionId);
			if (ps_inject(GetCurrentProcessId(), &DllBuffer, cCommandLine) != ERROR_SUCCESS)
			{
				BREAK_WITH_ERROR("[UI] desktop_screenshot. ps_inject current process failed", ERROR_ACCESS_DENIED);
			}
		}

		// Wait for at most 30 seconds for the screenshot to happen...
		// If we have injected our code via APC injection, it may take a while for the target
		// thread to enter an alertable state and get our queued APC executed.
		WaitForSingleObject(pPipeThread->handle, 30000);

		// signal our thread to terminate if it is still running.
		thread_sigterm(pPipeThread);

		// and wait for it to terminate...
		thread_join(pPipeThread);

		// get the exit code for our pthread
		if (!GetExitCodeThread(pPipeThread->handle, &dwResult))
		{
			BREAK_WITH_ERROR("[UI] desktop_screenshot. GetExitCodeThread failed", ERROR_INVALID_HANDLE);
		}

	} while (0);

	if (response)
	{
		packet_transmit_response(dwResult, remote, response);
	}

	if (pPipeThread)
	{
		thread_sigterm(pPipeThread);
		thread_join(pPipeThread);
		thread_destroy(pPipeThread);
	}

	return dwResult;
}