base/shell/cmd/assoc.c

   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
  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
  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
 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
 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
 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
 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                 // FIXME: Localize
 283                 ConErrPrintf(_T("Error occurred while processing: %s.\n"), param);
 284             }
 285             // retval = 1; /* Fixup the error value */
 286         }
 287     }
 288     else
 289     {
 290         /* No equals sign, print the association */
 291         retval = PrintAssociation(param);
 292         if (retval != ERROR_SUCCESS)
 293         {
 294             ConErrResPrintf(STRING_ASSOC_ERROR, param);
 295             retval = 1; /* Fixup the error value */
 296         }
 297     }
 298
 299 Quit:
 300     if (BatType != CMD_TYPE)
 301     {
 302         if (retval != 0)
 303             nErrorLevel = retval;
 304     }
 305     else
 306     {
 307         nErrorLevel = retval;
 308     }
 309
 310     return retval;
 311 }
 312
 313 #endif /* INCLUDE_CMD_ASSOC */