/* NetHalt - NetHalt config tool
 * 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 <ctype.h>

#include "nhconfig.h"
#include "gui.h"
#include "lib.h"

static HKEY regkey;
static HWND sdlist;
static HICON tools16, tools32;
static SC_HANDLE scm, svc;

static INT_PTR CALLBACK cfg_proc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);
static int invert_checkbox(HWND hwnd);
static void update_ui(HWND hwnd);
static DWORD reg_get_dword(char const *name, DWORD defval);
static int reg_set_dword(char const *name, DWORD value);
static void reg_get_string(char const *name, char *buf, DWORD size);
static int reg_set_string(char const *name, char const *value);
static INT_PTR CALLBACK add_proc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);
static int add_sdtime(int days, char const *str);
static void update_cfg(HWND hwnd);

int main(int argc, char **argv) {
	DWORD errnum;
	
	errnum = RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SOFTWARE\\NetHalt", 0, KEY_READ | KEY_WRITE, &regkey);
	if(errnum != ERROR_SUCCESS) {
		die("Failed to open registry key: %s", w32_error(errnum));
	}
	
	tools16 = LoadImage(
		GetModuleHandle(NULL),
		MAKEINTRESOURCE(ICO_TOOLS),
		IMAGE_ICON,
		16, 16,
		LR_DEFAULTCOLOR | LR_LOADTRANSPARENT
	);
	
	tools32 = LoadImage(
		GetModuleHandle(NULL),
		MAKEINTRESOURCE(ICO_TOOLS),
		IMAGE_ICON,
		32, 32,
		LR_DEFAULTCOLOR | LR_LOADTRANSPARENT
	);
	
	scm = OpenSCManager(NULL, NULL, SC_MANAGER_CONNECT);
	if(!scm) {
		die("Failed to connect to SCM: %s", w32_error(GetLastError()));
	}
	
	svc = OpenService(scm, "NetHalt", SERVICE_QUERY_STATUS | SERVICE_START | SERVICE_STOP);
	if(!svc) {
		die("Failed to open service: %s", w32_error(GetLastError()));
	}
	
	DialogBox(GetModuleHandle(NULL), MAKEINTRESOURCE(DLG_CONFIG), NULL, &cfg_proc);
	
	CloseServiceHandle(svc);
	CloseServiceHandle(scm);
	RegCloseKey(regkey);
	return 0;
}

static INT_PTR CALLBACK cfg_proc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam) {
	int control_id, sdnum;
	HWND control;
	DWORD cfg_interval, cfg_warn, cfg_abort, cfg_delay;
	char server_name[256], server_share[256];
	char sdtimes[1024], *sdtok;
	
	if(msg == WM_INITDIALOG) {
		SendMessage(hwnd, WM_SETICON, ICON_SMALL, (LPARAM)tools16);
		SendMessage(hwnd, WM_SETICON, ICON_BIG, (LPARAM)tools32);
		
		sdlist = GetDlgItem(hwnd, LST_SHUTDOWN);
		
		reg_get_string("ServerName", server_name, 256);
		reg_get_string("ServerShare", server_share, 256);
		reg_get_string("ShutdownTimes", sdtimes, 1024);
		
		cfg_interval = reg_get_dword("ServerInterval", 0);
		cfg_warn = reg_get_dword("Warning", 300);
		cfg_abort = reg_get_dword("EnableAbort", 0);
		cfg_delay = reg_get_dword("DelayTime", 0);
		
		gui_stext(hwnd, TXT_HOST, "%s", server_name);
		gui_stext(hwnd, TXT_SHARE, "%s", server_share);
		gui_stext(hwnd, TXT_INTERVAL, "%d", cfg_interval);
		gui_stext(hwnd, TXT_WARNING, "%d", cfg_warn ? cfg_warn : 300);
		gui_stext(hwnd, TXT_DELAY, "%d", cfg_delay ? cfg_delay : 300);
		
		if(server_name[0] && server_share[0]) {
			SendDlgItemMessage(hwnd, CHK_SERVER, BM_SETCHECK, BST_CHECKED, 0);
		}
		if(cfg_warn) {
			SendDlgItemMessage(hwnd, CHK_WARNING, BM_SETCHECK, BST_CHECKED, 0);
		}
		if(cfg_abort) {
			SendDlgItemMessage(hwnd, CHK_ABORT, BM_SETCHECK, BST_CHECKED, 0);
		}
		if(cfg_delay) {
			SendDlgItemMessage(hwnd, CHK_DELAY, BM_SETCHECK, BST_CHECKED, 0);
		}
		
		sdtok = strtok(sdtimes, " ");
		while(sdtok) {
			add_sdtime(0, sdtok);
			sdtok = strtok(NULL, " ");
		}
		
		update_ui(hwnd);
		return TRUE;
	}
	if(msg == WM_CLOSE) {
		EndDialog(hwnd, 0);
		return TRUE;
	}
	if(msg == WM_COMMAND) {
		control_id = LOWORD(wparam);
		control = (HANDLE)lparam;
		
		if(control_id == BTN_EXIT) {
			SendMessage(hwnd, WM_CLOSE, 0, 0);
		}
		if(control_id >= 200 && control_id < 300) {
			invert_checkbox(control);
			update_ui(hwnd);
		}
		if(control_id == BTN_ADD) {
			sdnum = SendMessage(sdlist, LB_GETCOUNT, 0, 0);
			
			if(sdnum < 78) {
				DialogBox(
					GetModuleHandle(NULL),
					MAKEINTRESOURCE(DLG_ADD),
					hwnd,
					&add_proc
				);
			}
		}
		if(control_id == BTN_DEL) {
			sdnum = SendMessage(sdlist, LB_GETCURSEL, 0, 0);
			
			if(sdnum != LB_ERR) {
				SendMessage(sdlist, LB_DELETESTRING, sdnum, 0);
			}
		}
		if(control_id == BTN_APPLY) {
			update_cfg(hwnd);
		}
		
		return TRUE;
	}
	
	return FALSE;
}

/* Invert a checkbox between the checked/unchecked state
 * Returns the new state
*/
static int invert_checkbox(HWND hwnd) {
	int state;
	
	state = SendMessage(hwnd, BM_GETCHECK, 0, 0);
	
	if(state == BST_CHECKED) {
		state = BST_UNCHECKED;
	}else if(state == BST_UNCHECKED) {
		state = BST_CHECKED;
	}
	
	SendMessage(hwnd, BM_SETCHECK, state, 0);
	
	return state;
}

