// 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);

if (!dd) {

OutputDebugString(L"PyShellExt::UpdateDropDescription - failed to lock DROPDESCRIPTION hGlobal");

ReleaseStgMedium(&medium);

return E_FAIL;

}

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;

}