/*
This program is distributed under the GPL which is available at www.gnu.org
*/

#include <stdio.h>
#include <stdlib.h>
#include <winsock2.h>
#include <windows.h>
#include <winnt.h>
#include <time.h> /* to get time */
#include <io.h> /* to create/write to log file */
#include <fcntl.h> /* to create/write to log file */
#include <sys/stat.h> /* to create/write to log file */

/* protocol information */
#define CURRENT_MAGIC "powerd-2.0"
#define ACT_OK		1
#define ACT_FAIL	0
#define ACT_ASK		254
/* this is the status code that is returned when select times out */
#define TIMEDOUT 255

struct powerd_message {
	unsigned char cmd;
	char magic[15];
};

/* function prototypes */
SOCKET initfd(char * host,unsigned short port);
int poll_server(SOCKET fd);
int waitformsg(SOCKET fd, long sec);
void take_action(unsigned char status);
void log_msg(const char *errormsg);
void powerfail();
void powerback();
int getshutdownpriv();
DWORD WINAPI waitforshutdown(LPVOID param);
DWORD WINAPI showmsg(LPVOID param);
void WINAPI service_main(void);
VOID WINAPI service_ctrl(DWORD dwCtrlCode);
VOID service_stop();
void mainsvc(void);
void error(char *msg);
/* */

/* the name of the program */
#define APPNAME TEXT("ntpowerd-1.0")
#define SERVICENAME TEXT("NTPowerd")

/* other stuff */
#define SERVICEFLAG TEXT("-service")

/* configuration */
/* where to store the logfile */
#define LOGFILE TEXT("c:/powerd.log")
/* name/port of powerd server */
#define DEFHOST TEXT("charger.brianweb.local")
#define PORT 2043
/* how long to wait between notify of power failure and actual shutdown */
#define SECONDS 180
/* seconds between polls of server */
#define POLLINTERVAL 240
/* how many times to try to connect to server on failed poll */
#define RETRIES 10
/* seconds between retries */
#define POLLRETRY 3
/* */

/* Global Vars */
char *servername;
SOCKET connection=0;
int logfd;
unsigned char current_status = ACT_OK;
unsigned char iamntsvc = 0;

/* these for wait thread in 95 */
HANDLE ThreadHandle;
DWORD ThreadID;

/* these for service on NT */
SERVICE_STATUS_HANDLE StatusHandle;
SERVICE_STATUS ServiceStatus;
/* */

int WINAPI WinMain( HINSTANCE hInstance,
HINSTANCE hPrevInstance, 
LPSTR lpCmdLine,
int nCmdShow)
{
	char *p;
	if((lpCmdLine[0] == '\0')||(lpCmdLine[0]=='-')){
		servername = malloc(strlen(DEFHOST)+1);
		if(!servername) error("Out of memory");
		strcpy(servername,DEFHOST);
		p = lpCmdLine;
	}
	else {
		p = strchr(lpCmdLine,' ');
		if(p) *p = '\0';
		servername = malloc(strlen(lpCmdLine)+1);
		if(!servername) error("Out of memory");
		strcpy(servername,lpCmdLine);
		if(p) p++; else p = lpCmdLine;
	}
	if(!strncmp(p, SERVICEFLAG, strlen(SERVICEFLAG))){
		if( (GetVersion() < 0x80000000) ){ /* NT */
			SERVICE_TABLE_ENTRY dispatchTable[] =
			{
				{SERVICENAME,(LPSERVICE_MAIN_FUNCTION)service_main},
				{NULL,NULL}
			};
			iamntsvc = 1;
			if(!StartServiceCtrlDispatcher(dispatchTable)) error("Error connecting to service ontrol manager.");
		}
		else { /* 95 */
			DWORD (WINAPI *lpfRegisterServiceProcess)(DWORD,DWORD);
			lpfRegisterServiceProcess = 
				(DWORD (WINAPI*)(DWORD,DWORD))
				GetProcAddress(GetModuleHandle("KERNEL32.DLL"),"RegisterServiceProcess");
			if(!lpfRegisterServiceProcess) error(0);
			if(!lpfRegisterServiceProcess(0,1)) error(0);
			mainsvc();
		}
	}
	else {
		mainsvc();
	}
	return(0);
}


