/*
 * Copyright (c) 2004 Andy Polyakov <appro@fy.chalmers.se>
 *
 * Build with:
 *
 * cl -Ox -GD -GF -Zl -MD -LD environ.c advapi32.lib kernel32.lib ntdll.lib
 *
 * Pre-load as:
 *
 * [HKEY_CURRENT_USER\Software\Microsoft\Windows NT\CurrentVersion\DLL_PRELOAD\WFICA32.EXE]
 * @="environ.dll"
 *
 * [HKEY_CURRENT_USER\Software\Microsoft\Windows NT\CurrentVersion\DLL_PRELOAD\WFICA32.EXE\environ.dll]
 * "_CLUSTER_NETWORK_NAME_"=hex(2):25,00,55,00,53,00,45,00,52,00,4e,00,41,00,4d,\
 *   00,45,00,25,00,00,00
 *
 * 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,er")
#pragma comment(linker,"/merge:.rdata=.text")

#define UNICODE
#define _UNICODE

#if defined(WIN32) && !defined(_WIN32)
#define _WIN32
#endif
#define _WIN32_WINNT 0x0500

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

#define SIZEOF_A(a) (sizeof(a)/sizeof(a[0]))

/*
 * Implement some trivia in order to be excused from linking with MSVCRT.
 */
static void _memmove (void *dst, void *src, size_t n)
{ unsigned char *d=dst,*s=src;
    while (n--) *d++ = *s++;
  return;
}

static size_t _lstrlen (TCHAR *str)
{ int len=0;
    while (*str) { str++, len++; }
  return len;
}

static TCHAR *_lstrrchr (TCHAR *str,TCHAR c)
{ TCHAR *p = NULL;
    while (*str) { if (*str == c) p = str; str++; }
  return p;
}

static TCHAR *_lstrchr (TCHAR *str,TCHAR c)
{ TCHAR *p = NULL;
    while (*str) { if (*str == c) { p = str; break; } str++; }
  return p;
}

static TCHAR *_lstrncpy (TCHAR *dst,TCHAR *src,size_t n)
{ TCHAR *ret=dst;
    while(--n && *src) { *dst++ = *src++; }
    *dst=_T('\0');
  return ret;
}

static void _setenv_by_key (HKEY hkey)
{ DWORD nv,maxnamelen,namelen,maxvaluelen,valuelen,type,ret;
  TCHAR *name,*value,expanded[1024];

    if (RegQueryInfoKey (hkey,NULL,NULL,NULL,NULL,NULL,NULL,
			&nv,&maxnamelen,&maxvaluelen,
			NULL,NULL) != ERROR_SUCCESS
	|| nv==0)
	return;

    name  = _alloca ((maxnamelen+1)*sizeof(TCHAR));	if (name==NULL) return;
    value = _alloca (maxvaluelen);			if (value==NULL) return;

    while (nv--)
    {	namelen  = maxnamelen+1;
	valuelen = maxvaluelen;
	if (RegEnumValue (hkey,nv,name,&namelen,NULL,&type,(void *)value,&valuelen)
		!= ERROR_SUCCESS)
	    break;

	if (type!=REG_SZ && type!=REG_EXPAND_SZ) continue;

	if (type==REG_EXPAND_SZ)
	{   ret = ExpandEnvironmentStrings (value,expanded,SIZEOF_A(expanded));
	    if (ret==0 || ret>SIZEOF_A(expanded)) continue;
	}

	SetEnvironmentVariable (name,type==REG_EXPAND_SZ?expanded:value);
    }
}

#define MY_REGKEY \
        _T("Software\\Microsoft\\Windows NT\\CurrentVersion\\DLL_PRELOAD")

BOOL WINAPI DllMain (HINSTANCE h, DWORD reason, LPVOID junk)
{ TCHAR		str[256],kname[256],*s;
  size_t	l;
  HKEY		hkey;

    while (reason==DLL_PROCESS_ATTACH)
    {	DisableThreadLibraryCalls(h);

	_lstrncpy (kname,MY_REGKEY,SIZEOF_A(kname)-1),
	l=_lstrlen (kname), kname[l++]=_T('\\');

	/* append executable name... */
	if (!GetModuleFileName (NULL,str,SIZEOF_A(str))) break;
	(s = _lstrrchr (str,_T('\\'))) ? s++ : (s=str);
	_lstrncpy (kname+l,s,SIZEOF_A(kname)-l-1),
	l=_lstrlen (kname), kname[l++]=_T('\\');

	/* append this module name... */
	if (!GetModuleFileName (h,str,SIZEOF_A(str))) break;
	(s = _lstrrchr (str,_T('\\'))) ? s++ : (s=str);
	_lstrncpy (kname+l,s,SIZEOF_A(kname)-l);

	if (RegOpenKeyEx (HKEY_CURRENT_USER,kname,0,KEY_QUERY_VALUE,&hkey)
		== ERROR_SUCCESS)
	    _setenv_by_key (hkey), RegCloseKey (hkey);

	if (RegOpenKeyEx (HKEY_LOCAL_MACHINE,kname,0,KEY_QUERY_VALUE,&hkey)
		== ERROR_SUCCESS)
	    _setenv_by_key (hkey), RegCloseKey (hkey);

	break;
    }

  return TRUE;
}