1 /* 2 * Assoc.C - assoc internal command. 3 * 4 * 5 * History: 6 * 7 * 14-Mar-2009 Lee C. Baker 8 * - initial implementation 9 * 10 * 15-Mar-2009 Lee C. Baker 11 * - Don't write to (or use) HKEY_CLASSES_ROOT directly 12 * - Externalize strings 13 * 14 * TODO: 15 * - PrintAllAssociations might could be optimized to not fetch all registry subkeys under 'Classes', just the ones that start with '.' 16 * - Make sure that non-administrator users can list associations, and get appropriate error messages when they don't have sufficient 17 * privileges to perform an operation 18 */ 19 20 #include "precomp.h" 21 22 #ifdef INCLUDE_CMD_ASSOC 23 24 static INT 25 PrintAssociation(LPTSTR extension) 26 { 27 DWORD return_val; 28 HKEY hKey = NULL, hInsideKey = NULL; 29 30 DWORD fileTypeLength = 0; 31 LPTSTR fileType = NULL; 32 33 return_val = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0, KEY_READ, &hKey); 34 35 if (return_val != ERROR_SUCCESS) 36 return -1; 37 38 return_val = RegOpenKeyEx(hKey, extension, 0, KEY_READ, &hInsideKey); 39 RegCloseKey(hKey); 40 41 if (return_val != ERROR_SUCCESS) 42 return 0; 43 44 /* obtain string length */ 45 return_val = RegQueryValueEx(hInsideKey, NULL, NULL, NULL, NULL, &fileTypeLength); 46 47 if (return_val == ERROR_FILE_NOT_FOUND) /* no default value, don't display */ 48 { 49 RegCloseKey(hInsideKey); 50 return 0; 51 } 52 53 if (return_val != ERROR_SUCCESS) 54 { 55 RegCloseKey(hInsideKey); 56 return -2; 57 } 58 59 fileType = cmd_alloc(fileTypeLength * sizeof(TCHAR)); 60 if (!fileType) 61 { 62 WARN("Cannot allocate memory for fileType!\n"); 63 RegCloseKey(hInsideKey); 64 return -2; 65 } 66 67 /* obtain actual file type */ 68 return_val = RegQueryValueEx(hInsideKey, NULL, NULL, NULL, (LPBYTE)fileType, &fileTypeLength); 69 RegCloseKey(hInsideKey); 70 71 if (return_val != ERROR_SUCCESS) 72 { 73 cmd_free(fileType); 74 return -2; 75 } 76 77 if (fileTypeLength != 0) /* if there is a default key, display relevant information */ 78 { 79 ConOutPrintf(_T("%s=%s\n"), extension, fileType); 80 } 81 82 cmd_free(fileType); 83 return 1; 84 } 85 86 static INT 87 PrintAllAssociations(VOID) 88 { 89 DWORD return_val = 0; 90 HKEY hKey = NULL; 91 DWORD numKeys = 0; 92 93 DWORD extLength = 0; 94 LPTSTR extName = NULL; 95 DWORD keyCtr = 0; 96 97 return_val = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0, KEY_READ, &hKey); 98 99 if (return_val != ERROR_SUCCESS) 100 return -1; 101 102 return_val = RegQueryInfoKey(hKey, NULL, NULL, NULL, &numKeys, &extLength, NULL, NULL, NULL, NULL, NULL, NULL); 103 104 if (return_val != ERROR_SUCCESS) 105 { 106 RegCloseKey(hKey); 107 return -2; 108 } 109 110 extLength++; 111 extName = cmd_alloc(extLength * sizeof(TCHAR)); 112 if (!extName) 113 { 114 WARN("Cannot allocate memory for extName!\n"); 115 RegCloseKey(hKey); 116 return -2; 117 } 118 119 for (keyCtr = 0; keyCtr < numKeys; keyCtr++) 120 { 121 DWORD buffer_size = extLength; 122 return_val = RegEnumKeyEx(hKey, keyCtr, extName, &buffer_size, NULL, NULL, NULL, NULL); 123 124 if (return_val == ERROR_SUCCESS || return_val == ERROR_MORE_DATA) 125 { 126 if (*extName == _T('.')) 127 PrintAssociation(extName); 128 } 129 else 130 { 131 cmd_free(extName); 132 RegCloseKey(hKey); 133 return -1; 134 } 135 } 136 137 RegCloseKey(hKey); 138 139 cmd_free(extName); 140 return numKeys; 141 } 142 143 static INT 144 AddAssociation(LPTSTR extension, LPTSTR type) 145 { 146 DWORD return_val; 147 HKEY hKey = NULL, insideKey = NULL; 148 149 return_val = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0, KEY_ALL_ACCESS, &hKey); 150 151 if (return_val != ERROR_SUCCESS) 152 return -1; 153 154 return_val = RegCreateKeyEx(hKey, extension, 0, NULL, REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, NULL, &insideKey, NULL); 155 RegCloseKey(hKey); 156 157 if (return_val != ERROR_SUCCESS) 158 return -1; 159 160 return_val = RegSetValueEx(insideKey, NULL, 0, REG_SZ, (LPBYTE)type, (_tcslen(type) + 1) * sizeof(TCHAR)); 161 RegCloseKey(insideKey); 162 163 if (return_val != ERROR_SUCCESS) 164 return -2; 165 166 return 0; 167 } 168 169 static int 170 RemoveAssociation(LPTSTR extension) 171 { 172 DWORD return_val; 173 HKEY hKey; 174 175 return_val = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0, KEY_ALL_ACCESS, &hKey); 176 177 if (return_val != ERROR_SUCCESS) 178 return -1; 179 180 return_val = RegDeleteKey(hKey, extension); 181 RegCloseKey(hKey); 182 183 if (return_val != ERROR_SUCCESS) 184 return -2; 185 186 return 0; 187 } 188 189 190 INT CommandAssoc (LPTSTR param) 191 { 192 /* print help */ 193 if (!_tcsncmp (param, _T("/?"), 2)) 194 { 195 ConOutResPaging(TRUE,STRING_ASSOC_HELP); 196 return 0; 197 } 198 199 nErrorLevel = 0; 200 201 if (_tcslen(param) == 0) 202 { 203 PrintAllAssociations(); 204 } 205 else 206 { 207 LPTSTR lpEqualSign = _tcschr(param, _T('=')); 208 if (lpEqualSign != NULL) 209 { 210 LPTSTR fileType = lpEqualSign + 1; 211 LPTSTR extension = cmd_alloc((lpEqualSign - param + 1) * sizeof(TCHAR)); 212 if (!extension) 213 { 214 WARN("Cannot allocate memory for extension!\n"); 215 error_out_of_memory(); 216 return 1; 217 } 218 219 _tcsncpy(extension, param, lpEqualSign - param); 220 extension[lpEqualSign - param] = _T('\0'); 221 222 if (_tcslen(fileType) == 0) 223 /* if the equal sign is the last character 224 in the string, then delete the key */ 225 { 226 RemoveAssociation(extension); 227 } 228 else 229 /* otherwise, add the key and print out the association*/ 230 { 231 AddAssociation(extension, fileType); 232 PrintAssociation(extension); 233 } 234 235 cmd_free(extension); 236 } 237 else 238 { 239 /* no equal sign, print all associations */ 240 INT retval = PrintAssociation(param); 241 242 if (retval == 0) /* if nothing printed out */ 243 ConOutResPrintf(STRING_ASSOC_ERROR, param); 244 } 245 } 246 247 return 0; 248 } 249 250 #endif /* INCLUDE_CMD_ASSOC */ 251