/*
 * Copyright (c) 2003 Andy Polyakov <appro@fy.chalmers.se>
 *
 * Build with:
 *
 * cl -Ox -GD -GF -Zl -MD -LD WINLOGOUT.c advapi32.lib kernel32.lib ntdll.lib
 *
 * See http://fy.chalmers.se/~appro/nt/DLL_PRELOAD/ for further details.
 *
 */
#ifndef _DLL
#error "_DLL is not defined."
#endif

#pragma comment(linker,"/entry:DllMain@12")
#pragma comment(linker,"/section:.text,erw")
#pragma comment(linker,"/merge:.rdata=.text")
#pragma comment(linker,"/merge:.data=.text")

#if defined(WIN32) && !defined(_WIN32)
#define _WIN32
#endif

#define _WIN32_WINNT 0x0500

#include <stdio.h>
#include <tchar.h>
#include <windows.h>
#include <winbase.h>

/* M$ headers are *complete* mess!!! */
#ifndef STATUS_INFO_LENGTH_MISMATCH
#define STATUS_INFO_LENGTH_MISMATCH	((NTSTATUS)0xC0000004L)
#endif

#define NtQuerySystemInformation _NtQuerySystemInformation
#include <winternl.h>
#undef NtQuerySystemInformation

#define SystemHandleInformation 16

typedef LONG KPRIORITY;

typedef struct {
    ULONG NextEntryOffset;
    ULONG NumberOfThreads;
    LARGE_INTEGER SpareLi1;
    LARGE_INTEGER SpareLi2;
    LARGE_INTEGER SpareLi3;
    LARGE_INTEGER CreateTime;
    LARGE_INTEGER UserTime;
    LARGE_INTEGER KernelTime;
    UNICODE_STRING ImageName;
    KPRIORITY BasePriority;
    HANDLE UniqueProcessId;
    HANDLE InheritedFromUniqueProcessId;
    ULONG HandleCount;
    ULONG SessionId;
    ULONG SpareUl3;
    SIZE_T PeakVirtualSize;
    SIZE_T VirtualSize;
    ULONG PageFaultCount;
    ULONG PeakWorkingSetSize;
    ULONG WorkingSetSize;
    SIZE_T QuotaPeakPagedPoolUsage;
    SIZE_T QuotaPagedPoolUsage;
    SIZE_T QuotaPeakNonPagedPoolUsage;
    SIZE_T QuotaNonPagedPoolUsage;
    SIZE_T PagefileUsage;
    SIZE_T PeakPagefileUsage;
    SIZE_T PrivatePageCount;
    LARGE_INTEGER ReadOperationCount;
    LARGE_INTEGER WriteOperationCount;
    LARGE_INTEGER OtherOperationCount;
    LARGE_INTEGER ReadTransferCount;
    LARGE_INTEGER WriteTransferCount;
    LARGE_INTEGER OtherTransferCount;
} MORE_SYSTEM_PROCESS_INFORMATION;

typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO {
    USHORT UniqueProcessId;
    USHORT CreatorBackTraceIndex;
    UCHAR ObjectTypeIndex;
    UCHAR HandleAttributes;
    USHORT HandleValue;
    PVOID Object;
    ULONG GrantedAccess;
} SYSTEM_HANDLE_TABLE_ENTRY_INFO, *PSYSTEM_HANDLE_TABLE_ENTRY_INFO;

typedef struct _SYSTEM_HANDLE_INFORMATION {
    ULONG NumberOfHandles;
    SYSTEM_HANDLE_TABLE_ENTRY_INFO Handles[ 1 ];
} SYSTEM_HANDLE_INFORMATION, *PSYSTEM_HANDLE_INFORMATION;

