/* NetHalt - Shutdown NT machines at specified time
 * Copyright (C) 2008 Daniel Collins <solemnwarning@solemnwarning.net>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *	* Redistributions of source code must retain the above copyright
 *	  notice, this list of conditions and the following disclaimer.
 *
 *	* Redistributions in binary form must reproduce the above copyright
 *	  notice, this list of conditions and the following disclaimer in the
 *	  documentation and/or other materials provided with the distribution.
 *
 *	* Neither the name of the author nor the names of its contributors may
 *	  be used to endorse or promote products derived from this software
 *	  without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <time.h>

#define VERSION "1.0"

#define TXT_FIXED 101
#define TXT_REMAIN 102
#define BTN_DELAY 103
#define BTN_ABORT 104
#define TIM_TIMER 105
#define TXT_ABOUT 106

#define DSECS_HOURS(dsecs) (dsecs / 3600)
#define DSECS_MINS(dsecs) ((dsecs - (3600*DSECS_HOURS(dsecs))) / 60)
#define DSECS_SECS(dsecs) (dsecs - (3600*DSECS_HOURS(dsecs)) - (60*DSECS_MINS(dsecs)))

static void WINAPI ServiceMain(DWORD argc, char **argv);
static void ServiceState(DWORD state);
static void ServiceCallback(DWORD signal);
static void ServiceExit(int exitcode);
static void LogError(char const *fmt, ...);
static LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam);
static DWORD RealMain(void *param);
static void sd_dialog(void);
static int expr_compare(char const *str, char const *expr);
static char *w32_error(DWORD errnum);
static int ncase_match(char const *str1, char const *str2);
static int read_dtime(char *dtime);
static void shutdown_now(void);

static SERVICE_STATUS_HANDLE svc_handle = 0;
static SERVICE_STATUS svc_status;
static SC_HANDLE svc_retlock;
static DWORD MainThread;
static BOOL ServiceRunning = TRUE;

static int warnsecs = 30;
static int delaysecs = 0;
static int edelay = FALSE, eabort = FALSE;

static int sdtime_dsecs = -1;
static time_t sdtime;

int main(int argc, char **argv) {
	SERVICE_TABLE_ENTRY serviceTable[] = {
		{"NetHalt", &ServiceMain},
		{NULL, NULL}
	};
	
	if(!(StartServiceCtrlDispatcher(serviceTable))) {
		LogError("StartServiceCtrlDispatcher() failed: %d", GetLastError());
		return 1;
	}
	
	return 0;
}

static void WINAPI ServiceMain(DWORD argc, char **argv) {
	svc_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS;
	svc_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SHUTDOWN;
	svc_status.dwWin32ExitCode = NO_ERROR;
	svc_status.dwCheckPoint = 0;
	svc_status.dwWaitHint = 10000;
	
	MainThread = GetCurrentThreadId();
	
	svc_handle = RegisterServiceCtrlHandler("NetHalt", (LPHANDLER_FUNCTION)&ServiceCallback);
	if(!svc_handle) {
		LogError("RegisterServiceCtrlHandler() failed: %s", w32_error(GetLastError()));
		return;
	}
	
	ServiceState(SERVICE_START_PENDING);
	
	svc_retlock = CreateEvent(0, TRUE, FALSE, 0);
	if(!svc_retlock) {
		LogError("CreateEvent() failed: %s", w32_error(GetLastError()));
		ServiceState(SERVICE_STOPPED);
		return;
	}
	
	if(!CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)&RealMain, NULL, 0, NULL)) {
		LogError("CreateThread() failed: %s", w32_error(GetLastError()));
		ServiceState(SERVICE_STOPPED);
		return;
	}
	
	ServiceState(SERVICE_RUNNING);
	
	WaitForSingleObject(svc_retlock, INFINITE);
	ServiceState(SERVICE_STOPPED);
}

static void ServiceState(DWORD state) {
	svc_status.dwCurrentState = state;
	svc_status.dwCheckPoint++;
	
	SetServiceStatus(svc_handle, &svc_status);
}

static void ServiceCallback(DWORD signal) {
	if(signal == SERVICE_CONTROL_SHUTDOWN || signal == SERVICE_CONTROL_STOP) {
		ServiceState(SERVICE_STOP_PENDING);
		ServiceExit(0);
		return;
	}
	
	ServiceState(svc_status.dwCurrentState);
}

static void ServiceExit(int exitcode) {
	svc_status.dwWin32ExitCode = exitcode;
	ServiceState(SERVICE_STOP_PENDING);
	
	ServiceRunning = FALSE;
	SetEvent(svc_retlock);
	
	if(GetCurrentThreadId() != MainThread) {
		ExitThread(exitcode);
	}
}

static void LogError(char const *fmt, ...) {
	char msgbuf[512];
	char *msgs[] = {msgbuf, NULL};
	static HANDLE evlog = NULL;
	va_list argv;
	
	if(!evlog) {
		if(!(evlog = OpenEventLog(NULL, "Application"))) {
			return;
		}
	}
	
	va_start(argv, fmt);
	vsnprintf(msgbuf, 512, fmt, argv);
	va_end(argv);
	
	ReportEvent(evlog, EVENTLOG_ERROR_TYPE, 0, 0, NULL, 1, 0, (LPCTSTR*)msgs, NULL);
}

static LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
	HFONT hfDefault;
	HWND txtFixed, txtRemain, btnDelay, btnAbort, txtAbout;
	char msgbuf[256];
	int rmins, rsecs, remain;
	struct tm *ltime;
	
	switch(msg) {
		case WM_CREATE:
		{	
			hfDefault = GetStockObject(DEFAULT_GUI_FONT);
			
			txtFixed = CreateWindow(
				"STATIC",
				"This machine is scheduled to shutdown\nAll unsaved work will be lost!",
				WS_CHILD | WS_VISIBLE,
				5, 5,
				300, 60,
				hwnd,
				(HMENU)TXT_FIXED,
				GetModuleHandle(NULL),
				NULL
			);
			if(txtFixed == NULL) {
				LogError("Failed to create txtFixed: %s", w32_error(GetLastError()));
			}
			
			txtRemain = CreateWindow(
				"STATIC",
				"",
				WS_CHILD | WS_VISIBLE,
				5, 45,
				300, 20,
				hwnd,
				(HMENU)TXT_REMAIN,
				GetModuleHandle(NULL),
				NULL
			);
			if(txtRemain == NULL) {
				LogError("Failed to create txtRemain: %s", w32_error(GetLastError()));
			}
			
			btnDelay = CreateWindow(
				"BUTTON",
				"Delay shutdown",
				WS_CHILD | WS_VISIBLE | BS_PUSHBUTTON,
				5, 65,
				140, 40,
				hwnd,
				(HMENU)BTN_DELAY,
				GetModuleHandle(NULL),
				NULL
			);
			if(btnDelay == NULL) {
				LogError("Failed to create btnDelay: %s", w32_error(GetLastError()));
			}
			
			btnAbort = CreateWindow(
				"BUTTON",
				"Abort shutdown",
				WS_CHILD | WS_VISIBLE | BS_PUSHBUTTON,
				150, 65,
				140, 40,
				hwnd,
				(HMENU)BTN_ABORT,
				GetModuleHandle(NULL),
				NULL
			);
			if(btnAbort == NULL) {
				LogError("Failed to create btnAbort: %s", w32_error(GetLastError()));
			}
			
			txtAbout = CreateWindow(
				"STATIC",
				"NetHalt v" VERSION ", Copyright (C) 2008 Daniel Collins",
				WS_CHILD | WS_VISIBLE,
				5, 110,
				300, 20,
				hwnd,
				(HMENU)TXT_ABOUT,
				GetModuleHandle(NULL),
				NULL
			);
			if(txtAbout == NULL) {
				LogError("Failed to create txtAbout: %s", w32_error(GetLastError()));
			}
			
			EnableWindow(btnDelay, edelay);
			EnableWindow(btnAbort, eabort);
			
			SetTimer(hwnd, TIM_TIMER, 1000, NULL);
			SendMessage(hwnd, WM_TIMER, TIM_TIMER, GetTickCount());
			
			SendMessage(txtFixed, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
			SendMessage(txtRemain, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
			SendMessage(btnDelay, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
			SendMessage(btnAbort, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
			SendMessage(txtAbout, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
			
			break;
		}
		case WM_CLOSE:
			break;
		case WM_DESTROY:
			PostQuitMessage(0);
			break;
		case WM_TIMER:
			txtRemain = GetDlgItem(hwnd, TXT_REMAIN);
			
			remain = sdtime - time(NULL);
			
			if(!ServiceRunning || remain <= 0) {
				DestroyWindow(hwnd);
				break;
			}
			rmins = remain / 60;
			rsecs = remain - (60*rmins);
			
			sprintf(msgbuf, "Time remaining: %02d:%02d\n", rmins, rsecs);
			SetWindowText(txtRemain, msgbuf);
			break;
		case WM_COMMAND:
			if(LOWORD(wParam) == BTN_DELAY) {
				sdtime += delaysecs;
				DestroyWindow(hwnd);
			}
			if(LOWORD(wParam) == BTN_ABORT) {
				sdtime = time(NULL);
				ltime = localtime(&sdtime);
				ltime->tm_hour = DSECS_HOURS(sdtime_dsecs);
				ltime->tm_min = DSECS_MINS(sdtime_dsecs);
				ltime->tm_sec = DSECS_SECS(sdtime_dsecs);
				sdtime = mktime(ltime) + 86400;
				
				DestroyWindow(hwnd);
			}
			
			break;
		default:
			return DefWindowProc(hwnd, msg, wParam, lParam);
	}
	
	return 0;
}

static DWORD RealMain(void *param) {
	WNDCLASSEX wc;
	FILE *inifile;
	char inipath[512];
	char linebuf[1024], *name, *value;
	int mblock = 0, n;
	char nbname[256];
	DWORD nbsize = 256;
	struct tm *ltime;
	
	wc.cbSize        = sizeof(WNDCLASSEX);
	wc.style         = 0;
	wc.lpfnWndProc   = &WndProc;
	wc.cbClsExtra    = 0;
	wc.cbWndExtra    = 0;
	wc.hInstance     = GetModuleHandle(NULL);
	wc.hIcon         = LoadIcon(NULL, IDI_APPLICATION);
	wc.hCursor       = LoadCursor(NULL, IDC_ARROW);
	wc.hbrBackground = GetSysColorBrush(COLOR_ACTIVEBORDER);
	wc.lpszMenuName  = NULL;
	wc.lpszClassName = "NetHaltDialog";
	wc.hIconSm       = LoadIcon(NULL, IDI_APPLICATION);
	
	if(!RegisterClassEx(&wc)) {
		LogError("Failed to register window class: %s", w32_error(GetLastError()));
		ServiceExit(1);
		return 1;
	}
	
	GetComputerName(nbname, &nbsize);
	
	GetModuleFileName(NULL, inipath, 512);
	
	n = strlen(inipath)-1;
	while(n >= 0 && inipath[n] != '\\') {
		inipath[n--] = '\0';
	}
	
	strcat(inipath, "nethalt.ini");
	
	inifile = fopen(inipath, "r");
	if(!inifile) {
		LogError("Failed to open INI file: %s", w32_error(GetLastError()));
		ServiceExit(1);
	}
	
	while(fgets(linebuf, 1024, inifile)) {
		name = linebuf+strspn(linebuf, " \r\t\n");
		
		if(name[0] == '#' || name[0] == ';') {
			continue;
		}
		if(name[0] == '[') {
			name++;
			name[strcspn(name, " \t\n\r]")] = '\0';
			mblock = expr_compare(nbname, name);
			
			continue;
		}
		if(!mblock) {
			continue;
		}
		
		value = name+strcspn(name, " \r\t\n=");
		
		if(value[0] != '\0') {
			value[0] = '\0';
			value++;
			
			value += strspn(value, " \r\t\n=");
		}
		
		if(ncase_match(name, "Shutdown")) {
			if(ncase_match(value, "off")) {
				sdtime_dsecs = -1;
				continue;
			}
			
			sdtime_dsecs = read_dtime(value);
			
			sdtime = time(NULL);
			ltime = localtime(&sdtime);
			ltime->tm_hour = DSECS_HOURS(sdtime_dsecs);
			ltime->tm_min = DSECS_MINS(sdtime_dsecs);
			ltime->tm_sec = DSECS_SECS(sdtime_dsecs);
			sdtime = mktime(ltime);
			
			if(sdtime < time(NULL)) {
				sdtime += 86400;
			}
		}
		if(ncase_match(name, "Warning")) {
			warnsecs = atoi(value);
		}
		if(ncase_match(name, "DelaySeconds")) {
			delaysecs = atoi(value);
			edelay = (delaysecs > 0 ? TRUE : FALSE);
		}
		if(ncase_match(name, "AllowAbort")) {
			if(atoi(value) == 0) {
				eabort = FALSE;
			}else{
				eabort = TRUE;
			}
		}
	}
	
	fclose(inifile);
	
	if(sdtime_dsecs == -1) {
		LogError("No shutdown time defined");
		ServiceExit(1);
	}
	
	while(ServiceRunning) {
		if(time(NULL) >= sdtime) {
			shutdown_now();
			break;
		}
		if(time(NULL) >= (sdtime-warnsecs)) {
			sd_dialog();
		}
		
		Sleep(1000);
	}
	
	ServiceExit(0);
	return 0;
}

static void sd_dialog(void) {
	MSG Msg;
	HWND hwnd;
	
	hwnd = CreateWindowEx(
		WS_EX_TOPMOST,
		"NetHaltDialog",
		"NetHalt - Timed shutdown",
		WS_DLGFRAME,
		CW_USEDEFAULT, CW_USEDEFAULT,
		300, 155,
		NULL, NULL, GetModuleHandle(NULL), NULL
	);
	
	if(hwnd == NULL) {
		LogError("Window Creation Failed: %s", w32_error(GetLastError()));
		return;
	}
	
	ShowWindow(hwnd, SW_SHOW);
	UpdateWindow(hwnd);
	
	while(GetMessage(&Msg, NULL, 0, 0) > 0) {
		TranslateMessage(&Msg);
		DispatchMessage(&Msg);
	}
}

/* Compare the supplied string and expression
 * Returns 1 upon match, zero otherwise.
*/
static int expr_compare(char const *str, char const *expr) {
	while(1) {
		if(expr[0] == '\0' && str[0] != '\0') {
			return 0;
		}
		if(str[0] == '\0') {
			if(expr[0] == '\0' || expr[0] == '*') {
				break;
			}
			
			return 0;
		}
		
		if(expr[0] == '*') {
			if(expr[1] == str[0]) {
				expr += 2;
			}
			str++;
			
			continue;
		}
		if(expr[0] == '?' && str[0] != '\0') {
			expr++;
			str++;
			
			continue;
		}
		if(expr[0] == '#' && isdigit(str[0])) {
			expr++;
			str++;
			
			continue;
		}
		if(tolower(expr[0]) == tolower(str[0])) {
			expr++;
			str++;
			
			continue;
		}
		
		return 0;
	}
	
	return 1;
}

