/*
 * awaked.c	awake daemon
 * Copyright(c)2022 by Hiroaki Sengoku <sengoku@gcd.org>
 * Version 0.0	Jun 18, 2022
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Emacs; see the file COPYING.  If not, write to
 * the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * Usage: awaked <opt>
 *   opt: -v                ; verbose
 *        -p <port>         ; listen port
 *        -M install <name> ; install service as <name>
 *        -M remove <name>  ; remove service <name>
 *
 * Make:
 *   cl awaked.c
 */
static char *CVS_ID =
"@(#) $Id: awaked.c,v 1.4 2022/07/02 13:29:53 sengoku Exp $";

#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "advapi32.lib")
#pragma comment(lib, "user32.lib")

#define STRMAX		256

WSADATA WSAData;
SERVICE_STATUS NTServiceStatus;
SERVICE_STATUS_HANDLE NTServiceStatusHandle;
#define NTServiceDisplayPrefix	"awake "
char *NTServiceDisplayName = NULL;
char *NTServiceName = NULL;
HANDLE NTServiceLog = NULL;
HANDLE NTServiceThreadHandle = NULL;

int Verbose = 0;
SOCKET WakerSd;

void waker(void) {
    for (;;) {
	EXECUTION_STATE es;
	fd_set fds;
	struct timeval tv;
	tv.tv_sec = 180;
	tv.tv_usec = 0;
	FD_ZERO(&fds);
	FD_SET(WakerSd, &fds);
	if (select((int)WakerSd+1, &fds, NULL, NULL, &tv) > 0) {
	    struct sockaddr_in from;
	    char buf[STRMAX+1];
	    int from_len = sizeof(from);
	    int len = recvfrom(WakerSd, buf, STRMAX, 0,
			       (struct sockaddr*)&from, &from_len);
	    if (len < 0) {
		fprintf(stderr, "RecvFrom failed err=%d\n", GetLastError());
		continue;
	    }
	    buf[len] = '\0';
	    es = SetThreadExecutionState(
		ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED | ES_CONTINUOUS);
	    if (Verbose) {
		fprintf(stderr, "es=%x %x(%d)%d:%s\n", es,
			ntohl((unsigned long)from.sin_addr.s_addr),
			ntohs(from.sin_port), len, buf);
	    }
	    if (0 == strncmp(buf, "es?", 3)) {
		len = snprintf(buf, STRMAX, "es=%x", es);
	    }
	    if (sendto(WakerSd, buf, len, 0,
		       (struct sockaddr*)&from, sizeof(from)) != len) {
		fprintf(stderr, "Sendto failed err=%d\n", GetLastError());
	    }
	} else {
	    es = SetThreadExecutionState(ES_CONTINUOUS);
	    if (Verbose > 1) {
		fprintf(stderr, "es=%x timeout\n", es);
	    }
	}
    }
}

size_t quoteToken(char *dst, size_t dlen, char *src) {
    char buf[STRMAX+1];
    size_t len;
    if (strchr(src, ' ')) {
	snprintf(buf, STRMAX, "\"%s\"", src);
	len = strlen(buf);
	if (dst) strncpy_s(dst, dlen, buf, len);
    } else {
	len = strlen(src);
	if (dst) strncpy_s(dst, dlen, src, len);
    }
    return len;
}

void installService(int argc, char *argv[]) {
    SC_HANDLE scManager;
    SC_HANDLE scService;
    char exeName[STRMAX+1];
    char *command;
    size_t len, j;
    int i;
    int state;
    char *p;
    if (!GetModuleFileName(0, exeName, sizeof(exeName))) {
	fprintf(stderr, "Can't determine exe name err=%d\n",
		(int)GetLastError());
	exit(1);
    }
    scManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scManager) {
	fprintf(stderr, "Can't open service control manager err=%d\n",
		(int)GetLastError());
	exit(1);
    }
    len = strlen(exeName);
    for (i=1; i < argc; i++) {
	len += 1 + quoteToken(NULL, 0, argv[i]);
    }
    len++;	/* for '\0' */
    command = malloc(len);
    if (!command) {
	fprintf(stderr, "Out of memory\n");
	exit(1);
    }
    strcpy_s(command, len, exeName);
    j = strlen(command);
    state = 0;
    for (i=1; i < argc; i++) {
	p = argv[i];
	switch(state) {
	case 0:
	    if (!strcmp(p, "-M")) state++;
	    break;
	case 1:
	    if (!strcmp(p, "install"))
		p = "run_svc";	/* assume same length */
	    break;
	}
	command[j++] = ' ';
	j += quoteToken(command+j, len-j, p);
    }
    command[j] = '\0';
    if (Verbose) fprintf(stderr, "install: %s\n", command);
    scService
	= CreateService(scManager, NTServiceName,
			NTServiceDisplayName,
			SERVICE_ALL_ACCESS, SERVICE_WIN32_OWN_PROCESS,
			SERVICE_DEMAND_START, SERVICE_ERROR_NORMAL,
			command, NULL, NULL, "TcpIp\0\0",
			NULL, NULL);
    if (!scService) {
	fprintf(stderr, "Can't install service: %s err=%d\n",
		NTServiceName, (int)GetLastError());
	CloseServiceHandle(scManager);
	exit(1);
    }
    printf("service installed: %s\n", NTServiceName);
    CloseServiceHandle(scService);
    CloseServiceHandle(scManager);
}