NTSTATUS NTAPI NtQuerySystemInformation (
    IN SYSTEM_INFORMATION_CLASS SystemInformationClass,
    OUT PVOID SystemInformation,
    IN ULONG SystemInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

typedef struct _OBJECT_TYPE_INFORMATION {
    UNICODE_STRING TypeName;
    ULONG TotalNumberOfObjects;
    ULONG TotalNumberOfHandles;
    ULONG TotalPagedPoolUsage;
    ULONG TotalNonPagedPoolUsage;
    ULONG TotalNamePoolUsage;
    ULONG TotalHandleTableUsage;
    ULONG HighWaterNumberOfObjects;
    ULONG HighWaterNumberOfHandles;
    ULONG HighWaterPagedPoolUsage;
    ULONG HighWaterNonPagedPoolUsage;
    ULONG HighWaterNamePoolUsage;
    ULONG HighWaterHandleTableUsage;
    ULONG InvalidAttributes;
    GENERIC_MAPPING GenericMapping;
    ULONG ValidAccessMask;
    BOOLEAN SecurityRequired;
    BOOLEAN MaintainHandleCount;
    ULONG PoolType;
    ULONG DefaultPagedPoolCharge;
    ULONG DefaultNonPagedPoolCharge;
} OBJECT_TYPE_INFORMATION, *POBJECT_TYPE_INFORMATION;

typedef struct _OBJECT_ALL_TYPES_INFORMATION {
    ULONG NumberOfTypes;
    OBJECT_TYPE_INFORMATION TypeInformation[1];
} OBJECT_ALL_TYPES_INFORMATION, *POBJECT_ALL_TYPES_INFORMATION;

typedef enum _OBJECT_INFORMATION_CLASS {
    ObjectBasicInformation,
    ObjectNameInformation,
    ObjectTypeInformation,
    ObjectAllTypesInformation,
    ObjectHandleInformation
} OBJECT_INFORMATION_CLASS;

NTSTATUS NTAPI NtQueryObject (
    IN HANDLE ObjectHandle,
    IN OBJECT_INFORMATION_CLASS ObjectInformationClass,
    OUT PVOID ObjectInformation,
    IN ULONG ObjectInformationLength,
    OUT PULONG ReturnLength OPTIONAL
    );

static DWORD find_csrpid()
{ SIZE_T RequiredLength;
  NTSTATUS s;
  static DWORD csrpid=0;

    if (csrpid) return csrpid;

    csrpid=(DWORD)-1;

    /* Dump processes and find this session's CSRSS.EXE */
    {	MORE_SYSTEM_PROCESS_INFORMATION *ProcTable,*Proc;
	DWORD sess;

	RequiredLength=64*1024; /* a bit large for ws, but ok for ts */ 
	ProcTable = VirtualAlloc(NULL,RequiredLength,MEM_COMMIT,PAGE_READWRITE);
	if (ProcTable == NULL) return -1;

	while ((s=NtQuerySystemInformation(SystemProcessInformation,
			ProcTable,
			(ULONG)RequiredLength,NULL))
			== STATUS_INFO_LENGTH_MISMATCH)
	{   RequiredLength += 16*1024;
	    VirtualFree(ProcTable,0,MEM_RELEASE);
            ProcTable = VirtualAlloc (NULL,RequiredLength,
			MEM_COMMIT,PAGE_READWRITE);
	    if (ProcTable == NULL) return -1;
	}
	if (s) return s;

	ProcessIdToSessionId (GetCurrentProcessId(),&sess);

	Proc=ProcTable;
	do {
	    if (Proc->SessionId == sess &&
		Proc->ImageName.Length==18 &&
		!_wcsnicmp(Proc->ImageName.Buffer,L"CSRSS.EXE",9)
	       )
	    {   csrpid=(DWORD)Proc->UniqueProcessId; break;   }
	    Proc=(MORE_SYSTEM_PROCESS_INFORMATION *)
			((char *)Proc+Proc->NextEntryOffset);
	} while (Proc->NextEntryOffset);

	VirtualFree (ProcTable,0,MEM_RELEASE);
    }

  return csrpid;
}

static NTSTATUS list_keys(DWORD csrpid,void (*cb)(DWORD,void *),void *cbdata)
{ SIZE_T RequiredLength;
  UCHAR  KeyType=0;
  SYSTEM_HANDLE_INFORMATION	*HandleTable=NULL;
  NTSTATUS s;
  int i;

    /* Retrieve Object Type table and find Key's index */
    {	OBJECT_ALL_TYPES_INFORMATION	*Types;
	OBJECT_TYPE_INFORMATION		*t;

	s=NtQueryObject(NULL,ObjectAllTypesInformation,NULL,0,&RequiredLength);

	if (s==STATUS_INFO_LENGTH_MISMATCH)
	{   Types = VirtualAlloc(NULL,RequiredLength,MEM_COMMIT,PAGE_READWRITE);
	    if (Types == NULL) return -1;

	    if (s=NtQueryObject(NULL,ObjectAllTypesInformation,
			Types,RequiredLength,&RequiredLength))
	    {	VirtualFree (Types,0,MEM_RELEASE);
		return s;
	    }

	    for (t=Types->TypeInformation,i=0;i<Types->NumberOfTypes;i++)
	    {   if (t->TypeName.Length==6 &&
		    !_wcsnicmp(t->TypeName.Buffer,L"Key",3))
		{   KeyType = i+1; break;   }
		t=(OBJECT_TYPE_INFORMATION *)(
			(char *)t->TypeName.Buffer +
			((t->TypeName.MaximumLength+3)&~3) );
	    }
	}
	else return s;

	VirtualFree (Types,0,MEM_RELEASE);
    }
    if (KeyType==0) return -1;

    /* Retrieve Handle Table */
    {	SYSTEM_HANDLE_INFORMATION Handle;

	s=NtQuerySystemInformation(SystemHandleInformation,
		&Handle,sizeof(Handle),(ULONG *)&RequiredLength);

	while (s==STATUS_INFO_LENGTH_MISMATCH)
	{   RequiredLength += 4*1024;
	    if (HandleTable) VirtualFree (HandleTable,0,MEM_RELEASE);
	    HandleTable = VirtualAlloc(NULL,RequiredLength,MEM_COMMIT,PAGE_READWRITE);
	    if (HandleTable == NULL)
		return -1;
	    s=NtQuerySystemInformation(SystemHandleInformation,
			HandleTable,(ULONG)RequiredLength,(ULONG *)&RequiredLength);
	}
    }
    if (s)
    {	VirtualFree (HandleTable,0,MEM_RELEASE);
	return s;
    }

    for (i=0;i<HandleTable->NumberOfHandles;i++)
    {	if (HandleTable->Handles[i].UniqueProcessId == csrpid)
	{   if (HandleTable->Handles[i].ObjectTypeIndex==KeyType)
		(*cb)(HandleTable->Handles[i].HandleValue,cbdata);
	}
    }

    VirtualFree (HandleTable,0,MEM_RELEASE);

  return 0;
}

typedef struct { HANDLE csr,log; } CB_DATA;

static void closeremotehandle (DWORD handle,CB_DATA *d)
{ HANDLE hcopy=NULL,me=GetCurrentProcess();
  NTSTATUS s;
  struct { UNICODE_STRING Name; wchar_t Buffer[96]; } NameQuery;

    while (d->csr)
    {	if (DuplicateHandle (d->csr,(HANDLE)handle,me,&hcopy,0,0,DUPLICATE_SAME_ACCESS) == 0)
	    break;
	s=NtQueryObject(hcopy,ObjectNameInformation,&NameQuery,sizeof(NameQuery),NULL);
        CloseHandle(hcopy);
	if (s && s!=STATUS_INFO_LENGTH_MISMATCH) break;

	if (NameQuery.Name.Length>=34 &&
	    !_wcsnicmp(NameQuery.Name.Buffer,L"\\REGISTRY\\User\\S-",17))
	{   DuplicateHandle (d->csr,(HANDLE)handle,GetCurrentProcess(),&hcopy,0,0,
			DUPLICATE_CLOSE_SOURCE|DUPLICATE_SAME_ACCESS);
            CloseHandle(hcopy);
#if 0
	    {	TCHAR str[256],*arr[]={str};
	    	_snprintf(str,sizeof(str)/sizeof(str[0])-1,
				"WINLOGOUT.dll: Closing %x, %*S",
				handle,NameQuery.Name.Length,NameQuery.Name.Buffer);
		str[sizeof(str)/sizeof(str[0])-1]=0;
		ReportEvent(d->log,EVENTLOG_INFORMATION_TYPE,0,1000,NULL,1,0,arr,NULL);
	    }
#endif
	}
	break;
    }
}

typedef BOOL (WINAPI *UnloadUserProfile_T) (
    IN HANDLE Token,
    IN HANDLE Profile
    );

static UnloadUserProfile_T _UnloadUserProfile=NULL,*__UnloadUserProfile=NULL;

static BOOL WINAPI UnloadUserProfile_ (
    IN HANDLE Token,
    IN HANDLE Profile
    )
{ static int once=1;

    if (once)
    {	HANDLE csr,applog;
	CB_DATA d;

	once=0;
	applog=OpenEventLog(NULL,_T("Application"));

	if (csr=OpenProcess(PROCESS_DUP_HANDLE,0,find_csrpid()))
	{   d.csr=csr;
	    d.log=applog;
	    list_keys(find_csrpid(),closeremotehandle,&d);
	    CloseHandle(csr);
	}
	else
	{   TCHAR str[64],*arr[]={str};
	    _snprintf(str,sizeof(str)/sizeof(str[0])-1,
			"WINLOGOUT.dll: Unable to OpenProcess(%d)",find_csrpid());
	    str[sizeof(str)/sizeof(str[0])-1]=0;
	    ReportEvent(applog,EVENTLOG_ERROR_TYPE,0,1000,NULL,1,0,arr,NULL);
	}
	CloseEventLog (applog);
    }

  return (*_UnloadUserProfile)(
            Token,
            Profile);
}

static int _lstricmp(const char *s1, const char *s2)
{ char c1,c2;
  int  ret;

    while (c1=*s1, c2=*s2, c1&&c2)
    {   c1|=0x20, c2|=0x20; /* lower the case */
        if (ret=c1-c2) return ret;
        s1++, s2++;
    }

  return c1-c2;
}

BOOL WINAPI DllMain (HINSTANCE h, DWORD reason, LPVOID junk)
{ DWORD                     acc;
  HMODULE                   hmod;
  IMAGE_DOS_HEADER         *dos_header;
  IMAGE_NT_HEADERS         *nt_headers;
  IMAGE_DATA_DIRECTORY     *dir;
  IMAGE_IMPORT_DESCRIPTOR  *idesc;
  IMAGE_THUNK_DATA         *thunk;

    switch (reason)
    { case DLL_PROCESS_ATTACH:

        DisableThreadLibraryCalls(h);

        if (!(hmod=GetModuleHandle(_T("USERENV.DLL"))))
        {   OutputDebugString(_T("USERENV.DLL not found?"));
            return FALSE;
        }

        _UnloadUserProfile=(UnloadUserProfile_T)GetProcAddress(hmod,"UnloadUserProfile");

        if (!(hmod=GetModuleHandle(NULL)))
        {   OutputDebugString(_T("NULL not found?"));
            return FALSE;
        }

        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);

        dir=&nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        idesc=(IMAGE_IMPORT_DESCRIPTOR *)((char *)hmod + dir->VirtualAddress);
        while (idesc->Name)
        {   if (!_lstricmp((char *)hmod+idesc->Name,"USERENV.DLL")) break;
            idesc++;
        }

        if (!idesc->Name)
        {   OutputDebugString(_T("Can't locate USERENV.DLL import descriptor"));
            return FALSE;
        }

        for (thunk=(IMAGE_THUNK_DATA *)((char *)hmod+idesc->FirstThunk);
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)_UnloadUserProfile;
             thunk++) ;
        __UnloadUserProfile=(UnloadUserProfile_T *)thunk;

        if (!VirtualProtect (__UnloadUserProfile,sizeof(void*),PAGE_EXECUTE_READWRITE,&acc))
        {   OutputDebugString(_T("Unable to unlock Thunk Table"));
            return FALSE;
        }
        *__UnloadUserProfile=UnloadUserProfile_;
        VirtualProtect (__UnloadUserProfile,sizeof(void*),acc,&acc);

        break;

      case DLL_PROCESS_DETACH:

        VirtualProtect (__UnloadUserProfile,sizeof(void*),PAGE_EXECUTE_READWRITE,&acc);
        *__UnloadUserProfile=_UnloadUserProfile;
        VirtualProtect (__UnloadUserProfile,sizeof(void*),acc,&acc);

        break;
    }

  return TRUE;
}