/* Display error in message box and exit */
void die_r(char const *file, int line, char const *fmt, ...) {
	char msg[1024];
	va_list argv;
	
	va_start(argv, fmt);
	vsnprintf(msg, 1024, fmt, argv);
	va_end(argv);
	
	msgboxf(MB_OK | MB_ICONERROR, "NetHalt - Error", "%s", msg);
	
	exit(1);
}

/* Enable/disable UI controls on the main window */
static void update_ui(HWND hwnd) {
	int server_state, warning_state, delay_state;
	
	server_state = SendDlgItemMessage(hwnd, CHK_SERVER, BM_GETCHECK, 0, 0);
	warning_state = SendDlgItemMessage(hwnd, CHK_WARNING, BM_GETCHECK, 0, 0);
	delay_state = SendDlgItemMessage(hwnd, CHK_DELAY, BM_GETCHECK, 0, 0);
	
	if(server_state == BST_CHECKED) {
		gui_enable(hwnd, TXT_HOST, TRUE);
		gui_enable(hwnd, TXT_SHARE, TRUE);
		gui_enable(hwnd, TXT_INTERVAL, TRUE);
		
		gui_enable(hwnd, LST_SHUTDOWN, FALSE);
		gui_enable(hwnd, BTN_ADD, FALSE);
		gui_enable(hwnd, BTN_DEL, FALSE);
		
		gui_enable(hwnd, CHK_WARNING, FALSE);
		gui_enable(hwnd, TXT_WARNING, FALSE);
		gui_enable(hwnd, CHK_ABORT, FALSE);
		gui_enable(hwnd, CHK_DELAY, FALSE);
		gui_enable(hwnd, TXT_DELAY, FALSE);
	}else{
		gui_enable(hwnd, TXT_HOST, FALSE);
		gui_enable(hwnd, TXT_SHARE, FALSE);
		gui_enable(hwnd, TXT_INTERVAL, FALSE);
		
		gui_enable(hwnd, LST_SHUTDOWN, TRUE);
		gui_enable(hwnd, BTN_ADD, TRUE);
		gui_enable(hwnd, BTN_DEL, TRUE);
		
		gui_enable(hwnd, CHK_WARNING, TRUE);
		
		if(warning_state == BST_CHECKED) {
			gui_enable(hwnd, CHK_ABORT, TRUE);
			gui_enable(hwnd, CHK_DELAY, TRUE);
			gui_enable(hwnd, TXT_WARNING, TRUE);
			
			if(delay_state == BST_CHECKED) {
				gui_enable(hwnd, TXT_DELAY, TRUE);
			}else{
				gui_enable(hwnd, TXT_DELAY, FALSE);
			}
		}else{
			gui_enable(hwnd, TXT_WARNING, FALSE);
			gui_enable(hwnd, CHK_ABORT, FALSE);
			gui_enable(hwnd, CHK_DELAY, FALSE);
			gui_enable(hwnd, TXT_DELAY, FALSE);
		}
	}
}