void WINAPI service_main()
{	
	StatusHandle = RegisterServiceCtrlHandler(SERVICENAME,service_ctrl);

	ServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
	ServiceStatus.dwCurrentState = SERVICE_START_PENDING;
	ServiceStatus.dwControlsAccepted = 0;
	ServiceStatus.dwWin32ExitCode = 0;
	ServiceStatus.dwServiceSpecificExitCode = 0;
	ServiceStatus.dwCheckPoint = 0;
	ServiceStatus.dwWaitHint = 4000;

	if(!SetServiceStatus(StatusHandle, &ServiceStatus))
		error("SetServiceStatus - in service_main()");
	mainsvc();
}

VOID WINAPI service_ctrl(DWORD dwCtrlCode){
    switch(dwCtrlCode){
		case SERVICE_CONTROL_STOP:
			ServiceStatus.dwCurrentState = SERVICE_STOP_PENDING;
			ServiceStatus.dwControlsAccepted = 0;
			ServiceStatus.dwCheckPoint = 0;
			ServiceStatus.dwWaitHint = 2000;
			if(connection){
				shutdown(connection,SD_BOTH);
				closesocket(connection);
			}
			break;
		default:
			break;
	}
	SetServiceStatus(StatusHandle, &ServiceStatus);
}

VOID service_stop(){
	log_msg("Shutdown");
	ServiceStatus.dwCurrentState = SERVICE_STOPPED;
	SetServiceStatus(StatusHandle, &ServiceStatus);
}

void mainsvc()
{
	SOCKET fd;
	WSADATA wsadata;
	unsigned char status;
	unsigned char failcount=0;

	logfd = _open(LOGFILE,_O_WRONLY|_O_CREAT|_O_APPEND,_S_IWRITE);
	if(logfd == -1) error("Error opening log file.");

	log_msg("Starting up...");

	if(GetVersion() < 0x80000000){ /* NT */
		if(!getshutdownpriv()) error("Error getting shutdown privileges.");
	}

	if(WSAStartup(MAKEWORD(2,0), &wsadata) != 0)
		error("Error initializing winsock");

	fd = initfd((char*)servername,PORT);
	status = poll_server(fd);
	if(status==TIMEDOUT){
		error("Failed to connect to server.");
	}
	else{
		log_msg("Contacted server.");
		if(iamntsvc){
			ServiceStatus.dwCurrentState = SERVICE_RUNNING;
			ServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
			ServiceStatus.dwCheckPoint = 0;
			if(!SetServiceStatus(StatusHandle, &ServiceStatus))
					error("SetServiceStatus - in mainsvc()");
		}
		take_action(status);
	}
	connection = fd;
	/* main loop */
	while (1){
		status = waitformsg(fd,POLLINTERVAL);
		if(ServiceStatus.dwCurrentState == SERVICE_STOP_PENDING) break;
		if(status==TIMEDOUT){
			while (1){
				status = poll_server(fd);
				if(status==TIMEDOUT){
					log_msg("Failed to reconnect to server.");
					failcount++;
					if(failcount > 2) error("Connection with server lost");
				}
				else {failcount = 0;take_action(status); break;}
			}
			
		}
		else take_action(status);
	}
	if(iamntsvc) service_stop();

}

void take_action(unsigned char status){
	if( (status==ACT_OK) && (current_status != ACT_OK) ){
			current_status = ACT_OK;
			powerback();
	}
	else if( (status==ACT_FAIL) && (current_status != ACT_FAIL) ){
			current_status = ACT_FAIL;	
			powerfail();
	}
}