static char *w32_error(DWORD errnum) {
	static char buf[1024] = {'\0'};
	
	FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, errnum, 0, buf, 1023, NULL);
	buf[strcspn(buf, "\r\n")] = '\0';
	return buf;	
}

static int ncase_match(char const *str1, char const *str2) {
	size_t pos = 0;
	
	while(tolower(str1[pos]) == tolower(str2[pos])) {
		if(str1[pos] == '\0') {
			return 1;
		}
		
		pos++;
	}
	
	return 0;
}

static int read_dtime(char *dtime) {
	char *hours, *mins, *secs;
	
	hours = dtime;
	mins = hours+strcspn(hours, ":");
	mins += strspn(mins, ":");
	secs = mins+strcspn(mins, ":");
	secs += strspn(secs, ":");
	
	return (atoi(hours)*3600) + (atoi(mins)*60) + atoi(secs);
}

static void shutdown_now(void) {
	HANDLE hToken;
	TOKEN_PRIVILEGES tkp;
	
	if(!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		LogError("OpenProcessToken() failed: %s", w32_error(GetLastError()));
		ServiceExit(1);
	}
	LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);
	
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	
	if(!AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, (PTOKEN_PRIVILEGES)NULL, 0)) {
		LogError("AdjustTokenPrivileges() failed: %s", w32_error(GetLastError()));
		ServiceExit(1);
	}
	
	if(!InitiateSystemShutdown(NULL, NULL, 0, TRUE, FALSE)) {
		LogError("InitiateShutdown() failed: %s", w32_error(GetLastError()));
		ServiceExit(1);
	}
}