/* Read a DWORD value from the registry
 *
 * If an error is encountered defval will be returned, if the value does not
 * exist, defval will be written to the registry before returning.
*/
static DWORD reg_get_dword(char const *name, DWORD defval) {
	DWORD errnum, size, value;
	
	size = sizeof(DWORD);
	errnum = RegQueryValueEx(regkey, name, NULL, NULL, (BYTE*)&value, &size);
	
	if(errnum == ERROR_FILE_NOT_FOUND) {
		reg_set_dword(name, defval);
		return defval;
	}
	
	if(errnum != ERROR_SUCCESS) {
		msgboxf(
			MB_OK | MB_ICONERROR, NULL,
			"Failed to read %s from registry: %s",
			name, w32_error(errnum)
		);
		
		return defval;
	}
	
	return value;
}

/* Write a DWORD value to the registry
 * Returns nonzero on success, zero on failure
*/
static int reg_set_dword(char const *name, DWORD value) {
	DWORD errnum;
	
	errnum = RegSetValueEx(regkey, name, 0, REG_DWORD, (BYTE*)&value, sizeof(DWORD));
	if(errnum != ERROR_SUCCESS) {
		msgboxf(
			MB_OK | MB_ICONERROR, NULL,
			"Failed to write %s to registry: %s",
			name, w32_error(errnum)
		);
		
		return 0;
	}
	
	return 1;
}

/* Read a string value from the registry */
static void reg_get_string(char const *name, char *buf, DWORD size) {
	DWORD errnum, rsize = size;
	
	errnum = RegQueryValueEx(regkey, name, NULL, NULL, buf, &rsize);
	if(errnum == ERROR_FILE_NOT_FOUND) {
		strcpy(buf, "");
		return;
	}
	if(errnum != ERROR_SUCCESS) {
		msgboxf(
			MB_OK | MB_ICONERROR, NULL,
			"Failed to read %s from registry: %s",
			name, w32_error(errnum)
		);
		
		strcpy(buf, "");
	}
	
	buf[size-1] = '\0';
}

/* Write a string value to the registry
 * Returns nonzero on success, zero on failure
*/
static int reg_set_string(char const *name, char const *value) {
	DWORD errnum;
	
	errnum = RegSetValueEx(regkey, name, 0, REG_SZ, value, strlen(value)+1);
	if(errnum != ERROR_SUCCESS) {
		msgboxf(
			MB_OK | MB_ICONERROR, NULL,
			"Failed to write %s to registry: %s",
			name, w32_error(errnum)
		);
		
		return 0;
	}
	
	return 1;
}

#define DAY_CHECKBOX(day, cbox) \
	if(SendDlgItemMessage(hwnd, cbox, BM_GETCHECK, 0, 0) == BST_CHECKED) { \
		days |= (1<<day); \
	}

