diff options
Diffstat (limited to 'PC/pyshellext.cpp')
| -rw-r--r-- | PC/pyshellext.cpp | 605 |
1 files changed, 605 insertions, 0 deletions
diff --git a/PC/pyshellext.cpp b/PC/pyshellext.cpp new file mode 100644 index 0000000..04fe61e --- /dev/null +++ b/PC/pyshellext.cpp @@ -0,0 +1,605 @@ +// Support back to Vista +#define _WIN32_WINNT _WIN32_WINNT_VISTA +#include <sdkddkver.h> + +// Use WRL to define a classic COM class +#define __WRL_CLASSIC_COM__ +#include <wrl.h> + +#include <windows.h> +#include <shlobj.h> +#include <shlwapi.h> +#include <olectl.h> +#include <strsafe.h> + +#include "pyshellext_h.h" + +#define DDWM_UPDATEWINDOW (WM_USER+3) + +static HINSTANCE hModule; +static CLIPFORMAT cfDropDescription; +static CLIPFORMAT cfDragWindow; + +static const LPCWSTR CLASS_SUBKEY = L"Software\\Classes\\CLSID\\{BEA218D2-6950-497B-9434-61683EC065FE}"; +static const LPCWSTR DRAG_MESSAGE = L"Open with %1"; + +using namespace Microsoft::WRL; + +HRESULT FilenameListCchLengthA(LPCSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) { + HRESULT hr = S_OK; + size_t count = 0; + size_t length = 0; + + while (pszSource && pszSource[0]) { + size_t oneLength; + hr = StringCchLengthA(pszSource, cchMax - length, &oneLength); + if (FAILED(hr)) { + return hr; + } + count += 1; + length += oneLength + (strchr(pszSource, ' ') ? 3 : 1); + pszSource = &pszSource[oneLength + 1]; + } + + *pcchCount = count; + *pcchLength = length; + return hr; +} + +HRESULT FilenameListCchLengthW(LPCWSTR pszSource, size_t cchMax, size_t *pcchLength, size_t *pcchCount) { + HRESULT hr = S_OK; + size_t count = 0; + size_t length = 0; + + while (pszSource && pszSource[0]) { + size_t oneLength; + hr = StringCchLengthW(pszSource, cchMax - length, &oneLength); + if (FAILED(hr)) { + return hr; + } + count += 1; + length += oneLength + (wcschr(pszSource, ' ') ? 3 : 1); + pszSource = &pszSource[oneLength + 1]; + } + + *pcchCount = count; + *pcchLength = length; + return hr; +} + +HRESULT FilenameListCchCopyA(STRSAFE_LPSTR pszDest, size_t cchDest, LPCSTR pszSource, LPCSTR pszSeparator) { + HRESULT hr = S_OK; + size_t count = 0; + size_t length = 0; + + while (pszSource[0]) { + STRSAFE_LPSTR newDest; + + hr = StringCchCopyExA(pszDest, cchDest, pszSource, &newDest, &cchDest, 0); + if (FAILED(hr)) { + return hr; + } + pszSource += (newDest - pszDest) + 1; + pszDest = PathQuoteSpacesA(pszDest) ? newDest + 2 : newDest; + + if (pszSource[0]) { + hr = StringCchCopyExA(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0); + if (FAILED(hr)) { + return hr; + } + pszDest = newDest; + } + } + + return hr; +} + +HRESULT FilenameListCchCopyW(STRSAFE_LPWSTR pszDest, size_t cchDest, LPCWSTR pszSource, LPCWSTR pszSeparator) { + HRESULT hr = S_OK; + size_t count = 0; + size_t length = 0; + + while (pszSource[0]) { + STRSAFE_LPWSTR newDest; + + hr = StringCchCopyExW(pszDest, cchDest, pszSource, &newDest, &cchDest, 0); + if (FAILED(hr)) { + return hr; + } + pszSource += (newDest - pszDest) + 1; + pszDest = PathQuoteSpacesW(pszDest) ? newDest + 2 : newDest; + + if (pszSource[0]) { + hr = StringCchCopyExW(pszDest, cchDest, pszSeparator, &newDest, &cchDest, 0); + if (FAILED(hr)) { + return hr; + } + pszDest = newDest; + } + } + + return hr; +} + + +class PyShellExt : public RuntimeClass< + RuntimeClassFlags<ClassicCom>, + IDropTarget, + IPersistFile +> +{ + LPOLESTR target, target_dir; + DWORD target_mode; + + IDataObject *data_obj; + +public: + PyShellExt() : target(NULL), target_dir(NULL), target_mode(0), data_obj(NULL) { + OutputDebugString(L"PyShellExt::PyShellExt"); + } + + ~PyShellExt() { + if (target) { + CoTaskMemFree(target); + } + if (target_dir) { + CoTaskMemFree(target_dir); + } + if (data_obj) { + data_obj->Release(); + } + } + +private: + HRESULT UpdateDropDescription(IDataObject *pDataObj) { + STGMEDIUM medium; + FORMATETC fmt = { + cfDropDescription, + NULL, + DVASPECT_CONTENT, + -1, + TYMED_HGLOBAL + }; + + auto hr = pDataObj->GetData(&fmt, &medium); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to get DROPDESCRIPTION format"); + return hr; + } + if (!medium.hGlobal) { + OutputDebugString(L"PyShellExt::UpdateDropDescription - DROPDESCRIPTION format had NULL hGlobal"); + ReleaseStgMedium(&medium); + return E_FAIL; + } + auto dd = (DROPDESCRIPTION*)GlobalLock(medium.hGlobal); + StringCchCopy(dd->szMessage, sizeof(dd->szMessage) / sizeof(dd->szMessage[0]), DRAG_MESSAGE); + StringCchCopy(dd->szInsert, sizeof(dd->szInsert) / sizeof(dd->szInsert[0]), PathFindFileNameW(target)); + dd->type = DROPIMAGE_MOVE; + + GlobalUnlock(medium.hGlobal); + ReleaseStgMedium(&medium); + + return S_OK; + } + + HRESULT GetDragWindow(IDataObject *pDataObj, HWND *phWnd) { + HRESULT hr; + HWND *pMem; + STGMEDIUM medium; + FORMATETC fmt = { + cfDragWindow, + NULL, + DVASPECT_CONTENT, + -1, + TYMED_HGLOBAL + }; + + hr = pDataObj->GetData(&fmt, &medium); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::GetDragWindow - failed to get DragWindow format"); + return hr; + } + if (!medium.hGlobal) { + OutputDebugString(L"PyShellExt::GetDragWindow - DragWindow format had NULL hGlobal"); + ReleaseStgMedium(&medium); + return E_FAIL; + } + + pMem = (HWND*)GlobalLock(medium.hGlobal); + if (!pMem) { + OutputDebugString(L"PyShellExt::GetDragWindow - failed to lock DragWindow hGlobal"); + ReleaseStgMedium(&medium); + return E_FAIL; + } + + *phWnd = *pMem; + + GlobalUnlock(medium.hGlobal); + ReleaseStgMedium(&medium); + + return S_OK; + } + + HRESULT GetArguments(IDataObject *pDataObj, LPCWSTR *pArguments) { + HRESULT hr; + DROPFILES *pdropfiles; + + STGMEDIUM medium; + FORMATETC fmt = { + CF_HDROP, + NULL, + DVASPECT_CONTENT, + -1, + TYMED_HGLOBAL + }; + + hr = pDataObj->GetData(&fmt, &medium); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::GetArguments - failed to get CF_HDROP format"); + return hr; + } + if (!medium.hGlobal) { + OutputDebugString(L"PyShellExt::GetArguments - CF_HDROP format had NULL hGlobal"); + ReleaseStgMedium(&medium); + return E_FAIL; + } + + pdropfiles = (DROPFILES*)GlobalLock(medium.hGlobal); + if (!pdropfiles) { + OutputDebugString(L"PyShellExt::GetArguments - failed to lock CF_HDROP hGlobal"); + ReleaseStgMedium(&medium); + return E_FAIL; + } + + if (pdropfiles->fWide) { + LPCWSTR files = (LPCWSTR)((char*)pdropfiles + pdropfiles->pFiles); + size_t len, count; + hr = FilenameListCchLengthW(files, 32767, &len, &count); + if (SUCCEEDED(hr)) { + LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); + if (args) { + hr = FilenameListCchCopyW(args, 32767, files, L" "); + if (SUCCEEDED(hr)) { + *pArguments = args; + } else { + CoTaskMemFree(args); + } + } else { + hr = E_OUTOFMEMORY; + } + } + } else { + LPCSTR files = (LPCSTR)((char*)pdropfiles + pdropfiles->pFiles); + size_t len, count; + hr = FilenameListCchLengthA(files, 32767, &len, &count); + if (SUCCEEDED(hr)) { + LPSTR temp = (LPSTR)CoTaskMemAlloc(sizeof(CHAR) * (len + 1)); + if (temp) { + hr = FilenameListCchCopyA(temp, 32767, files, " "); + if (SUCCEEDED(hr)) { + int wlen = MultiByteToWideChar(CP_ACP, 0, temp, (int)len, NULL, 0); + if (wlen) { + LPWSTR args = (LPWSTR)CoTaskMemAlloc(sizeof(WCHAR) * (wlen + 1)); + if (MultiByteToWideChar(CP_ACP, 0, temp, (int)len, args, wlen + 1)) { + *pArguments = args; + } else { + OutputDebugString(L"PyShellExt::GetArguments - failed to convert multi-byte to wide-char path"); + CoTaskMemFree(args); + hr = E_FAIL; + } + } else { + OutputDebugString(L"PyShellExt::GetArguments - failed to get length of wide-char path"); + hr = E_FAIL; + } + } + CoTaskMemFree(temp); + } else { + hr = E_OUTOFMEMORY; + } + } + } + + GlobalUnlock(medium.hGlobal); + ReleaseStgMedium(&medium); + + return hr; + } + + HRESULT NotifyDragWindow(HWND hwnd) { + LRESULT res; + + if (!hwnd) { + return S_FALSE; + } + + res = SendMessage(hwnd, DDWM_UPDATEWINDOW, 0, NULL); + + if (res) { + OutputDebugString(L"PyShellExt::NotifyDragWindow - failed to post DDWM_UPDATEWINDOW"); + return E_FAIL; + } + + return S_OK; + } + +public: + // IDropTarget implementation + + STDMETHODIMP DragEnter(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { + HWND hwnd; + + OutputDebugString(L"PyShellExt::DragEnter"); + + pDataObj->AddRef(); + data_obj = pDataObj; + + *pdwEffect = DROPEFFECT_MOVE; + + if (FAILED(UpdateDropDescription(data_obj))) { + OutputDebugString(L"PyShellExt::DragEnter - failed to update drop description"); + } + if (FAILED(GetDragWindow(data_obj, &hwnd))) { + OutputDebugString(L"PyShellExt::DragEnter - failed to get drag window"); + } + if (FAILED(NotifyDragWindow(hwnd))) { + OutputDebugString(L"PyShellExt::DragEnter - failed to notify drag window"); + } + + return S_OK; + } + + STDMETHODIMP DragLeave() { + return S_OK; + } + + STDMETHODIMP DragOver(DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { + return S_OK; + } + + STDMETHODIMP Drop(IDataObject *pDataObj, DWORD grfKeyState, POINTL pt, DWORD *pdwEffect) { + LPCWSTR args; + + OutputDebugString(L"PyShellExt::Drop"); + *pdwEffect = DROPEFFECT_NONE; + + if (pDataObj != data_obj) { + OutputDebugString(L"PyShellExt::Drop - unexpected data object"); + return E_FAIL; + } + + data_obj->Release(); + data_obj = NULL; + + if (SUCCEEDED(GetArguments(pDataObj, &args))) { + OutputDebugString(args); + ShellExecute(NULL, NULL, target, args, target_dir, SW_NORMAL); + + CoTaskMemFree((LPVOID)args); + } else { + OutputDebugString(L"PyShellExt::Drop - failed to get launch arguments"); + } + + return S_OK; + } + + // IPersistFile implementation + + STDMETHODIMP GetCurFile(LPOLESTR *ppszFileName) { + HRESULT hr; + size_t len; + + if (!ppszFileName) { + return E_POINTER; + } + + hr = StringCchLength(target, STRSAFE_MAX_CCH - 1, &len); + if (FAILED(hr)) { + return E_FAIL; + } + + *ppszFileName = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); + if (!*ppszFileName) { + return E_OUTOFMEMORY; + } + + hr = StringCchCopy(*ppszFileName, len + 1, target); + if (FAILED(hr)) { + CoTaskMemFree(*ppszFileName); + *ppszFileName = NULL; + return E_FAIL; + } + + return S_OK; + } + + STDMETHODIMP IsDirty() { + return S_FALSE; + } + + STDMETHODIMP Load(LPCOLESTR pszFileName, DWORD dwMode) { + HRESULT hr; + size_t len; + + OutputDebugString(L"PyShellExt::Load"); + OutputDebugString(pszFileName); + + hr = StringCchLength(pszFileName, STRSAFE_MAX_CCH - 1, &len); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::Load - failed to get string length"); + return hr; + } + + if (target) { + CoTaskMemFree(target); + } + if (target_dir) { + CoTaskMemFree(target_dir); + } + + target = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); + if (!target) { + OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY"); + return E_OUTOFMEMORY; + } + target_dir = (LPOLESTR)CoTaskMemAlloc(sizeof(WCHAR) * (len + 1)); + if (!target_dir) { + OutputDebugString(L"PyShellExt::Load - E_OUTOFMEMORY"); + return E_OUTOFMEMORY; + } + + hr = StringCchCopy(target, len + 1, pszFileName); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::Load - failed to copy string"); + return hr; + } + + hr = StringCchCopy(target_dir, len + 1, pszFileName); + if (FAILED(hr)) { + OutputDebugString(L"PyShellExt::Load - failed to copy string"); + return hr; + } + if (!PathRemoveFileSpecW(target_dir)) { + OutputDebugStringW(L"PyShellExt::Load - failed to remove filespec from target"); + return E_FAIL; + } + + OutputDebugString(target); + target_mode = dwMode; + OutputDebugString(L"PyShellExt::Load - S_OK"); + return S_OK; + } + + STDMETHODIMP Save(LPCOLESTR pszFileName, BOOL fRemember) { + return E_NOTIMPL; + } + + STDMETHODIMP SaveCompleted(LPCOLESTR pszFileName) { + return E_NOTIMPL; + } + + STDMETHODIMP GetClassID(CLSID *pClassID) { + *pClassID = CLSID_PyShellExt; + return S_OK; + } +}; + +CoCreatableClass(PyShellExt); + +STDAPI DllGetClassObject(REFCLSID rclsid, REFIID riid, _COM_Outptr_ void** ppv) { + return Module<InProc>::GetModule().GetClassObject(rclsid, riid, ppv); +} + +STDAPI DllCanUnloadNow() { + return Module<InProc>::GetModule().Terminate() ? S_OK : S_FALSE; +} + +STDAPI DllRegisterServer() { + LONG res; + SECURITY_ATTRIBUTES secattr = { sizeof(SECURITY_ATTRIBUTES), NULL, FALSE }; + LPSECURITY_ATTRIBUTES psecattr = NULL; + HKEY key, ipsKey; + WCHAR modname[MAX_PATH]; + DWORD modname_len; + + OutputDebugString(L"PyShellExt::DllRegisterServer"); + if (!hModule) { + OutputDebugString(L"PyShellExt::DllRegisterServer - module handle was not set"); + return SELFREG_E_CLASS; + } + modname_len = GetModuleFileName(hModule, modname, MAX_PATH); + if (modname_len == 0 || + (modname_len == MAX_PATH && GetLastError() == ERROR_INSUFFICIENT_BUFFER)) { + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to get module file name"); + return SELFREG_E_CLASS; + } + + DWORD disp; + res = RegCreateKeyEx(HKEY_LOCAL_MACHINE, CLASS_SUBKEY, 0, NULL, 0, + KEY_ALL_ACCESS, psecattr, &key, &disp); + if (res == ERROR_ACCESS_DENIED) { + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to write per-machine registration. Attempting per-user instead."); + res = RegCreateKeyEx(HKEY_CURRENT_USER, CLASS_SUBKEY, 0, NULL, 0, + KEY_ALL_ACCESS, psecattr, &key, &disp); + } + if (res != ERROR_SUCCESS) { + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create class key"); + return SELFREG_E_CLASS; + } + + res = RegCreateKeyEx(key, L"InProcServer32", 0, NULL, 0, + KEY_ALL_ACCESS, psecattr, &ipsKey, NULL); + if (res != ERROR_SUCCESS) { + RegCloseKey(key); + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to create InProcServer32 key"); + return SELFREG_E_CLASS; + } + + res = RegSetValueEx(ipsKey, NULL, 0, + REG_SZ, (LPBYTE)modname, modname_len * sizeof(modname[0])); + + if (res != ERROR_SUCCESS) { + RegCloseKey(ipsKey); + RegCloseKey(key); + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set server path"); + return SELFREG_E_CLASS; + } + + res = RegSetValueEx(ipsKey, L"ThreadingModel", 0, + REG_SZ, (LPBYTE)(L"Apartment"), sizeof(L"Apartment")); + + RegCloseKey(ipsKey); + RegCloseKey(key); + if (res != ERROR_SUCCESS) { + OutputDebugString(L"PyShellExt::DllRegisterServer - failed to set threading model"); + return SELFREG_E_CLASS; + } + + SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL); + + OutputDebugString(L"PyShellExt::DllRegisterServer - S_OK"); + return S_OK; +} + +STDAPI DllUnregisterServer() { + LONG res_lm, res_cu; + + res_lm = RegDeleteTree(HKEY_LOCAL_MACHINE, CLASS_SUBKEY); + if (res_lm != ERROR_SUCCESS && res_lm != ERROR_FILE_NOT_FOUND) { + OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-machine registration"); + return SELFREG_E_CLASS; + } + + res_cu = RegDeleteTree(HKEY_CURRENT_USER, CLASS_SUBKEY); + if (res_cu != ERROR_SUCCESS && res_cu != ERROR_FILE_NOT_FOUND) { + OutputDebugString(L"PyShellExt::DllUnregisterServer - failed to delete per-user registration"); + return SELFREG_E_CLASS; + } + + if (res_lm == ERROR_FILE_NOT_FOUND && res_cu == ERROR_FILE_NOT_FOUND) { + OutputDebugString(L"PyShellExt::DllUnregisterServer - extension was not registered"); + return SELFREG_E_CLASS; + } + + SHChangeNotify(SHCNE_ASSOCCHANGED, SHCNF_IDLIST, NULL, NULL); + + OutputDebugString(L"PyShellExt::DllUnregisterServer - S_OK"); + return S_OK; +} + +STDAPI_(BOOL) DllMain(_In_opt_ HINSTANCE hinst, DWORD reason, _In_opt_ void*) { + if (reason == DLL_PROCESS_ATTACH) { + hModule = hinst; + + cfDropDescription = RegisterClipboardFormat(CFSTR_DROPDESCRIPTION); + if (!cfDropDescription) { + OutputDebugString(L"PyShellExt::DllMain - failed to get CFSTR_DROPDESCRIPTION format"); + } + cfDragWindow = RegisterClipboardFormat(L"DragWindow"); + if (!cfDragWindow) { + OutputDebugString(L"PyShellExt::DllMain - failed to get DragWindow format"); + } + + DisableThreadLibraryCalls(hinst); + } + return TRUE; +}
\ No newline at end of file |