registry.c

Virtual Floppy Driver

/*
	registry.c

	Generic registry manipulating functions
    Copyright (C) 2003 Kenji Kato
*/

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>

#define _CRTDBG_MAP_ALLOC
#include <stdlib.h>
#include <crtdbg.h>

#include "registry.h"

//
//	local functions
//
static DWORD ReadRegStrValue(LPCTSTR subkey, LPCTSTR value, LPTSTR buf);
static DWORD DeleteRegistryTree(HKEY hKey, LPCTSTR subkey);
static DWORD CopyRegistryTree(HKEY hKeyTo, HKEY hKeyFrom);

//
//	Get associated program name for extension
//
DWORD GetAssociatedProgram(LPCTSTR ext, LPTSTR prog)
{
	char	subkey[MAX_PATH];
	char	buf[MAX_PATH];
	DWORD	ret;
	char	*top, *tail = NULL;

	*prog = '\0';

	//	Read [HKEY_CLASSES_ROOT\.ext] -- file type
	ret = ReadRegStrValue(ext, NULL, subkey);

	if (ret != ERROR_SUCCESS) {
		return ret == ERROR_FILE_NOT_FOUND ? ERROR_SUCCESS : ret;
	}

	//	Read [HKEY_CLASSES_ROOT\<file type>\shell] -- default verb

	strcat(subkey, "\\shell");

	ret = ReadRegStrValue(subkey, NULL, buf);

	if (ret != ERROR_SUCCESS && ret != ERROR_FILE_NOT_FOUND) {
		return ret;
	}

	//	Read [HKEY_CLASSES_ROOT\<file type>\shell\<verb>\command] -- command line

	sprintf(subkey + strlen(subkey), "\\%s\\command", buf[0] ? buf : "open");

	if ((ret = ReadRegStrValue(subkey, NULL, buf)) != ERROR_SUCCESS) {
		return ret == ERROR_FILE_NOT_FOUND ? ERROR_SUCCESS : ret;
	}

	//	Extract program name from string

	if (buf[0] == '\"') {
		// program path is quoted 
		tail = &buf[1];

		// search trailing quote
		while (*tail != '\"' && *tail != '\0') {
			tail++;
		}

		// find top of program file name
		top = tail;

		while (top > buf && *(top - 1) != '\\' && *(top - 1) != '\"') {
			top--;
		}
	}
	else {
		BOOL found = FALSE;

		// program path is not quoted
		tail = buf;

		// first search for extension

		while (*tail) {

			// search blank or end of string
			while (*tail != '\0' && *tail != ' ') {
				tail++;
			}

			if (tail - buf >= 4) {
				if (!strnicmp(tail - 4, ".exe", 4) ||
					!strnicmp(tail - 4, ".cmd", 4) ||
					!strnicmp(tail - 4, ".bat", 4) ||
					!strnicmp(tail - 4, ".com", 4) ||
					!strnicmp(tail - 4, ".vbs", 4) ||
					!strnicmp(tail - 4, ".vbe", 4) ||
					!strnicmp(tail - 4, ".jse", 4) ||
					!strnicmp(tail - 4, ".wsf", 4) ||
					!strnicmp(tail - 4, ".wsh", 4))
				{
					found = TRUE;
					break;
				}
			}
			else if (tail - buf >= 3) {
				if (!strnicmp(tail - 3, ".js",  3)) {
					found = TRUE;
					break;
				}
			}

			tail++;
		}

		//	extension was not found in the text, so switch to simpler way

		if (!found) {
			tail = buf;

			// search first blank

			while (*tail != ' ' && *tail != '\0') {
				tail++;
			}
		}

		// find top of the program file name

		top = tail;

		while (top > buf && *(top - 1) != '\\') {
			top --;
		}
	}

	//	copy program file name to output buffer
	*tail = '\0';
	strcpy(prog, top);

	return ERROR_SUCCESS;
}