void removeService(void) {
    SC_HANDLE scManager;
    SC_HANDLE scService;
    scManager = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scManager) {
	fprintf(stderr, "Can't open service control manager err=%d\n",
		(int)GetLastError());
	exit(1);
    }
    scService = OpenService(scManager, NTServiceName,
			    SERVICE_ALL_ACCESS);
    if (!scService) {
	fprintf(stderr, "Can't open service: %s err=%d\n",
		NTServiceName, (int)GetLastError());
	CloseServiceHandle(scManager);
	exit(1);
    }
    if (ControlService(scService, SERVICE_CONTROL_STOP, &NTServiceStatus)) {
	do {
	    Sleep(1);
	} while (QueryServiceStatus(scService, &NTServiceStatus),
		 NTServiceStatus.dwCurrentState == SERVICE_STOP_PENDING);
	if (NTServiceStatus.dwCurrentState == SERVICE_STOPPED) {
	    printf("%s stopped\n", NTServiceName);
	} else {
	    fprintf(stderr, "failed to stop %s\n", NTServiceName);
	}
    }
    if (!DeleteService(scService)) {
	fprintf(stderr, "failed to remove service: %s err=%d\n",
		NTServiceName, (int)GetLastError());
	CloseServiceHandle(scService);
	CloseServiceHandle(scManager);
	exit(1);
    }
    CloseServiceHandle(scService);
    CloseServiceHandle(scManager);
    printf("service removed: %s\n", NTServiceName);
}

void addEventSource(char *name) {
    HKEY hk;
    char key[STRMAX+1];
    char exeName[STRMAX+1];
    DWORD data;
    snprintf(key, STRMAX, "SYSTEM\\CurrentControlSet\\Services\\EventLog\\Application\\%s", name);
    if (RegCreateKey(HKEY_LOCAL_MACHINE, key, &hk)) return;
    if (!GetModuleFileName(0, exeName, sizeof(exeName))) return;
    if (RegSetValueEx(hk, "EventMessageFile", 0, REG_EXPAND_SZ,
		      (BYTE*)exeName, (DWORD)strlen(exeName)+1)) return;
    data = (EVENTLOG_ERROR_TYPE | EVENTLOG_WARNING_TYPE |
	    EVENTLOG_INFORMATION_TYPE);
    if (RegSetValueEx(hk, "TypesSupported", 0, REG_DWORD,
		      (LPBYTE)&data, sizeof(DWORD))) return;
    RegCloseKey(hk);
}

void scReportStatus(DWORD curState, DWORD exitCode, DWORD hint) {
    static DWORD checkPoint = 1;
    if (curState == SERVICE_START_PENDING)
	NTServiceStatus.dwControlsAccepted = 0;
    else
	NTServiceStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    NTServiceStatus.dwCurrentState = curState;
    NTServiceStatus.dwWin32ExitCode = exitCode;
    NTServiceStatus.dwWaitHint = hint;
    if ((curState == SERVICE_RUNNING) || (curState == SERVICE_STOPPED))
	NTServiceStatus.dwCheckPoint = 0;
    else
	NTServiceStatus.dwCheckPoint = checkPoint++;
    SetServiceStatus(NTServiceStatusHandle, &NTServiceStatus);
}

void WINAPI serviceCtrl(DWORD code) {
    switch(code) {
    case SERVICE_CONTROL_STOP:
	scReportStatus(SERVICE_STOP_PENDING, NO_ERROR, 0);
	printf("Service stopping..");
	if (WaitForSingleObject(NTServiceThreadHandle, 1000) == WAIT_TIMEOUT)
	    TerminateThread(NTServiceThreadHandle, 0);
	break;
    default:
	break;
    }
}

DWORD WINAPI serviceThread(LPVOID lpParms) {
    UNREFERENCED_PARAMETER(lpParms);
    do {
	waker();
    } while (NTServiceStatus.dwCurrentState == SERVICE_RUNNING);
    ExitThread(0);
}

