/**********************************************************************
 *                                                                    *
 * INITLO.C, version 0.1                                              *
 *                                                                    *
 * This program for Windows NT logs off the user who initiates it.    *
 * Behavior is the same as if the user clicks                         *
 * Start -> Shut Down... -> Close all programs and log on			  *
 *                          as a different user?                      *
 * except that non-responding tasks or DOS windows will be forced     *
 * to close (due to the EWX_FORCE switch in ExitWindowsEx).           *
 * It is designed to accompany INITLO, which initializes a log off    *
 * from the SYSTEM account, which displays a warning message box      *
 * before log off. If the user does not close it, the box will        *
 * remain visible indefinitely. Therefore LOGOFF looks for this       *
 * window (with title 'System Message') and sends a WM_CLOSE message  *
 * to it.                                                             *
 *                                                                    *
 * Copyright (C) 1997 by Alexander Frink                              *
 *                       Hermann Schauss Str. 8                       *
 *                       65232 Taunusstein                            *
 *                       Germany                                      *
 *                       Alexander.Frink@Uni-Mainz.DE                 *
 *                                                                    *
 * 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     *
 * of the License, 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 (file name: LICENSE) for more details.  *
 *                                                                    *
 * You should have received a copy of the GNU General Public License  * 
 * along with this program; if not, write to the Free Software        *
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.          *
 *                                                                    *
 **********************************************************************/

#include <windows.h>
#include <lm.h>
#include <stdio.h>
#include <stdlib.h>
#include <process.h>
#include <tchar.h>
#include <string.h>

#include "tasklist.h"
#include "privs.h"
#include "error.h"

#define DEFAULT_GRACE                0
#define DEFAULT_PROGRAM_TO_EXECUTE   "logoff.exe"
#define DEFAULT_WARNING_TEXT         "Please log off within %i seconds!"
#define DEFAULT_WARNING_TEXT_ESCAPED "Please log off within %%i seconds!"

#define REGKEY_WINLOGON     "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Winlogon"
#define REGSUBKEY_WINLOGON_SHELL  "Shell"

#define DOMAIN_AND_ACCOUNT_LENGTH 200
#define SZ_SD_BUF                 100 
#define MAXPATHLEN                1024

typedef struct _UINFO {
  UCHAR szAccountName[DOMAIN_AND_ACCOUNT_LENGTH];
  UCHAR szDomainName[DOMAIN_AND_ACCOUNT_LENGTH];
  BYTE  bHours[21];
  HANDLE hShellToken;
} UINFO, * PUINFO;


UCHAR                ucEvrSDBuf[SZ_SD_BUF] = ""; 
PSECURITY_DESCRIPTOR psdEveryoneSD         = (PSECURITY_DESCRIPTOR)&ucEvrSDBuf; 

int dwGrace;  // grace time in seconds between warning and logoff
UCHAR szWarningText[1024];
UCHAR szProgramToExecute[MAXPATHLEN];

UINFO UInfo;

int iNumShellsToWatch=3;
UCHAR szShellsToWatchArray[3][50] = { "explorer.exe",
                                      "progman.exe",
                                      "winfile.exe" };

BOOL
GetNameOfWinlogonShell(
    CHAR *ShellBuf,
    DWORD BufLen
    )

/*++

Routine Description:

    Retrieves the name of the shell program from the Winlogon Key

Arguments:

    ShellBuf          - a buffer for the name of the shell
    BufLen            - length of the buffer

Return Value:

    TRUE                    - success
    FALSE                   - failure

--*/

{
    DWORD                        rc;
    HKEY                         hKeyWinlogon;
    DWORD                        dwType;

    ShellBuf[0]=0;

    //
    // open the Winlogon key
    //
    rc = RegOpenKeyEx( HKEY_LOCAL_MACHINE,
                       REGKEY_WINLOGON,
                       0,
                       KEY_READ,
                       &hKeyWinlogon
                     );
    if (rc != ERROR_SUCCESS) {
        goto exit;
    }

    //
    // read the name of the shell from the registry
    //
    rc = RegQueryValueEx( hKeyWinlogon,
                          REGSUBKEY_WINLOGON_SHELL,
                          NULL,
                          &dwType,
                          ShellBuf,
                          &BufLen
                        );

    if (rc != ERROR_SUCCESS) {
        goto exit;
    }

exit:
    RegCloseKey(hKeyWinlogon);

    return ShellBuf[0]!=0;
}

