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 could be optimized to not fetch all registry subkeys under 'Classes', just the ones that start with '.'
16 */
17
18 #include "precomp.h"
19
20 #ifdef INCLUDE_CMD_ASSOC
21
22 static LONG
PrintAssociationEx(IN HKEY hKeyClasses,IN PCTSTR pszExtension)23 PrintAssociationEx(
24 IN HKEY hKeyClasses,
25 IN PCTSTR pszExtension)
26 {
27 LONG lRet;
28 HKEY hKey;
29 DWORD dwFileTypeLen = 0;
30 PTSTR pszFileType;
31
32 lRet = RegOpenKeyEx(hKeyClasses, pszExtension, 0, KEY_QUERY_VALUE, &hKey);
33 if (lRet != ERROR_SUCCESS)
34 {
35 if (lRet != ERROR_FILE_NOT_FOUND)
36 ErrorMessage(lRet, NULL);
37 return lRet;
38 }
39
40 /* Obtain the string length */
41 lRet = RegQueryValueEx(hKey, NULL, NULL, NULL, NULL, &dwFileTypeLen);
42
43 /* If there is no default value, don't display it */
44 if (lRet == ERROR_FILE_NOT_FOUND)
45 {
46 RegCloseKey(hKey);
47 return lRet;
48 }
49 if (lRet != ERROR_SUCCESS)
50 {
51 ErrorMessage(lRet, NULL);
52 RegCloseKey(hKey);
53 return lRet;
54 }
55
56 ++dwFileTypeLen;
57 pszFileType = cmd_alloc(dwFileTypeLen * sizeof(TCHAR));
58 if (!pszFileType)
59 {
60 WARN("Cannot allocate memory for pszFileType!\n");
61 RegCloseKey(hKey);
62 return ERROR_NOT_ENOUGH_MEMORY;
63 }
64
65 /* Obtain the actual file type */
66 lRet = RegQueryValueEx(hKey, NULL, NULL, NULL, (LPBYTE)pszFileType, &dwFileTypeLen);
67 RegCloseKey(hKey);
68
69 if (lRet != ERROR_SUCCESS)
70 {
71 ErrorMessage(lRet, NULL);
72 cmd_free(pszFileType);
73 return lRet;
74 }
75
76 /* If there is a default key, display the relevant information */
77 if (dwFileTypeLen != 0)
78 {
79 ConOutPrintf(_T("%s=%s\n"), pszExtension, pszFileType);
80 }
81
82 cmd_free(pszFileType);
83 return ERROR_SUCCESS;
84 }
85
86 static LONG
PrintAssociation(IN PCTSTR pszExtension)87 PrintAssociation(
88 IN PCTSTR pszExtension)
89 {
90 LONG lRet;
91 HKEY hKeyClasses;
92
93 lRet = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0,
94 KEY_ENUMERATE_SUB_KEYS, &hKeyClasses);
95 if (lRet != ERROR_SUCCESS)
96 {
97 ErrorMessage(lRet, NULL);
98 return lRet;
99 }
100
101 lRet = PrintAssociationEx(hKeyClasses, pszExtension);
102
103 RegCloseKey(hKeyClasses);
104 return lRet;
105 }
106
107 static LONG
PrintAllAssociations(VOID)108 PrintAllAssociations(VOID)
109 {
110 LONG lRet;
111 HKEY hKeyClasses;
112 DWORD dwKeyCtr;
113 DWORD dwNumKeys = 0;
114 DWORD dwExtLen = 0;
115 PTSTR pszExtName;
116
117 lRet = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0,
118 KEY_QUERY_VALUE | KEY_ENUMERATE_SUB_KEYS, &hKeyClasses);
119 if (lRet != ERROR_SUCCESS)
120 {
121 ErrorMessage(lRet, NULL);
122 return lRet;
123 }
124
125 lRet = RegQueryInfoKey(hKeyClasses, NULL, NULL, NULL, &dwNumKeys, &dwExtLen,
126 NULL, NULL, NULL, NULL, NULL, NULL);
127 if (lRet != ERROR_SUCCESS)
128 {
129 ErrorMessage(lRet, NULL);
130 RegCloseKey(hKeyClasses);
131 return lRet;
132 }
133
134 ++dwExtLen;
135 pszExtName = cmd_alloc(dwExtLen * sizeof(TCHAR));
136 if (!pszExtName)
137 {
138 WARN("Cannot allocate memory for pszExtName!\n");
139 RegCloseKey(hKeyClasses);
140 return ERROR_NOT_ENOUGH_MEMORY;
141 }
142
143 for (dwKeyCtr = 0; dwKeyCtr < dwNumKeys; ++dwKeyCtr)
144 {
145 DWORD dwBufSize = dwExtLen;
146 lRet = RegEnumKeyEx(hKeyClasses, dwKeyCtr, pszExtName, &dwBufSize,
147 NULL, NULL, NULL, NULL);
148
149 if (lRet == ERROR_SUCCESS || lRet == ERROR_MORE_DATA)
150 {
151 /* Name starts with '.': this is an extension */
152 if (*pszExtName == _T('.'))
153 PrintAssociationEx(hKeyClasses, pszExtName);
154 }
155 else
156 {
157 ErrorMessage(lRet, NULL);
158 cmd_free(pszExtName);
159 RegCloseKey(hKeyClasses);
160 return lRet;
161 }
162 }
163
164 RegCloseKey(hKeyClasses);
165
166 cmd_free(pszExtName);
167 return ERROR_SUCCESS;
168 }
169
170 static LONG
AddAssociation(IN PCTSTR pszExtension,IN PCTSTR pszType)171 AddAssociation(
172 IN PCTSTR pszExtension,
173 IN PCTSTR pszType)
174 {
175 LONG lRet;
176 HKEY hKeyClasses, hKey;
177
178 lRet = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0,
179 KEY_CREATE_SUB_KEY, &hKeyClasses);
180 if (lRet != ERROR_SUCCESS)
181 {
182 ErrorMessage(lRet, NULL);
183 return lRet;
184 }
185
186 lRet = RegCreateKeyEx(hKeyClasses, pszExtension, 0, NULL, REG_OPTION_NON_VOLATILE,
187 KEY_SET_VALUE, NULL, &hKey, NULL);
188 RegCloseKey(hKeyClasses);
189
190 if (lRet != ERROR_SUCCESS)
191 {
192 ErrorMessage(lRet, NULL);
193 return lRet;
194 }
195
196 lRet = RegSetValueEx(hKey, NULL, 0, REG_SZ,
197 (LPBYTE)pszType, (DWORD)(_tcslen(pszType) + 1) * sizeof(TCHAR));
198 RegCloseKey(hKey);
199
200 if (lRet != ERROR_SUCCESS)
201 {
202 ErrorMessage(lRet, NULL);
203 return lRet;
204 }
205
206 return ERROR_SUCCESS;
207 }
208
209 static LONG
RemoveAssociation(IN PCTSTR pszExtension)210 RemoveAssociation(
211 IN PCTSTR pszExtension)
212 {
213 LONG lRet;
214 HKEY hKeyClasses;
215
216 lRet = RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\Classes"), 0,
217 KEY_QUERY_VALUE, &hKeyClasses);
218 if (lRet != ERROR_SUCCESS)
219 {
220 ErrorMessage(lRet, NULL);
221 return lRet;
222 }
223
224 lRet = RegDeleteKey(hKeyClasses, pszExtension);
225 RegCloseKey(hKeyClasses);
226
227 if (lRet != ERROR_SUCCESS)
228 {
229 if (lRet != ERROR_FILE_NOT_FOUND)
230 ErrorMessage(lRet, NULL);
231 return lRet;
232 }
233
234 return ERROR_SUCCESS;
235 }
236
237
CommandAssoc(LPTSTR param)238 INT CommandAssoc(LPTSTR param)
239 {
240 INT retval = 0;
241 PTCHAR pEqualSign;
242
243 /* Print help */
244 if (!_tcsncmp(param, _T("/?"), 2))
245 {
246 ConOutResPaging(TRUE, STRING_ASSOC_HELP);
247 return 0;
248 }
249
250 /* Print all associations if no parameter has been specified */
251 if (!*param)
252 {
253 PrintAllAssociations();
254 goto Quit;
255 }
256
257 pEqualSign = _tcschr(param, _T('='));
258 if (pEqualSign != NULL)
259 {
260 PTSTR pszFileType = pEqualSign + 1;
261
262 /* NULL-terminate at the equals sign */
263 *pEqualSign = 0;
264
265 /* If the equals sign is the last character
266 * in the string, delete the association. */
267 if (*pszFileType == 0)
268 {
269 retval = RemoveAssociation(param);
270 }
271 else
272 /* Otherwise, add the association and print it out */
273 {
274 retval = AddAssociation(param, pszFileType);
275 PrintAssociation(param);
276 }
277
278 if (retval != ERROR_SUCCESS)
279 {
280 if (retval != ERROR_FILE_NOT_FOUND)
281 {
282 ConErrResPrintf(STRING_ERROR_WHILE_PROCESSING, param);
283 }
284 // retval = 1; /* Fixup the error value */
285 }
286 }
287 else
288 {
289 /* No equals sign, print the association */
290 retval = PrintAssociation(param);
291 if (retval != ERROR_SUCCESS)
292 {
293 ConErrResPrintf(STRING_ASSOC_ERROR, param);
294 retval = 1; /* Fixup the error value */
295 }
296 }
297
298 Quit:
299 if (BatType != CMD_TYPE)
300 {
301 if (retval != 0)
302 nErrorLevel = retval;
303 }
304 else
305 {
306 nErrorLevel = retval;
307 }
308
309 return retval;
310 }
311
312 #endif /* INCLUDE_CMD_ASSOC */
313