//
//	read a string value from a registry key
//
DWORD ReadRegStrValue(LPCTSTR subkey, LPCTSTR value, LPTSTR buf)
{
	HKEY hKey;
	DWORD size = MAX_PATH;
	DWORD type;
	LONG ret;

	*buf = '\0';

	ret = RegOpenKeyEx(HKEY_CLASSES_ROOT, subkey, 0, KEY_QUERY_VALUE, &hKey);

	if (ret != ERROR_SUCCESS) {
		if (ret != ERROR_FILE_NOT_FOUND) {
			_RPTF1(_CRT_WARN,
				"ReadRegStrValue : RegOpenKeyEx(\"%s\")\n",
				subkey);
		}

		return ret;
	}

	ret = RegQueryValueEx(hKey, value, NULL, &type, (LPBYTE)buf, &size);

	RegCloseKey(hKey);

	if (ret == ERROR_SUCCESS) {
		if (type == REG_SZ || type == REG_MULTI_SZ || type == REG_EXPAND_SZ) {
			// ensure that the string is properly terminated

			if (*(buf + size)) {
				*(buf + size++) = '\0';
			}
			return ERROR_SUCCESS;
		}
		else {
			_RPTF2(_CRT_WARN,
				"ReadRegStrValue : RegQueryValueEx(\"[%s] %s\") returned non-text data\n",
				subkey, value ? value : "@");

			return ERROR_INVALID_DATA;
		}
	}
	else if (ret != ERROR_FILE_NOT_FOUND) {
		_RPTF2(_CRT_WARN,
			"ReadRegStrValue : RegQueryValueEx(\"[%s] %s\")\n",
			subkey, value ? value : "@");
	}
	
	return ret;
}

//
//	remove current association and restore
//	the backed up association if present
//
DWORD RestoreAssociation(LPCTSTR ext, LPCTSTR backup_value)
{
	TCHAR current[MAX_PATH], backup[MAX_PATH];
	DWORD size;
	HKEY hKey;
	LONG ret;

	//
	//	Get current filetype and backup filetype
	//

	ret = RegOpenKeyEx(HKEY_CLASSES_ROOT, ext, 0, KEY_ALL_ACCESS, &hKey);

	if (ret != ERROR_SUCCESS) {
		_RPTF1(_CRT_WARN,
			"RemoveAssociation : RegOpenKeyEx(\"HKCR\\%s\")\n",
			ext);

		return ret;
	}

	// read current filetype

	size = sizeof(current);

	ret = RegQueryValueEx(hKey, NULL, NULL, NULL, (LPBYTE)current, &size);

	if (ret == ERROR_FILE_NOT_FOUND) {
		size = 0;
	}
	else if (ret != ERROR_SUCCESS) {
		_RPTF1(_CRT_WARN,
			"RemoveAssociation : RegQueryValueEx(\"[HKCR\\%s] @\")\n",
			ext);

		RegCloseKey(hKey);
		return ret;
	}

	// make sure that the string is terminated

	if (current[size]) {
		current[size] = '\0';
	}

	// read backed up filetype

	if (backup_value && *backup_value) {
		size = sizeof(backup);

		ret = RegQueryValueEx(hKey, backup_value, NULL, NULL, (LPBYTE)backup, &size);

		if (ret == ERROR_FILE_NOT_FOUND) {
			size = 0;
		}
		else if (ret != ERROR_SUCCESS) {
			_RPTF2(_CRT_WARN,
				"RemoveAssociation : RegQueryValueEx(\"[HKCR\\%s] %s\")\n",
				ext, backup_value);

			RegCloseKey(hKey);
			return ret;
		}

		// make sure that the string is terminated

		if (backup[size]) {
			backup[size++] = '\0';
		}

	}
	else {
		backup[0] = '\0';
	}

	if (backup[0]) {
		//
		// buckup filetype exists so revert to it
		//

		ret = RegSetValueEx(hKey, NULL, 0, REG_SZ, (LPBYTE)backup, size);

		if (ret != ERROR_SUCCESS) {
			_RPTF2(_CRT_WARN,
				"RemoveAssociation : RegSetValueEx(\"[HKCR\\%s] @ = %s\")\n",
				ext, backup);

			RegCloseKey(hKey);
			return ret;
		}

		// delte backup value

		ret = RegDeleteValue(hKey, backup_value);

		if (ret != ERROR_SUCCESS) {
			_RPTF2(_CRT_WARN,
				"RemoveAssociation : RegDeleteValue(\"[HKCR\\%s] %s\")\n",
				ext, backup_value);

			RegCloseKey(hKey);
			return ret;
		}
	}
	else {
		// simply delete current filetype

		ret = RegDeleteValue(hKey, NULL);

		if (ret != ERROR_SUCCESS) {
			_RPTF1(_CRT_WARN,
				"RemoveAssociation : RegDeleteValue(\"[HKCR\\%s] @\")\n",
				ext);

			RegCloseKey(hKey);
			return ret;
		}
	}

	RegCloseKey(hKey);

	if (backup[0] == '\0') {

		// either backup filetype does not exist or backup value was not
		// specified, so delete the extension key

		ret = RegDeleteKey(HKEY_CLASSES_ROOT, ext);

		if (ret != ERROR_SUCCESS && ret != ERROR_ACCESS_DENIED) {
			_RPTF1(_CRT_WARN,
				"RemoveAssociation : RegDeleteKey(\"[HKCR\\%s]\")\n",
				ext);

			// this is not really fatal
		}
	}

	if (current[0]) {

		// successfully reverted or removed extention key, so now
		// delte filetype subkey

		return DeleteRegistryTree(HKEY_CLASSES_ROOT, current);
	}

	return ERROR_SUCCESS;
}

