// Douglas Thrift
//
// CCS Computer Science
//
// Windows Finger Daemon

#include <windows.h>
#include <tchar.h>
#include <signal.h>
#include <wtsapi32.h>

#include <iostream>
#include <string>
#include <sstream>
#include <vector>

inline std::string Utf8(const std::wstring &wstring)
{
	std::string string(WideCharToMultiByte(CP_UTF8, 0, wstring.data(), int(wstring.size()), NULL, 0, NULL, NULL), '\0');

	WideCharToMultiByte(CP_UTF8, 0, wstring.data(), int(wstring.size()), const_cast<LPSTR>(string.data()), int(string.size()), NULL, NULL);

	return string;
}

class Finger
{
private:
	std::ostringstream stream;

public:
	Finger(bool full)
	{
		if (!full)
			stream << "Login\t\tName\t\tTTY\r\n";

		PWTS_SESSION_INFO sessions;
		DWORD count;

		WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &sessions, &count);

		for (PWTS_SESSION_INFO session = sessions; session != sessions + count; ++session)
		{
			LPTSTR name;
			DWORD size;

			WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, session->SessionId, WTSUserName, &name, &size);

			std::string name_(Utf8(name));

			if (!name_.empty())
				stream << name_ << "\r\n";

			WTSFreeMemory(name);

			//PWTS_CLIENT_ADDRESS address;
			//WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, session->SessionId, WTSClientAddress, reinterpret_cast<LPTSTR *>(&address), &size);
			//WTSFreeMemory(address);
		}

		WTSFreeMemory(sessions);
	}

	Finger(const std::string &name)
	{
		//NetQueryDisplayInformation
		//NetUserGetInfo
	}

	inline operator std::string()
	{
		return stream.str();
	}
};

LPTSTR name;
SERVICE_STATUS status;
SERVICE_STATUS_HANDLE handle;
HANDLE stop;
WSADATA data;
std::vector<HANDLE> threads;

void FingerDaemon();
DWORD WINAPI FingerListen(LPVOID server_);
DWORD WINAPI FingerDo(LPVOID client_);
void FingerStop(int);
void WINAPI FingerMain(DWORD argc, LPTSTR *argv);
void WINAPI FingerControl(DWORD control);

int _tmain(int argc, TCHAR *argv[])
{
	SERVICE_TABLE_ENTRY entry[] = { { TEXT("CCS Finger Daemon"), LPSERVICE_MAIN_FUNCTION(FingerMain) }, { NULL, NULL } };

	if (!StartServiceCtrlDispatcher(entry))
	{
		DWORD error(GetLastError());

		switch (error)
		{
		case ERROR_FAILED_SERVICE_CONTROLLER_CONNECT:
			goto go;

		case ERROR_INVALID_DATA:
			std::cerr << "ERROR_INVALID_DATA" << std::endl;

			break;

		case ERROR_SERVICE_ALREADY_RUNNING:
			std::cerr << "ERROR_SERVICE_ALREADY_RUNNING" << std::endl;

			break;

		default:
			std::cerr << error << std::endl;
		}

		return 1;
	}

	return 0;

go:
	stop = CreateEvent(NULL, TRUE, FALSE, NULL);

	signal(SIGINT, FingerStop);

	try
	{
		FingerDaemon();
	}
	catch (...) {}

	return 0;
}

void FingerDaemon()
{
	WSAStartup(MAKEWORD(2, 0), &data);

	//

	SOCKET server(socket(PF_INET, SOCK_STREAM, IPPROTO_TCP));
	SOCKADDR_IN service;
	
	service.sin_family = AF_INET;
	service.sin_addr.s_addr = inet_addr("0.0.0.0");
	service.sin_port = htons(79);

	bind(server, (SOCKADDR *)(&service), sizeof (service));
	listen(server, SOMAXCONN);

	threads.push_back(CreateThread(NULL, 0, FingerListen, &server, 0, NULL));

	while(WaitForSingleObject(stop, 1000) != WAIT_OBJECT_0);

	//

	closesocket(server);

	WSACleanup();
}

DWORD WINAPI FingerListen(LPVOID server_)
{
	SOCKET &server(*reinterpret_cast<SOCKET *>(server_));

	while (true)
	{
		SOCKET client(accept(server, NULL, NULL));

		threads.push_back(CreateThread(NULL, 0, FingerDo, &client, 0, NULL));
	}

	return 0;
}

DWORD WINAPI FingerDo(LPVOID client_)
{
	SOCKET &client(*reinterpret_cast<SOCKET *>(client_));
	char buffer[1024];
	std::istringstream stream(std::string(buffer, recv(client, buffer, sizeof buffer, 0)));
	std::string line;

	std::getline(stream, line);

	stream.str(line);

	std::getline(stream, line, '\r');

	stream.str(line);

	std::string name;
	bool full(false);

	while (stream >> line)
		if (line == "/W")
			full = true;
		else
			name = line;

	std::string finger(name.empty() ? Finger(full) : Finger(name));

	send(client, finger.data(), int(finger.size()), 0);
	closesocket(client);

	return 0;
}

void FingerStop(int)
{
	SetEvent(stop);
}

void WINAPI FingerMain(DWORD argc, LPTSTR *argv)
{
	name = argv[0];
	handle = RegisterServiceCtrlHandler(name, LPHANDLER_FUNCTION(FingerControl));
	stop = CreateEvent(NULL, TRUE, FALSE, NULL);

	//

	status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
	status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
	status.dwServiceSpecificExitCode = 0;
	status.dwCurrentState = SERVICE_RUNNING;
	status.dwWin32ExitCode = NO_ERROR;
	status.dwCheckPoint = 0;
	status.dwWaitHint = 0;

	SetServiceStatus(handle, &status);

	try
	{
		FingerDaemon();
	}
	catch (...) {}

	status.dwCurrentState = SERVICE_STOPPED;
	status.dwWin32ExitCode = NO_ERROR;
	status.dwCheckPoint = 0;
	status.dwWaitHint = 0;

	SetServiceStatus(handle, &status);
}

void WINAPI FingerControl(DWORD control)
{
	switch (control)
	{
	case SERVICE_CONTROL_STOP:
	case SERVICE_CONTROL_SHUTDOWN:
		status.dwCurrentState = SERVICE_STOP_PENDING;

		SetEvent(stop);

		break;

	case SERVICE_CONTROL_INTERROGATE:
		break;
	}

	SetServiceStatus(handle, &status);
}