static INT_PTR CALLBACK add_proc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam) {
	int control_id, days = 0;
	HWND control;
	TCHAR sdtime[32];
	
	if(msg == WM_INITDIALOG) {
		SendMessage(hwnd, WM_SETICON, ICON_SMALL, (LPARAM)tools16);
		SendMessage(hwnd, WM_SETICON, ICON_BIG, (LPARAM)tools32);
		
		gui_stext(hwnd, TXT_TIME, "HH:MM:SS");
		return TRUE;
	}
	if(msg == WM_CLOSE) {
		EndDialog(hwnd, 0);
		return TRUE;
	}
	if(msg == WM_COMMAND) {
		control_id = LOWORD(wparam);
		control = (HANDLE)lparam;
		
		if(control_id == BTN_ADD) {
			DAY_CHECKBOX(1, CHK_MON);
			DAY_CHECKBOX(2, CHK_TUE);
			DAY_CHECKBOX(3, CHK_WED);
			DAY_CHECKBOX(4, CHK_THU);
			DAY_CHECKBOX(5, CHK_FRI);
			DAY_CHECKBOX(6, CHK_SAT);
			DAY_CHECKBOX(0, CHK_SUN);
			
			if(days == 0) {
				MessageBox(hwnd, "You selected no days", NULL, MB_OK);
				return TRUE;
			}
			
			SendDlgItemMessage(hwnd, TXT_TIME, WM_GETTEXT, 32, (LPARAM)sdtime);
			
			if(add_sdtime(days, sdtime)) {
				EndDialog(hwnd, 0);
			}else{
				MessageBox(hwnd, "You entered an invalid time", NULL, MB_OK);
			}
		}
		if(control_id == BTN_EXIT) {
			SendMessage(hwnd, WM_CLOSE, 0, 0);
		}
		if(control_id >= 200 && control_id < 300) {
			invert_checkbox(control);
		}
		
		return TRUE;
	}
	
	return FALSE;
}

/* Add a shutdown time to the list
 * The shutdown time string may contain a days field ONLY if days is zero
 * Returns zero if an invalid time was passed, nonzero otherwise
*/
static int add_sdtime(int days, char const *str) {
	int step = 0, n, hour, min, sec = 0;
	char sdstring[64];
	
	if(!isdigit(str[0])) {
		return 0;
	}
	
	for(n = 0; str[n]; n++) {
		if(isdigit(str[n])) {
			continue;
		}
		if(days == 0 && str[n] == ';' && isdigit(str[n-1]) && isdigit(str[n+1]) && step == 0) {
			step = 1;
			continue;
		}
		if(str[n] == ':' && isdigit(str[n-1]) && isdigit(str[n+1]) && step < 3) {
			step = (step == 2 ? 3 : 2);
			continue;
		}
		
		return 0;
	}
	
	if(step < 2) {
		return 0;
	}
	
	if(strchr(str, ';')) {
		days = atoi(str);
		str = strchr(str, ';')+1;
		
		if(days > 127 || days < 1) {
			return 0;
		}
	}else if(days == 0) {
		days = 127;
	}
	
	hour = atoi(str);
	str = strchr(str, ':')+1;
	min = atoi(str);
	
	if(step == 3) {
		str = strchr(str, ':')+1;
		sec = atoi(str);
	}
	
	if(hour < 0 || hour > 23 || min < 0 || min > 59 || sec < 0 || sec > 59) {
		return 0;
	}
	
	if(sec) {
		sprintf(sdstring, "%02d:%02d:%02d,", hour, min, sec);
	}else{
		sprintf(sdstring, "%02d:%02d,", hour, min);
	}
	
	if(days == 127) {
		strcat(sdstring, " Every day");
	}else{
		if(days & (1<<1)) { strcat(sdstring, " Mon"); }
		if(days & (1<<2)) { strcat(sdstring, " Tue"); }
		if(days & (1<<3)) { strcat(sdstring, " Wed"); }
		if(days & (1<<4)) { strcat(sdstring, " Thu"); }
		if(days & (1<<5)) { strcat(sdstring, " Fri"); }
		if(days & (1<<6)) { strcat(sdstring, " Sat"); }
		if(days & (1<<0)) { strcat(sdstring, " Sun"); }
	}
	
	SendMessage(sdlist, LB_ADDSTRING, 0, (LPARAM)sdstring);
	return 1;
}