void WINAPI serviceMain(DWORD argc, LPTSTR *argv) {
    UNREFERENCED_PARAMETER(argc);
    UNREFERENCED_PARAMETER(argv);
    DWORD thid;
    NTServiceStatusHandle
	= RegisterServiceCtrlHandler(NTServiceName, serviceCtrl);
    if (!NTServiceStatusHandle) {
//	fprintf(stderr, "Can't register ServiceCtrlHandler\n");
	return;
    }
    NTServiceStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    NTServiceStatus.dwServiceSpecificExitCode = 0;
    scReportStatus(SERVICE_START_PENDING, NO_ERROR, 3000);
//  printf("Service started\n");
    scReportStatus(SERVICE_RUNNING, NO_ERROR, 0);
    NTServiceThreadHandle = CreateThread(0, 0, serviceThread, NULL, 0, &thid);
    if (NTServiceThreadHandle) {
	WaitForSingleObject(NTServiceThreadHandle, INFINITE);
	CloseHandle(NTServiceThreadHandle);
    }
//  printf("Service stopped\n");
    scReportStatus(SERVICE_STOPPED, NO_ERROR, 0);
}

void help(void) {
    fprintf(stderr,
	    "Usage: awaked <opt>\n"
	    "  opt: -v               ; verbose\n"
	    "       -p <port>        ; listen port\n"
	    "      -M install <name> ; install service as <name>\n"
	    "      -M remove <name>  ; remove service <name>\n"
	    );
    exit(1);
}

int main(int argc, char *argv[]) {
    int i;
    size_t len;
    struct sockaddr_in sin;
    memset((char *)&sin, 0, sizeof(sin));
    sin.sin_family = AF_INET;
    for (i=1; i < argc; i++) {
	char *p = argv[i];
	if (*p == '-') {
	    p++;
	    do {
		switch (*p) {
		case 'v':
		    Verbose++;
		    break;
		case 'p':
		    sin.sin_port = htons((u_short)atoi(argv[++i]));
		    break;
		case 'M':
		    if (argc <= i+2) {
			fprintf(stderr, "Illegal Option: -M without args\n");
			help();
		    }
		    NTServiceName = _strdup(argv[i+2]);
		    len = strlen(NTServiceName)
			+ strlen(NTServiceDisplayPrefix) + 1;
		    NTServiceDisplayName = malloc(len);
		    if (!NTServiceDisplayName) {
			fprintf(stderr, "Out of memory\n");
			exit(1);
		    }
		    strcpy_s(NTServiceDisplayName, len,
			     NTServiceDisplayPrefix);
		    strcat_s(NTServiceDisplayName, len, NTServiceName);
		    ++i;
		    if (!strcmp(argv[i], "install")) {
			installService(argc, argv);
			exit(0);
		    } else if (!strcmp(argv[i], "remove")) {
			removeService();
			exit(0);
		    } else if (!strcmp(argv[i], "run_svc")) {
			addEventSource(NTServiceName);
			NTServiceLog
			    = RegisterEventSource(NULL, NTServiceName);
		    } else {
			fprintf(stderr, "Illegal Option: -M %s %s\n",
				argv[i], argv[i+1]);
			exit(1);
		    }
		    i++;
		    break;
		default:
		    help();
		}
		p++;
	    } while(*p);
	} else {
	    help();
	}
    }
    if (Verbose) {
	fprintf(stderr, "verbose=%d\n", Verbose);
    }
    if (WSAStartup(MAKEWORD(1, 1), &WSAData)) {
	fprintf(stderr, "Can't find winsock\n");
	exit(1);
    }
    if ((WakerSd=socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) < 0) {
	fprintf(stderr, "Can't get socket.\n");
	exit(1);
    }
    if (sin.sin_port <= 0) {
	help();
    }
    if (bind(WakerSd, (struct sockaddr*)&sin, sizeof(sin)) < 0) {
	fprintf(stderr, "Can't bind.\n");
	exit(1);
    }
    if (NTServiceName) {
	SERVICE_TABLE_ENTRY dispatchTable[] =
	    {
		{ NTServiceName, (LPSERVICE_MAIN_FUNCTION)serviceMain },
		{ NULL, NULL }
	    };
	if (!StartServiceCtrlDispatcher(dispatchTable))
	    fprintf(stderr, "StartServiceCtrlDispatcher failed\n");
	return 0;
    }
    waker();
}

/*
  For Gnu Emacs.
  Local Variables:
  tab-width: 8
  c-basic-offset: 4
  End:
*/