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