/* Validate new settings and write them to the registry */
static void update_cfg(HWND hwnd) {
	int server_cbox, warning_cbox, abort_cbox, delay_cbox;
	char server_name[256], server_share[256], sdtime[64], sdtimes[1024];
	int interval, cfg_warning, cfg_delay;
	int sdcount, sdnum, days, hour, min, sec;
	SERVICE_STATUS svc_status;
	DWORD errnum;
	
	server_cbox = gui_checked(hwnd, CHK_SERVER);
	warning_cbox = gui_checked(hwnd, CHK_WARNING);
	abort_cbox = gui_checked(hwnd, CHK_ABORT);
	delay_cbox = gui_checked(hwnd, CHK_DELAY);
	
	SendDlgItemMessage(hwnd, TXT_HOST, WM_GETTEXT, 256, (LPARAM)server_name);
	SendDlgItemMessage(hwnd, TXT_SHARE, WM_GETTEXT, 256, (LPARAM)server_share);
	
	interval = gui_pint(hwnd, TXT_INTERVAL);
	cfg_warning = gui_pint(hwnd, TXT_WARNING);
	cfg_delay = gui_pint(hwnd, TXT_DELAY);
	
	if(interval == -1) {
		MessageBox(hwnd, "Interval must be a non-negative integer", NULL, MB_OK);
		return;
	}
	if(warning_cbox && cfg_warning < 1) {
		MessageBox(hwnd, "Warning time must be a positive integer", NULL, MB_OK);
		return;
	}
	if(delay_cbox && cfg_delay < 1) {
		MessageBox(hwnd, "Delay time must be a positive integer", NULL, MB_OK);
		return;
	}
	
	sdcount = SendMessage(sdlist, LB_GETCOUNT, 0, 0);
	sdtimes[0] = '\0';
	
	for(sdnum = 0; sdnum < sdcount; sdnum++) {
		SendMessage(sdlist, LB_GETTEXT, sdnum, (LPARAM)sdtime);
		
		days = 0;
		hour = atoi(sdtime);
		min = atoi(sdtime+3);
		sec = (sdtime[5] == ':' ? atoi(sdtime+6) : 0);
		
		if(strstr(sdtime, "Mon")) { days |= (1<<1); }
		if(strstr(sdtime, "Tue")) { days |= (1<<2); }
		if(strstr(sdtime, "Wed")) { days |= (1<<3); }
		if(strstr(sdtime, "Thu")) { days |= (1<<4); }
		if(strstr(sdtime, "Fri")) { days |= (1<<5); }
		if(strstr(sdtime, "Sat")) { days |= (1<<6); }
		if(strstr(sdtime, "Sun")) { days |= (1<<0); }
		
		sprintf(sdtime, "%d;%d:%d:%d ", days, hour, min, sec);
		strcat(sdtimes, sdtime);
	}
	
	if(!server_cbox) {
		server_name[0] = '\0';
	}
	if(!warning_cbox) {
		cfg_warning = 0;
	}
	if(!delay_cbox) {
		cfg_delay = 0;
	}
	
	if(
		!reg_set_string("ServerName", server_name) ||
		!reg_set_string("ServerShare", server_share) ||
		!reg_set_dword("ServerInterval", interval) ||
		
		!reg_set_string("ShutdownTimes", sdtimes) ||
		!reg_set_dword("Warning", cfg_warning) ||
		!reg_set_dword("EnableAbort", abort_cbox) ||
		!reg_set_dword("DelayTime", cfg_delay)
	) {
		return;
	}
	
	if(!ControlService(svc, SERVICE_CONTROL_STOP, &svc_status)) {
		errnum = GetLastError();
		
		if(errnum != ERROR_SERVICE_NOT_ACTIVE) {
			msgboxf(MB_OK, NULL, "Failed to stop service: %s", w32_error(GetLastError()));
			return;
		}
	}
	
	while(svc_status.dwCurrentState != SERVICE_STOPPED) {
		Sleep(100);
		
		if(!QueryServiceStatus(svc, &svc_status)) {
			msgboxf(MB_OK, NULL, "Failed to query service: %s", w32_error(GetLastError()));
			return;
		}
	}
	
	/* Make sure the process has actually terminated */
	Sleep(100);
	
	if(!StartService(svc, 0, NULL)) {
		msgboxf(MB_OK, NULL, "Failed to start service: %s", w32_error(GetLastError()));
		return;
	}
}