1 // Support back to Vista 2 #define _WIN32_WINNT _WIN32_WINNT_VISTA 3 #include <sdkddkver.h> 4 5 // Use WRL to define a classic COM class 6 #define __WRL_CLASSIC_COM__ 7 #include <wrl.h> 8 9 #include <windows.h> 10 #include <shlobj.h> 11 #include <shlwapi.h> 12 #include <olectl.h> 13 #include <strsafe.h> 14 15 #include "pyshellext_h.h" 16 17 #define DDWM_UPDATEWINDOW (WM_USER+3) 18 19 static HINSTANCE hModule; 20 static CLIPFORMAT cfDropDescription; 21 static CLIPFORMAT cfDragWindow; 22 23 static const LPCWSTR CLASS_SUBKEY = L"Software\\Classes\\CLSID\\{BEA218D2-6950-497B-9434-61683EC065FE}"; 24 static const LPCWSTR DRAG_MESSAGE = L"Open with %1"; 25 26 using namespace Microsoft::WRL; 27 28 HRESULT FilenameListCchLengthA(LPCSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) { 29 HRESULT hr = S_OK; 30 size_t count = 0; 31 size_t length = 0; 32 33 while (pszSource && pszSource[0]) { 34 size_t oneLength; 35 hr = StringCchLengthA(pszSource, cchMax - length, &oneLength); 36 if (FAILED(hr)) { 37 return hr; 38 } 39 count += 1; 40 length += oneLength + (strchr(pszSource, ' ') ? 3 : 1); 41 pszSource = &pszSource[oneLength + 1]; 42 } 43 44 *pcchCount = count; 45 *pcchLength = length; 46 return hr; 47 } 48 49 HRESULT FilenameListCchLengthW(LPCWSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) { 50 HRESULT hr = S_OK; 51 size_t count = 0; 52 size_t length = 0; 53 54 while (pszSource && pszSource[0]) { 55 size_t oneLength; 56 hr = StringCchLengthW(pszSource, cchMax - length, &oneLength); 57 if (FAILED(hr)) { 58 return hr; 59 } 60 count += 1; 61 length += oneLength + (wcschr(pszSource, ' ') ? 3 : 1); 62 pszSource = &pszSource[oneLength + 1]; 63 } 64 65 *pcchCount = count; 66 *pcchLength = length; 67 return hr; 68 } 69 70 HRESULT FilenameListCchCopyA(STRSAFE_LPSTR pszDest, size_t cchDest, LPCSTR pszSource, LPCSTR pszSeparator) { 71 HRESULT hr = S_OK; 72 size_t count = 0; 73 size_t length = 0; 74 75 while (pszSource[0]) { 76 STRSAFE_LPSTR newDest; 77 78 hr = StringCchCopyExA(pszDest, cchDest, pszSource, &newDest, &cchDest, 0); 79 if (FAILED(hr)) { 80 return hr; 81 } 82 pszSource += (newDest - pszDest) + 1; 83 pszDest = PathQuoteSpacesA(pszDest) ? newDest + 2 : newDest; 84 85 if (pszSource[0]) { 86 hr = StringCchCopyExA(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0); 87 if (FAILED(hr)) { 88 return hr; 89 } 90 pszDest = newDest; 91 } 92 } 93 94 return hr; 95 } 96 97 HRESULT FilenameListCchCopyW(STRSAFE_LPWSTR pszDest, size_t cchDest, LPCWSTR pszSource, LPCWSTR pszSeparator) { 98 HRESULT hr = S_OK; 99 size_t count = 0; 100 size_t length = 0; 101 102 while (pszSource[0]) { 103 STRSAFE_LPWSTR newDest; 104 105 hr = StringCchCopyExW(pszDest, cchDest, pszSource, &newDest, &cchDest, 0); 106 if (FAILED(hr)) { 107 return hr; 108 } 109 pszSource += (newDest - pszDest) + 1; 110 pszDest = PathQuoteSpacesW(pszDest) ? newDest + 2 : newDest; 111 112 if (pszSource[0]) { 113 hr = StringCchCopyExW(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0); 114 if (FAILED(hr)) { 115 return hr; 116 } 117 pszDest = newDest; 118 } 119 } 120 121 return hr; 122 } 123 124 125 class PyShellExt : public RuntimeClass< 126 RuntimeClassFlags<ClassicCom>, 127 IDropTarget, 128 IPersistFile 129 > 130 { 131 LPOLESTR target, target_dir; 132 DWORD target_mode; 133 134 IDataObject *data_obj; 135 136 public: 137 PyShellExt() : target(NULL), target_dir(NULL), target_mode(0), data_obj(NULL) { 138 OutputDebugString(L"PyShellExt::PyShellExt"); 139 } 140 141 ~PyShellExt() { 142 if (target) { 143 CoTaskMemFree(target); 144 } 145 if (target_dir) { 146 CoTaskMemFree(target_dir); 147 } 148 if (data_obj) { 149 data_obj->Release(); 150 } 151 } 152 153 private: 154 HRESULT UpdateDropDescription(IDataObject *pDataObj) { 155 STGMEDIUM medium; 156 FORMATETC fmt = { 157 cfDropDescription, 158 NULL, 159 DVASPECT_CONTENT, 160 -1, 161 TYMED_HGLOBAL 162 }; 163 164 auto hr = pDataObj->GetData(&fmt, &medium); 165 if (FAILED(hr)) { 166 OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to get DROPDESCRIPTION format"); 167 return hr; 168 } 169 if (!medium.hGlobal) { 170 OutputDebugString(L"PyShellExt::UpdateDropDescription - DROPDESCRIPTION format had NULL hGlobal"); 171 ReleaseStgMedium(&medium); 172 return E_FAIL; 173 } 174 auto dd = (DROPDESCRIPTION*)GlobalLock(medium.hGlobal); 175 StringCchCopy(dd->szMessage, sizeof(dd->szMessage) / sizeof(dd->szMessage[0]), DRAG_MESSAGE); 176 StringCchCopy(dd->szInsert, sizeof(dd->szInsert) / sizeof(dd->szInsert[0]), PathFindFileNameW(target)); 177 dd->type = DROPIMAGE_MOVE; 178 179 GlobalUnlock(medium.hGlobal); 180 ReleaseStgMedium(&medium); 181 182 return S_OK; 183 } 184 185 HRESULT GetDragWindow(IDataObject *pDataObj, HWND *phWnd) { 186 HRESULT hr; 187 HWND *pMem; 188 STGMEDIUM medium; 189 FORMATETC fmt = { 190 cfDragWindow, 191 NULL, 192 DVASPECT_CONTENT, 193 -1, 194 TYMED_HGLOBAL 195 }; 196 197 hr = pDataObj->GetData(&fmt, &medium); 198 if (FAILED(hr)) { 199 OutputDebugString(L"PyShellExt::GetDragWindow - failed to get DragWindow format"); 200 return hr; 201 } 202 if (!medium.hGlobal) { 203 OutputDebugString(L"PyShellExt::GetDragWindow - DragWindow format had NULL hGlobal"); 204 ReleaseStgMedium(&medium); 205 return E_FAIL; 206 } 207 208 pMem = (HWND*)GlobalLock(medium.hGlobal); 209 if (!pMem) { 210 OutputDebugString(L"PyShellExt::GetDragWindow - failed to lock DragWindow hGlobal"); 211 ReleaseStgMedium(&medium); 212 return E_FAIL; 213 } 214 215 *phWnd = *pMem; 216 217 GlobalUnlock(medium.hGlobal); 218 ReleaseStgMedium(&medium); 219 220 return S_OK; 221 } 222 223 HRESULT GetArguments(IDataObject *pDataObj, LPCWSTR *pArguments) { 224 HRESULT hr; 225 DROPFILES *pdropfiles; 226 227 STGMEDIUM medium; 228 FORMATETC fmt = { 229 CF_HDROP, 230 NULL, 231 DVASPECT_CONTENT, 232 -1, 233 TYMED_HGLOBAL 234 }; 235 236 hr = pDataObj->GetData(&fmt, &medium); 237 if (FAILED(hr)) { 238 OutputDebugString(L"PyShellExt::GetArguments - failed to get CF_HDROP format"); 239 return hr; 240 } 241 if (!medium.hGlobal) { 242 OutputDebugString(L"PyShellExt::GetArguments - CF_HDROP format had NULL hGlobal"); 243 ReleaseStgMedium(&medium); 244 return E_FAIL; 245 } 246 247 pdropfiles = (DROPFILES*)GlobalLock(medium.hGlobal); 248 if (!pdropfiles) { 249 OutputDebugString(L"PyShellExt::GetArguments - failed to lock CF_HDROP hGlobal"); 250 ReleaseStgMedium(&medium); 251 return E_FAIL; 252 } 253 254 if (pdropfiles->fWide) { 255 LPCWSTR files = (LPCWSTR)((char*)pdropfiles + pdropfiles->pFiles); 256 size_t len, count; 257 hr = FilenameListCchLengthW(files, 32767, &len, &count); 258 if (SUCCEEDED(hr)) { 259 LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); 260 if (args) { 261 hr = FilenameListCchCopyW(args, 32767, files, L" "); 262 if (SUCCEEDED(hr)) { 263 *pArguments = args; 264 } else { 265 CoTaskMemFree(args); 266 } 267 } else { 268 hr = E_OUTOFMEMORY; 269 } 270 } 271 } else { 272 LPCSTR files = (LPCSTR)((char*)pdropfiles + pdropfiles->pFiles); 273 size_t len, count; 274 hr = FilenameListCchLengthA(files, 32767, &len, &count); 275 if (SUCCEEDED(hr)) { 276 LPSTR temp = (LPSTR)CoTaskMemAlloc(sizeof(CHAR) * (len + 1)); 277 if (temp) { 278 hr = FilenameListCchCopyA(temp, 32767, files, " "); 279 if (SUCCEEDED(hr)) { 280 int wlen = MultiByteToWideChar(CP_ACP, 0, temp, (int)len, NULL, 0); 281 if (wlen) { 282 LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (wlen + 1)); 283 if (MultiByteToWideChar(CP_ACP, 0, temp, (int)len, args, wlen + 1)) { 284 *pArguments = args; 285 } else { 286 OutputDebugString(L"PyShellExt::GetArguments - failed to convert multi-byte to wide-char path"); 287 CoTaskMemFree(args); 288 hr = E_FAIL; 289 } 290 } else { 291 OutputDebugString(L"PyShellExt::GetArguments - failed to get length of wide-char path"); 292 hr = E_FAIL; 293 } 294 } 295 CoTaskMemFree(temp); 296 } else { 297 hr = E_OUTOFMEMORY; 298 } 299 } 300 } 301 302 GlobalUnlock(medium.hGlobal); 303 ReleaseStgMedium(&medium); 304 305 return hr; 306 } 307 308 HRESULT NotifyDragWindow(HWND hwnd) { 309 LRESULT res; 310 311 if (!hwnd) { 312 return S_FALSE; 313 } 314 315 res = SendMessage(hwnd, DDWM_UPDATEWINDOW, 0, NULL); 316 317 if (res) { 318 OutputDebugString(L"PyShellExt::NotifyDragWindow - failed to post DDWM_UPDATEWINDOW"); 319 return E_FAIL; 320 } 321 322 return S_OK; 323 } 324 325 public: 326 // IDropTarget implementation 327 328 STDMETHODIMP DragEnter(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { 329 HWND hwnd; 330 331 OutputDebugString(L"PyShellExt::DragEnter"); 332 333 pDataObj->AddRef(); 334 data_obj = pDataObj; 335 336 *pdwEffect = DROPEFFECT_MOVE; 337 338 if (FAILED(UpdateDropDescription(data_obj))) { 339 OutputDebugString(L"PyShellExt::DragEnter - failed to update drop description"); 340 } 341 if (FAILED(GetDragWindow(data_obj, &hwnd))) { 342 OutputDebugString(L"PyShellExt::DragEnter - failed to get drag window"); 343 } 344 if (FAILED(NotifyDragWindow(hwnd))) { 345 OutputDebugString(L"PyShellExt::DragEnter - failed to notify drag window"); 346 } 347 348 return S_OK; 349 } 350 351 STDMETHODIMP DragLeave() { 352 return S_OK; 353 } 354 355 STDMETHODIMP DragOver(DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { 356 return S_OK; 357 } 358 359 STDMETHODIMP Drop(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { 360 LPCWSTR args; 361 362 OutputDebugString(L"PyShellExt::Drop"); 363 *pdwEffect = DROPEFFECT_NONE; 364 365 if (pDataObj != data_obj) { 366 OutputDebugString(L"PyShellExt::Drop - unexpected data object"); 367 return E_FAIL; 368 } 369 370 data_obj->Release(); 371 data_obj = NULL; 372 373 if (SUCCEEDED(GetArguments(pDataObj, &args))) { 374 OutputDebugString(args); 375 ShellExecute(NULL, NULL, target, args, target_dir, SW_NORMAL); 376 377 CoTaskMemFree((LPVOID)args); 378 } else { 379 OutputDebugString(L"PyShellExt::Drop - failed to get launch arguments"); 380 } 381 382 return S_OK; 383 } 384 385 // IPersistFile implementation 386 387 STDMETHODIMP GetCurFile(LPOLESTR *ppszFileName) { 388 HRESULT hr; 389 size_t len; 390 391 if (!ppszFileName) { 392 return E_POINTER; 393 } 394 395 hr = StringCchLength(target, STRSAFE_MAX_CCH - 1, &len); 396 if (FAILED(hr)) { 397 return E_FAIL; 398 } 399 400 *ppszFileName = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); 401 if (!*ppszFileName) { 402 return E_OUTOFMEMORY; 403 } 404 405 hr = StringCchCopy(*ppszFileName, len + 1, target); 406 if (FAILED(hr)) { 407 CoTaskMemFree(*ppszFileName); 408 *ppszFileName = NULL; 409 return E_FAIL; 410 } 411 412 return S_OK; 413 } 414 415 STDMETHODIMP IsDirty() { 416 return S_FALSE; 417 } 418 419 STDMETHODIMP Load(LPCOLESTR pszFileName, DWORD dwMode) { 420 HRESULT hr; 421 size_t len; 422 423 OutputDebugString(L"PyShellExt::Load"); 424 OutputDebugString(pszFileName); 425 426 hr = StringCchLength(pszFileName, STRSAFE_MAX_CCH - 1, &len); 427 if (FAILED(hr)) { 428 OutputDebugString(L"PyShellExt::Load - failed to get string length"); 429 return hr; 430 } 431 432 if (target) { 433 CoTaskMemFree(target); 434 } 435 if (target_dir) { 436 CoTaskMemFree(target_dir); 437 } 438 439 target = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); 440 if (!target) { 441 OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY"); 442 return E_OUTOFMEMORY; 443 } 444 target_dir = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); 445 if (!target_dir) { 446 OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY"); 447 return E_OUTOFMEMORY; 448 } 449 450 hr = StringCchCopy(target, len + 1, pszFileName); 451 if (FAILED(hr)) { 452 OutputDebugString(L"PyShellExt::Load - failed to copy string"); 453 return hr; 454 } 455 456 hr = StringCchCopy(target_dir, len + 1, pszFileName); 457 if (FAILED(hr)) { 458 OutputDebugString(L"PyShellExt::Load - failed to copy string"); 459 return hr; 460 } 461 if (!PathRemoveFileSpecW(target_dir)) { 462 OutputDebugStringW(L"PyShellExt::Load - failed to remove filespec from target"); 463 return E_FAIL; 464 } 465 466 OutputDebugString(target); 467 target_mode = dwMode; 468 OutputDebugString(L"PyShellExt::Load - S_OK"); 469 return S_OK; 470 } 471 472 STDMETHODIMP Save(LPCOLESTR pszFileName, BOOL fRemember) { 473 return E_NOTIMPL; 474 } 475 476 STDMETHODIMP SaveCompleted(LPCOLESTR pszFileName) { 477 return E_NOTIMPL; 478 } 479 480 STDMETHODIMP GetClassID(CLSID *pClassID) { 481 *pClassID = CLSID_PyShellExt; 482 return S_OK; 483 } 484 }; 485 486 CoCreatableClass(PyShellExt); 487 488 STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) { 489 return Module<InProc>::GetModule().GetClassObject(rclsid, riid, ppv); 490 } 491 492 STDAPI DllCanUnloadNow() { 493 return Module<InProc>::GetModule().Terminate() ? S_OK : S_FALSE; 494 } 495 496 STDAPI DllRegisterServer() { 497 LONG res; 498 SECURITY_ATTRIBUTES secattr = { sizeof(SECURITY_ATTRIBUTES), NULL, FALSE }; 499 LPSECURITY_ATTRIBUTES psecattr = NULL; 500 HKEY key, ipsKey; 501 WCHAR modname[MAX_PATH]; 502 DWORD modname_len; 503 504 OutputDebugString(L"PyShellExt::DllRegisterServer"); 505 if (!hModule) { 506 OutputDebugString(L"PyShellExt::DllRegisterServer - module handle was not set"); 507 return SELFREG_E_CLASS; 508 } 509 modname_len = GetModuleFileName(hModule, modname, MAX_PATH); 510 if (modname_len == 0 || 511 (modname_len == MAX_PATH && GetLastError() == ERROR_INSUFFICIENT_BUFFER)) { 512 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to get module file name"); 513 return SELFREG_E_CLASS; 514 } 515 516 DWORD disp; 517 res = RegCreateKeyEx(HKEY_LOCAL_MACHINE, CLASS_SUBKEY, 0, NULL, 0, 518 KEY_ALL_ACCESS, psecattr, &key, &disp); 519 if (res == ERROR_ACCESS_DENIED) { 520 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to write per-machine registration. Attempting per-user instead."); 521 res = RegCreateKeyEx(HKEY_CURRENT_USER, CLASS_SUBKEY, 0, NULL, 0, 522 KEY_ALL_ACCESS, psecattr, &key, &disp); 523 } 524 if (res != ERROR_SUCCESS) { 525 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create class key"); 526 return SELFREG_E_CLASS; 527 } 528 529 res = RegCreateKeyEx(key, L"InProcServer32", 0, NULL, 0, 530 KEY_ALL_ACCESS, psecattr, &ipsKey, NULL); 531 if (res != ERROR_SUCCESS) { 532 RegCloseKey(key); 533 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create InProcServer32 key"); 534 return SELFREG_E_CLASS; 535 } 536 537 res = RegSetValueEx(ipsKey, NULL, 0, 538 REG_SZ, (LPBYTE)modname, modname_len * sizeof(modname[0])); 539 540 if (res != ERROR_SUCCESS) { 541 RegCloseKey(ipsKey); 542 RegCloseKey(key); 543 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set server path"); 544 return SELFREG_E_CLASS; 545 } 546 547 res = RegSetValueEx(ipsKey, L"ThreadingModel", 0, 548 REG_SZ, (LPBYTE)(L"Apartment"), sizeof(L"Apartment")); 549 550 RegCloseKey(ipsKey); 551 RegCloseKey(key); 552 if (res != ERROR_SUCCESS) { 553 OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set threading model"); 554 return SELFREG_E_CLASS; 555 } 556 557 SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL); 558 559 OutputDebugString(L"PyShellExt::DllRegisterServer - S_OK"); 560 return S_OK; 561 } 562 563 STDAPI DllUnregisterServer() { 564 LONG res_lm, res_cu; 565 566 res_lm = RegDeleteTree(HKEY_LOCAL_MACHINE, CLASS_SUBKEY); 567 if (res_lm != ERROR_SUCCESS && res_lm != ERROR_FILE_NOT_FOUND) { 568 OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-machine registration"); 569 return SELFREG_E_CLASS; 570 } 571 572 res_cu = RegDeleteTree(HKEY_CURRENT_USER, CLASS_SUBKEY); 573 if (res_cu != ERROR_SUCCESS && res_cu != ERROR_FILE_NOT_FOUND) { 574 OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-user registration"); 575 return SELFREG_E_CLASS; 576 } 577 578 if (res_lm == ERROR_FILE_NOT_FOUND && res_cu == ERROR_FILE_NOT_FOUND) { 579 OutputDebugString(L"PyShellExt::DllUnregisterServer - extension was not registered"); 580 return SELFREG_E_CLASS; 581 } 582 583 SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL); 584 585 OutputDebugString(L"PyShellExt::DllUnregisterServer - S_OK"); 586 return S_OK; 587 } 588 589 STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) { 590 if (reason == DLL_PROCESS_ATTACH) { 591 hModule = hinst; 592 593 cfDropDescription = RegisterClipboardFormat(CFSTR_DROPDESCRIPTION); 594 if (!cfDropDescription) { 595 OutputDebugString(L"PyShellExt::DllMain - failed to get CFSTR_DROPDESCRIPTION format"); 596 } 597 cfDragWindow = RegisterClipboardFormat(L"DragWindow"); 598 if (!cfDragWindow) { 599 OutputDebugString(L"PyShellExt::DllMain - failed to get DragWindow format"); 600 } 601 602 DisableThreadLibraryCalls(hinst); 603 } 604 return TRUE; 605 }