//
//	recursively delete a registry subtree
//
DWORD	 DeleteRegistryTree(HKEY hKey, LPCTSTR subkey)
{
	HKEY hSubKey;
	LONG ret;
	DWORD keysize;

	//	open subkey to delete

	ret = RegOpenKeyEx(hKey, subkey, 0, KEY_READ, &hSubKey);

	if (ret != ERROR_SUCCESS) {
		_RPTF1(_CRT_WARN,
			"DeleteRegistryTree : RegOpenKeyEx(\"%s\")\n", subkey);

		return ret;
	}

	//	get maximum subkey name length

	ret = RegQueryInfoKey(hSubKey, NULL, NULL, NULL, NULL, &keysize, NULL, NULL, NULL, NULL, NULL, NULL);

	if (ret != ERROR_SUCCESS) {
		_RPTF1(_CRT_WARN,
			"DeleteRegistryTree : RegOpenKeyEx(\"%s\")\n", subkey);

		goto cleanup;
	}
		
	if (keysize) {
		//	allocate buffer for subkey name

		LPTSTR keyname = (LPTSTR)malloc(++keysize);

		if (!keyname) {
			ret = GetLastError();

			_RPTF1(_CRT_WARN,
				"DeleteRegistryTree : malloc(%lu)\n", keysize);

			goto cleanup;
		}

		// enumerate each sub-subkey

		for (;;) {
			ret = RegEnumKey(hSubKey, 0, keyname, keysize);
			
			if (ret == ERROR_NO_MORE_ITEMS) {
				break;
			}
			else if (ret != ERROR_SUCCESS) {
				_RPTF1(_CRT_WARN,
					"DeleteRegistryTree(\"%s\"): RegEnumKey\n", subkey);
				break;
			}

			// delete sub-subkey

			ret = DeleteRegistryTree(hSubKey, keyname);
			
			if (ret != ERROR_SUCCESS) {
				_RPTF2(_CRT_WARN,
					"DeleteRegistryTree(\"%s\"): DeleteRegistryTree(\"%s\")\n",
					subkey, keyname);
				break;
			}
		}

		free(keyname);
	}

cleanup:
	// close target subkey

	RegCloseKey(hSubKey);

	//	delete target subkey

	if (ret == ERROR_SUCCESS || ret == ERROR_NO_MORE_ITEMS) {
		ret = RegDeleteKey(hKey, subkey);

		if (ret != ERROR_SUCCESS) {
			_RPTF1(_CRT_WARN,
				"DeleteRegistryTree : RegDeleteKey(\"%s\")\n",
				subkey);
		}
	}

	return ret;
}

//
//	Add filetype association
//	-- preserve current association as much as possible
//
DWORD AddAssociation(
	LPCTSTR ext,
	LPCTSTR type_prefix,
	LPCTSTR type_desc,
	LPCTSTR verb,
	LPCTSTR verb_desc,
	LPCTSTR program,
	int icon_idx)
{
	HKEY hExtKey, hNewKey, hSub;
	char filetype[MAX_PATH], backup[MAX_PATH], buf[MAX_PATH];