BOOL
GetTokenToShell(PHANDLE hShellToken)
/*++

Routine Description:

    This routine gets a token of the currently running login shell
	(file name from the registry or standard file names).
    If no shell is running, it is assumed that no user is logged in
    interactively, and NULL is returned.

Arguments:

    hShellToken             - Token of the login shell
	                          (NULL if not running)

Return Value:

    TRUE                    - success (shell running or not running)
    FALSE                   - failure

--*/

{
#define BUFLEN 256
    DWORD Pid;
    CHAR ShellBuf[BUFLEN];
    HANDLE hShellProcess;
    int i;
    
    if (!GetNameOfWinlogonShell(ShellBuf,BUFLEN)) {
        ErrorHandler("GetTokenToShell: Error reading RegKey for Winlogon Shell");
		return FALSE;
    }
    if (!FindPid(ShellBuf,&Pid)) {
        // the Winlogon Shell is not running, try the list of possible shells instead
        for (i=0; i<iNumShellsToWatch; i++) {
            if (FindPid(szShellsToWatchArray[i],&Pid)) {
                break;
            }
        }
        if (i==iNumShellsToWatch) {
            // no shell seems to be running, no-one is logged on
		    *hShellToken=NULL;
		    return TRUE;
        }
	}
    hShellProcess=OpenProcess(PROCESS_ALL_ACCESS,
                              TRUE,
                              Pid);
    if (hShellProcess==NULL) {
        ErrorHandler("GetTokenToShell: OpenProcess for Shell failed");
        return FALSE;
    }
    if (!OpenProcessToken(hShellProcess,
                          TOKEN_ALL_ACCESS,
                          hShellToken)) {
        CloseHandle(hShellProcess);
        ErrorHandler("GetTokenToShell: OpenProcessToken for Shell failed");
        return FALSE;
    }

    CloseHandle(hShellProcess);
    return TRUE;
}

BOOL GetLoggedOnUserInformation(PUINFO UInfo)
{
    /*
        returns TRUE if everything is o.k.
                FALSE if an error occured
    */
    
    UCHAR InfoBuffer[1000];
    PTOKEN_USER pTokenUser = (PTOKEN_USER)InfoBuffer;
    DWORD dwInfoBufferSize;
    DWORD dwAccountSize=DOMAIN_AND_ACCOUNT_LENGTH;
    DWORD dwDomainSize=DOMAIN_AND_ACCOUNT_LENGTH;
    SID_NAME_USE snu;
    PUSER_INFO_11 pUserInfo11;
    WCHAR  wszAccountName[DOMAIN_AND_ACCOUNT_LENGTH]; // Unicode user name
    WCHAR  wszDomainName[DOMAIN_AND_ACCOUNT_LENGTH];
    LPBYTE pDCName; // Domain Controller
    int i;

    // initialize
    UInfo->szAccountName[0]=0;
    UInfo->szDomainName[0]=0;
    for (i=0; i<21; i++) UInfo->bHours[i]=255;

    if (!GetTokenToShell(&(UInfo->hShellToken))) {
		return FALSE;
	}
    if (UInfo->hShellToken==NULL) {
        // no-one logged on
		return TRUE;
	}
    if (!GetTokenInformation(UInfo->hShellToken,
							 TokenUser,
							 InfoBuffer,
							 1000,
							 &dwInfoBufferSize)) {
        ErrorHandler("GetLoggedOnUserInformation: GetTokenInformation failed");
		return FALSE;
	}
	if (!LookupAccountSid(NULL,
                          pTokenUser->User.Sid,
                          UInfo->szAccountName,
                          &dwAccountSize,
                          UInfo->szDomainName,
                          &dwDomainSize,
                          &snu)) {
        ErrorHandler("GetLoggedOnUserInformation: LookupAccountSid failed");
        return FALSE;
    }

    MultiByteToWideChar(CP_ACP,
                        0,
                        UInfo->szAccountName, 
                        strlen(UInfo->szAccountName)+1,
                        wszAccountName,
                        sizeof(wszAccountName));

    MultiByteToWideChar(CP_ACP,
                        0,
                        UInfo->szDomainName, 
                        strlen(UInfo->szDomainName)+1,
                        wszDomainName,
                        sizeof(wszDomainName));

	if (NetGetAnyDCName(NULL,wszDomainName,&pDCName)!=NERR_Success)
	{
        //ErrorHandler("GetLoggedOnUserInformation: NetGetAnyDCName failed");
        NetApiBufferFree((LPBYTE)pUserInfo11);
        //return FALSE;
    }
	else
	{
		if (NetUserGetInfo((LPWSTR) pDCName,
						   (LPWSTR) wszAccountName,
						   11,
						   (LPBYTE *)&pUserInfo11)!=NERR_Success) {
			ErrorHandler("GetLoggedOnUserInformation: NetUserGetInfo failed");
			NetApiBufferFree((LPBYTE)pUserInfo11);
			NetApiBufferFree(pDCName);
			//return FALSE;
		}
		else
		{
			for (i=0; i<21; i++) {
				UInfo->bHours[i]=pUserInfo11->usri11_logon_hours[i];
			}
			NetApiBufferFree((LPBYTE)pUserInfo11);
			NetApiBufferFree(pDCName);
		}
	}
    return TRUE;
}