SOCKET initfd(char * host, unsigned short port){
	SOCKET fd;
	struct hostent *hostptr;
	struct sockaddr_in remote;

	hostptr = gethostbyname(host);
	if(!hostptr){
		log_msg("error in gethostbyname()");
		log_msg(host);
		exit(1);
	}
	ZeroMemory(&remote,sizeof(remote));
	CopyMemory(&remote.sin_addr, hostptr->h_addr, hostptr->h_length);
	remote.sin_family = AF_INET;
	remote.sin_port = htons(port);
	fd = socket(AF_INET, SOCK_DGRAM, 0);
	if(fd==-1){
		log_msg("error in socket()");
		exit(1);
	}
	if(connect(fd, (struct sockaddr *) &remote, sizeof(remote)) != 0){
		log_msg("error in connect()");
		exit(1);
	}
	return(fd);
}
int poll_server(SOCKET fd) {
	FD_SET set;
	struct powerd_message buf;
	struct timeval to = {POLLRETRY,0};
	int i = 0;
	log_msg("polling server...");

	FD_ZERO(&set);
	strcpy(buf.magic, CURRENT_MAGIC);
	buf.cmd = ACT_ASK; 
	for(i=0;i<RETRIES;i++){
		send(fd, (char *)&buf, sizeof(buf), 0);
		FD_SET(fd, &set);
		if(select(0, &set,0,0,&to)){
			ZeroMemory(&buf, sizeof(buf));
			if(recv(fd, (char *)&buf, sizeof(buf),0) == sizeof(buf))
				return(buf.cmd);
			else i--;
		}
		if(iamntsvc){
			ServiceStatus.dwCheckPoint++;
			ServiceStatus.dwWaitHint = (POLLRETRY + 1)*1000;
			if(!SetServiceStatus(StatusHandle, &ServiceStatus))
				error("SetServiceStatus - in poll_server()");
		}
		log_msg("Failed to connect... retrying...");
	}
	return TIMEDOUT;
}
int waitformsg(SOCKET fd, long sec){
	FD_SET set;
	struct powerd_message buf;
	struct timeval to;
	ZeroMemory(&to,sizeof(to));
	to.tv_sec = sec;

	FD_ZERO(&set); FD_SET(fd, &set);
	strcpy(buf.magic, CURRENT_MAGIC);
	while(1){
		if(select(0, &set,0,0,&to)){
			ZeroMemory(&buf, sizeof(buf));
			if(recv(fd, (char *)&buf, sizeof(buf),0) == sizeof(buf)){
					return(buf.cmd);
			}
			else break;
		}
		else break;
	}
	return TIMEDOUT;
}

void log_msg(const char * errormsg){
	char date[9],time[9];
	char *buf;

	buf = malloc(22+strlen(errormsg)+1); /* 9 + 9 + 4 = 22 */
	if(!buf){log_msg("Out of Memory"); exit(1);}
	_strdate(date);
	_strtime(time);
	sprintf(buf,"%s %s: %s\n", date,time,errormsg);
	_write(logfd,buf,strlen(buf));
	free(buf);
}

void powerfail() {
	long *sec;
	sec = malloc(sizeof(long));
	if(!sec){ log_msg("Out of memory"); exit(1); }
	*sec = SECONDS;
	log_msg("The power is failing!");
	if(GetVersion() < 0x80000000){
		if(!InitiateSystemShutdown(0,"The power has failed!",*sec,1,0)) error(0);
	}
	else {
		ThreadHandle = CreateThread(0,0,&waitforshutdown,sec,0, &ThreadID);
		if(!ThreadHandle) error(0);
	}

}

DWORD WINAPI waitforshutdown(LPVOID param){
	long *sec;
	sec = param;
	Sleep(*sec * 1000);
	if(!ExitWindowsEx(EWX_SHUTDOWN|EWX_FORCE,0)) error(0);
	return 0;
}

void powerback() {
	log_msg("The power is back on.");
	if(GetVersion() < 0x80000000){
		if(!AbortSystemShutdown(0)) error(0);
	}
	else {
		if(!TerminateThread(ThreadHandle,0)) error(0);
	}
}


int getshutdownpriv(){
	HANDLE prochandle;
	HANDLE accesstoken;
	LUID luid;
	TOKEN_PRIVILEGES tokenpriv;
	
	prochandle = GetCurrentProcess();
	if(!OpenProcessToken(prochandle,TOKEN_ADJUST_PRIVILEGES |TOKEN_QUERY,&accesstoken)) return(0);
	if(!LookupPrivilegeValue(0,SE_SHUTDOWN_NAME,&luid)) return(0);
	tokenpriv.PrivilegeCount=1;
	tokenpriv.Privileges[0].Luid = luid;
	tokenpriv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if(!AdjustTokenPrivileges(accesstoken,0,&tokenpriv,0,0,0)) return(0);

	return(1);
}

void error(char *msg)
{
	char *p;
	UINT flags=0;
	if(!msg){
		FormatMessage(
			FORMAT_MESSAGE_ALLOCATE_BUFFER|FORMAT_MESSAGE_FROM_SYSTEM,
			0,GetLastError(),
			MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
			(LPTSTR)&p,0,0);
	}
	else {
		p = msg;
	}
	if(iamntsvc) flags=0x00200000L; /* makes msgbox appear on active windowstation/desktop */
	MessageBox(0,p,"Error",flags);
	if(iamntsvc) service_stop();
	exit(1);
}