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

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

#include <algorithm>
#include <iostream>
#include <iomanip>
#include <string>
#include <sstream>
#include <vector>
#include <map>
#include <cwctype>

std::ostream &operator<<(std::ostream &out, WTS_CONNECTSTATE_CLASS status)
{
	switch (status)
	{
	case WTSActive:
		return out << "Active";

	case WTSConnected:
		return out << "Connected";

	case WTSConnectQuery:
		return out << "Connecting";

	case WTSShadow:
		return out << "Shadowing";

	case WTSDisconnected:
		return out << "Disconnected";

	case WTSIdle:
		return out << "Idle";

	case WTSListen:
		return out << "Listening";

	case WTSReset:
		return out << "Reseting";

	case WTSDown:
		return out << "Down";

	case WTSInit:
		return out << "Initializing";

	default:
		return out << "Unknown";
	}
}

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
{
	class Login
	{
		std::wstring login;
		mutable USER_INFO_11 *info;

		inline void GetInfo() const
		{
			if (!info)
				::NetUserGetInfo(NULL, login.c_str(), 11, reinterpret_cast<LPBYTE *>(&info));
		}
	public:
		std::string session;
		DWORD id;
		WTS_CONNECTSTATE_CLASS status;

		Login(const std::wstring &login, DWORD id, const std::wstring &session, WTS_CONNECTSTATE_CLASS status) : login(login), info(NULL), id(id), session(Utf8(session)), status(status)
		{
		}

		~Login()
		{
			if (!info)
				::NetApiBufferFree(info);
		}

		std::string GetShortName() const
		{
			GetInfo();

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

		std::string GetName() const
		{
			GetInfo();

			return Utf8(info->usri11_full_name);
		}

		std::string GetDirectory() const
		{
			HANDLE token;

			if (::WTSQueryUserToken(id, &token))
			{
				TCHAR directory[MAX_PATH];

				::SHGetFolderPath(NULL, CSIDL_PROFILE, token, SHGFP_TYPE_CURRENT, directory);
				::CloseHandle(token);

				return Utf8(directory);
			}
			else
			{
				TCHAR directory_[MAX_PATH];

				::SHGetFolderPath(NULL, CSIDL_PROFILE, NULL, SHGFP_TYPE_CURRENT, directory_);

				std::wstring directory(directory_);

				directory.replace(directory.rfind('\\') + 1, std::wstring::npos, login);

				return Utf8(directory);
			}
		}
	};

	std::ostringstream stream;
	std::multimap<std::string, Login> logins;

	void Full()
	{
		for (std::multimap<std::string, Login>::const_iterator login(logins.begin()); login != logins.end(); login = logins.upper_bound(login->first))
		{
			stream << "Login: " << login->first << std::string(33 - login->first.size(), ' ') << "Name: " << login->second.GetName() << "\r\nDirectory: " << login->second.GetDirectory() << "\r\n";
		}
	}

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

			std::wstring name_(name);

			::WTSFreeMemory(name);

			if (!name_.empty())
			{
				std::transform(name_.begin(), name_.end(), name_.begin(), std::towlower);

				LPTSTR session_;

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

				Login login(name_, session->SessionId, session_, session->State);

				::WTSFreeMemory(session_);

				logins.insert(std::make_pair(Utf8(name_), login));
			}
		}

		::WTSFreeMemory(sessions);

		if (!full)
		{
			stream << "Login" << std::string(11, ' ') << "Name" << std::string(17, ' ') << "Id     Session       Status\r\n";

			for (std::multimap<std::string, Login>::const_iterator login(logins.begin()); login != logins.end(); ++login)
				stream << login->first << std::string(16 - login->first.size(), ' ') << login->second.GetShortName() << std::string(21 - login->second.GetShortName().size(), ' ') << std::setw(5) << std::left << login->second.id << "  " << login->second.session << std::string(14 - login->second.session.size(), ' ') << login->second.status << "\r\n";
		}
		else
			Full();
	}

	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("CCSFinger"), 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:
	for (int index(1); index != argc; ++index)
	{
		std::wstring arg(argv[index]);

		if (arg == TEXT("create"))
		{
			TCHAR file[MAX_PATH];

			::GetModuleFileName(NULL, file, MAX_PATH);

			SC_HANDLE manager(::OpenSCManager(NULL, SERVICES_ACTIVE_DATABASE, SC_MANAGER_CREATE_SERVICE)), service(::CreateService(manager, TEXT("CCSFinger"), TEXT("CCS Finger Daemon"), 0, SERVICE_WIN32_OWN_PROCESS, SERVICE_DEMAND_START, SERVICE_ERROR_NORMAL, file, NULL, NULL, TEXT("TermService\0"), NULL, TEXT("")));

			if (!service)
				std::cerr << ::GetLastError() << std::endl;
			else
				::CloseServiceHandle(service);

			::CloseServiceHandle(manager);

			return 0;
		}
		else if (arg == TEXT("delete"))
		{
			SC_HANDLE manager(::OpenSCManager(NULL, SERVICES_ACTIVE_DATABASE, DELETE)), service(::OpenService(manager, TEXT("CCSFinger"), DELETE));

			if (!::DeleteService(service))
				std::cerr << ::GetLastError() << std::endl;

			::CloseServiceHandle(service);
			::CloseServiceHandle(manager);

			return 0;
		}
	}

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