BOOL LogOffUser(PUINFO UInfo)
{
    HANDLE hShellTokenDuplicate;
    SECURITY_ATTRIBUTES     sa; 
    STARTUPINFO si; 
    PROCESS_INFORMATION pi;
    UCHAR szPathAndProgramToExecute[MAXPATHLEN];
    UCHAR *p;
	if (!SetPrivilege(NULL,SE_ASSIGNPRIMARYTOKEN_NAME,TRUE)) {
        ErrorHandler("LogOffUser: SetPrivilege SE_ASSIGNPRIMARYTOKEN_NAME");
		return FALSE;
	}
	if (!SetPrivilege(NULL,SE_CREATE_TOKEN_NAME,TRUE)) {
        ErrorHandler("LogOffUser: SetPrivilege SE_CREATE_TOKEN_NAME");
		return FALSE;
	}
	if (!SetPrivilege(NULL,SE_SECURITY_NAME,TRUE)) {
        ErrorHandler("LogOffUser: SetPrivilege SE_SECURITY_NAME");
		return FALSE;
	}
    if (!InitializeSecurityDescriptor(psdEveryoneSD,SECURITY_DESCRIPTOR_REVISION)) 
    { 
        ErrorHandler("LogOffUser: InitializeSecurityDescriptor"); 
    } 
    if (!SetSecurityDescriptorDacl(psdEveryoneSD, 
                                   TRUE,      // fDaclPresent flag 
                                   (PACL)NULL, 
                                   FALSE))    // not a default DACL 
    {
        ErrorHandler("LogOffUser: SetSecurityDescriptorDacl"); 
    } 
    sa.nLength = sizeof(sa); 
    sa.lpSecurityDescriptor = psdEveryoneSD; 
    sa.bInheritHandle = FALSE;
    if (!DuplicateTokenEx(UInfo->hShellToken,
		                  TOKEN_ALL_ACCESS,
						  &sa,
						  SecurityImpersonation,
						  TokenPrimary,
						  &hShellTokenDuplicate)) {
        ErrorHandler("LogOffUser: DuplicateTokenEx");
		return FALSE;
	}

    // cf. code samples/win32/winnt/security/rcmd/server.c
    si.cb = sizeof(STARTUPINFO); 
    si.lpReserved = NULL; 
    si.lpDesktop = "WinSta0\\Default"; 
    si.lpTitle = NULL; 
    si.cbReserved2 = 0; 
    si.lpReserved2 = NULL; 
    si.dwFlags = 0; 
    si.hStdOutput = NULL; 
    si.hStdError = NULL; 
    si.hStdInput = NULL; 

    printf("trying %s\n",szProgramToExecute);
    
    if (!CreateProcessAsUser(hShellTokenDuplicate,szProgramToExecute,NULL, // command line
                             NULL, // security descriptor for process
                             NULL, // security descriptor for thread
                             FALSE,NORMAL_PRIORITY_CLASS,NULL, // environment block
                             NULL, // current directory
                             &si,&pi)) 
	{
        // if this fails, try the current path
        if (GetModuleFileName(NULL,szPathAndProgramToExecute,MAXPATHLEN)!=0)
		{
            p=strrchr(szPathAndProgramToExecute,'\\');
            if (p!=NULL) 
				p[1]='\0';
            strcat(szPathAndProgramToExecute,szProgramToExecute);
            printf("trying %s\n",szPathAndProgramToExecute);
            if (!CreateProcessAsUser(hShellTokenDuplicate,
                                     szPathAndProgramToExecute,
                                     NULL, // command line
                                     NULL, // security descriptor for process
                                     NULL, // security descriptor for thread
                                     FALSE,
                                     NORMAL_PRIORITY_CLASS,
                                     NULL, // environment block
                                     NULL, // current directory
                                     &si,
                                     &pi)) 
			{
                ErrorHandler("LogOffUser: CreateProcessAsUser (probably program to execute not found)");
	            return FALSE;
            }
			else
			{
				ErrorHandler("LogOffUser: CreateProcessAsUser (OKOKOK)");
			}
        }
		else
		{
            ErrorHandler("LogOffUser: CreateProcessAsUser (probably program to execute not found)");
	        return FALSE;
        }
	}

    // CloseHandle(hShellToken);
    CloseHandle(hShellTokenDuplicate);
    
    return TRUE;
}

