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

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

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

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:
	struct Login
	{
		std::string login, name, session;
		Login(const std::wstring &login, const std::wstring &session) : login(Utf8(login)), session(Utf8(session))
		{
			USER_INFO_11 *info;

			::NetUserGetInfo(NULL, login.c_str(), 11, reinterpret_cast<LPBYTE *>(&info));

			name = Utf8(std::wstring(info->usri11_full_name).substr(0, 20));

			::NetApiBufferFree(info);
		}
		bool operator<(const Login &login) const
		{
			if (this->login == login.login)
				return session < login.session;
			else
				return this->login < login.login;
		}
	};
	std::ostringstream stream;
	std::set<Login> logins;

public:
	Finger(bool full)
	{
		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);

			if (!std::wstring(name).empty())
			{
				LPTSTR session_;

				::WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, session->SessionId, WTSWinStationName, &session_, &size);

				Login login(name, session_);

				::WTSFreeMemory(session_);

				logins.insert(login);
			}

			::WTSFreeMemory(name);
		}

		::WTSFreeMemory(sessions);

		if (!full)
		{
			stream << "Login\t\tName" << std::string(17, ' ') << "Session\r\n";

			for (std::set<Login>::const_iterator login(logins.begin()); login != logins.end(); ++login)
				stream << login->login << std::string(16 - login->login.size(), ' ') << login->name << std::string(21 - login->name.size(), ' ') << login->session << "\r\n";
		}
	}

	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);
}