DWORD WINAPI WarningThread(LPVOID dummy)
{
    TCHAR szOut[256];

    _stprintf(szOut,szWarningText,dwGrace);
    //MessageBox(NULL,szOut,TEXT("System Message"),MB_SERVICE_NOTIFICATION);
	//MessageBox(NULL,szOut,TEXT("System Message"),MB_OK);
    ExitThread(0);
    return 0;
}

void DisplayWarningMessage(void)
{
    HANDLE hThread;
    DWORD dwThreadId;

    hThread=CreateThread(NULL,
                         0,
                         WarningThread,
                         0,
                         0,
                         &dwThreadId);
    Sleep(5000);
    TerminateThread(hThread,0);
    CloseHandle(hThread);
}

void Usage(char *cmdname)
{
    printf("Usage: %s [-g grace] [-p program] [-t text]\n",cmdname);
    printf("       -g grace       set time in seconds between warning and logoff\n"
           "                      (default: %i, 0=immediate logoff)\n",DEFAULT_GRACE);
    printf("       -p program     program to execute after grace period\n"
           "                      (default: " DEFAULT_PROGRAM_TO_EXECUTE ")\n"
           "       -t text        text to display in warning message\n"
           "                      (default: " DEFAULT_WARNING_TEXT_ESCAPED "\n\n"
           "       (C) 1997 by Alexander Frink (Alexander.Frink@Uni-Mainz.DE)\n"
           "       All rights reserved.\n");
    exit(1);
}

void ProcessCmdLine(int argc, char *argv[])
{
    char c, *p, *arg0;

    // algorithm for command line processing taken from SDK sample 
    // Win32/WinNT/Floppy/MFMT.C

    arg0=argv[0];

    if ( argc > 1 ) { 
        while (--argc > 0 ) { 
            p = *++argv; 
            if (*p == '/' || *p == '-') { 
                while (c = *++p) {
                    switch (tolower(c)) { 
                    case '?': 
                        Usage(arg0);
                        break; 
                    case 'g': 
                        if (argc<2) Usage(arg0);
                        argc--, argv++; 
                        dwGrace=atoi(*argv);
                        break;
                    case 'p': 
                        if (argc<2) Usage(arg0);
                        argc--, argv++; 
                        strcpy(szProgramToExecute,*argv); 
                        break;
                    case 't': 
                        if (argc<2) Usage(arg0);
                        argc--, argv++; 
                        strcpy(szWarningText,*argv); 
                        break;
                    default: 
                        Usage(arg0); 
                        break; 
                    } 
                } 
            }
        } 
    } 
}


int main(int argc, char *argv[])
{
    dwGrace=DEFAULT_GRACE;
    strcpy(szWarningText,DEFAULT_WARNING_TEXT);
    strcpy(szProgramToExecute,DEFAULT_PROGRAM_TO_EXECUTE);

    ProcessCmdLine(argc,argv);
    
    if (GetLoggedOnUserInformation(&UInfo)) 
	{
        if (dwGrace!=0) DisplayWarningMessage();
        Sleep(dwGrace*1000); // 1000 milliseconds per second
        if (GetLoggedOnUserInformation(&UInfo))
		{
            LogOffUser(&UInfo);
            ErrorHandler("User has been forced to log off.");
        }
		else
		{
            ErrorHandler("User has deliberately logged off during grace period.");
        }
    } 
	else 
	{
        ErrorHandler("No-one is currently logged on to your system.");
    }
    return 0;
}

