source: trunk/eraser6/ShellExt/CtxMenu.cpp @ 1369

Revision 1369, 28.2 KB checked in by lowjoel, 5 years ago (diff)

Manually parse the string table to get the language resource we want when calling CCtxMenu::LoadString? as we separate the string tables according to language in our module resource. This allows us to specify the language we want the resources to load in (to be implemented).

  • Property svn:eol-style set to native
  • Property svn:keywords set to Id Rev
Line 
1/*
2 * $Id$
3 * Copyright 2008-2009 The Eraser Project
4 * Original Author: Kasra Nassiri <cjax@users.sourceforge.net>
5 * Modified By: Joel Low <lowjoel@users.sourceforge.net>
6 *
7 * This file is part of Eraser.
8 *
9 * Eraser is free software: you can redistribute it and/or modify it under the
10 * terms of the GNU General Public License as published by the Free Software
11 * Foundation, either version 3 of the License, or (at your option) any later
12 * version.
13 *
14 * Eraser is distributed in the hope that it will be useful, but WITHOUT ANY
15 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
16 * A PARTICULAR PURPOSE. See the GNU General Public License for more details.
17 *
18 * A copy of the GNU General Public License can be found at
19 * <http://www.gnu.org/licenses/>.
20 */
21
22#include "stdafx.h"
23#include "CtxMenu.h"
24#include "DllMain.h"
25#include <sstream>
26
27extern "C"
28{
29    typedef LONG NTSTATUS;
30    enum KEY_INFORMATION_CLASS
31    {
32        KeyBasicInformation,
33        KeyNodeInformation,
34        KeyFullInformation,
35        KeyNameInformation,
36        KeyCachedInformation,
37        KeyVirtualizationInformation
38    };
39
40    struct KEY_BASIC_INFORMATION
41    {
42        LARGE_INTEGER LastWriteTime;
43        ULONG TitleIndex;
44        ULONG NameLength;
45        WCHAR Name[1];
46    };
47
48    struct KEY_NODE_INFORMATION
49    {
50        LARGE_INTEGER LastWriteTime;
51        ULONG TitleIndex;
52        ULONG ClassOffset;
53        ULONG ClassLength;
54        ULONG NameLength;
55        WCHAR Name[1];
56    };
57
58    typedef NTSTATUS (__stdcall *pNtQueryKey)(HANDLE KeyHandle, KEY_INFORMATION_CLASS KeyInformationClass,
59        PVOID KeyInformation, ULONG Length, PULONG ResultLength);
60    pNtQueryKey NtQueryKey = NULL;
61}
62
63template<typename handleType> class Handle
64{
65public:
66    Handle()
67    {
68        Object = NULL;
69    }
70
71    Handle(handleType handle)
72    {
73        Object = handle;
74    }
75
76    ~Handle()
77    {
78        DeleteObject(Object);
79    }
80
81    operator handleType&()
82    {
83        return Object;
84    }
85
86private:
87    handleType Object;
88};
89
90Handle<HICON>::~Handle()
91{
92    DestroyIcon(Object);
93}
94
95Handle<HANDLE>::~Handle()
96{
97    CloseHandle(Object);
98}
99
100Handle<HKEY>::~Handle()
101{
102    CloseHandle(Object);
103}
104
105namespace Eraser {
106    HRESULT CCtxMenu::FinalConstruct()
107    {
108        //Initialise member variables.
109        MenuID = 0;
110        std::wstring menuTitle(LoadString(IDS_ERASER));
111        MenuTitle = new wchar_t[menuTitle.length() + 1];
112        wcscpy_s(MenuTitle, menuTitle.length() + 1, menuTitle.c_str());
113
114        //Check if the shell extension has been disabled.
115        Handle<HKEY> eraserKey;
116        LONG openKeyResult = RegOpenKeyEx(HKEY_CURRENT_USER,
117            L"Software\\Eraser\\Eraser 6\\3460478d-ed1b-4ecc-96c9-2ca0e8500557\\", 0,
118            KEY_READ, &static_cast<HKEY&>(eraserKey));
119
120        switch (openKeyResult)
121        {
122        case ERROR_FILE_NOT_FOUND:
123            //No settings defined: we default to enabling the shell extension.
124            return S_OK;
125
126        case ERROR_SUCCESS:
127            break;
128           
129        default:
130            return E_FAIL;
131        }
132
133        //Check the value of the IntegrateWithShell value.
134        DWORD value = 0;
135        DWORD valueType = 0;
136        DWORD valueSize = sizeof(value);
137        DWORD error = RegQueryValueEx(eraserKey, L"IntegrateWithShell", NULL, &valueType,
138            reinterpret_cast<BYTE*>(&value), &valueSize);
139        if (error == ERROR_SUCCESS && value == 0)
140        {
141            return E_FAIL;
142        }
143
144        return S_OK;
145    }
146
147    HRESULT CCtxMenu::FinalRelease()
148    {
149        delete[] MenuTitle;
150        return S_OK;
151    }
152
153    HRESULT CCtxMenu::Initialize(LPCITEMIDLIST pidlFolder, LPDATAOBJECT pDataObj,
154                                 HKEY hProgID)
155    {
156        //Determine where the shell extension was invoked from.
157        if (GetHKeyPath(hProgID) == L"{645FF040-5081-101B-9F08-00AA002F954E}")
158        {
159            InvokeReason = INVOKEREASON_RECYCLEBIN;
160
161            //We can't do much other processing: the LPDATAOBJECT parameter contains
162            //data that is a private type so we don't know how to query for it.
163            return S_OK;
164        }
165
166        //Check pidlFolder for the drop path, if it exists. This is for drag-and-drop
167        //context menus.
168        else if (pidlFolder != NULL)
169        {
170            InvokeReason = INVOKEREASON_DRAGDROP;
171
172            //Translate the drop path to a location on the filesystem.
173            wchar_t dropTargetPath[MAX_PATH];
174            if (!SHGetPathFromIDList(pidlFolder, dropTargetPath))
175                return E_FAIL;
176
177            DragDropDestinationDirectory = dropTargetPath;
178        }
179
180        //Okay, everything else is a simple context menu for a set of selected files/
181        //folders/drives.
182        else
183            InvokeReason = INVOKEREASON_FILEFOLDER;
184
185        //Look for CF_HDROP data in the data object.
186        FORMATETC fmt = { CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL };
187        STGMEDIUM stg = { TYMED_HGLOBAL };
188        if (FAILED(pDataObj->GetData(&fmt, &stg)))
189            //Nope! Return an "invalid argument" error back to Explorer.
190            return E_INVALIDARG;
191
192        //Get a pointer to the actual data.
193        HDROP hDrop = static_cast<HDROP>(GlobalLock(stg.hGlobal));
194        if (hDrop == NULL)
195            return E_INVALIDARG;
196
197        //Sanity check - make sure there is at least one filename.
198        UINT uNumFiles = DragQueryFile(hDrop, 0xFFFFFFFF, NULL, 0);
199        if (!uNumFiles)
200        {
201            GlobalUnlock(stg.hGlobal);
202            ReleaseStgMedium(&stg);
203            return E_INVALIDARG;
204        }
205
206        //Collect all the files which have been selected.
207        HRESULT hr = S_OK;
208        WCHAR buffer[MAX_PATH] = {0};
209        for (UINT i = 0; i < uNumFiles; i++)
210        {
211            UINT charsWritten = DragQueryFile(hDrop, i, buffer, sizeof(buffer) / sizeof(buffer[0]));
212            if (!charsWritten)
213            {
214                hr = E_INVALIDARG;
215                continue;
216            }
217
218            SelectedFiles.push_back(std::wstring(buffer, charsWritten));
219        }
220
221        //Clean up.
222        GlobalUnlock(stg.hGlobal);
223        ReleaseStgMedium(&stg);
224        return hr;
225    }
226
227    HRESULT CCtxMenu::QueryContextMenu(HMENU hmenu, UINT uMenuIndex, UINT uidFirstCmd,
228                                       UINT /*uidLastCmd*/, UINT uFlags)
229    {
230        //First check if we're running on Vista or later
231        bool isVistaOrLater = false;
232        {
233            //Set the bitmap for the registered item. Vista machines will be set using a DIB,
234            //older machines will be ownerdrawn.
235            OSVERSIONINFO osvi;
236            ZeroMemory(&osvi, sizeof(osvi));
237            osvi.dwOSVersionInfoSize = sizeof(osvi);
238
239            isVistaOrLater = GetVersionEx(&osvi) && osvi.dwPlatformId == VER_PLATFORM_WIN32_NT &&
240                osvi.dwMajorVersion >= 6;
241        }
242
243        //If the flags include CMF_DEFAULTONLY then we shouldn't do anything.
244        if (uFlags & CMF_DEFAULTONLY)
245            return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, 0);
246
247        //First, create and populate a submenu.
248        UINT uID = uidFirstCmd;
249        HMENU hSubmenu = CreatePopupMenu();
250
251        //Create the submenu, following the order defined in the CEraserLPVERB enum, creating
252        //only items which are applicable.
253        Actions applicableActions = GetApplicableActions();
254        VerbMenuIndices.clear();
255        if (applicableActions & ACTION_ERASE)
256        {
257            InsertMenu(hSubmenu, ACTION_ERASE, MF_BYPOSITION, uID++,
258                LoadString(IDS_ACTION_ERASE).c_str());              //Erase
259            VerbMenuIndices.push_back(ACTION_ERASE);
260        }
261        if (applicableActions & ACTION_ERASE_ON_RESTART)
262        {
263            InsertMenu(hSubmenu, ACTION_ERASE_ON_RESTART, MF_BYPOSITION, uID++,
264                LoadString(IDS_ACTION_ERASERESTART).c_str());       //Erase on Restart
265            VerbMenuIndices.push_back(ACTION_ERASE_ON_RESTART);
266        }
267        if (applicableActions & ACTION_ERASE_UNUSED_SPACE)
268        {
269            MENUITEMINFO mii = { sizeof(MENUITEMINFO) };
270            mii.wID = uID++;
271            mii.fMask = MIIM_STRING | MIIM_ID;
272
273            std::wstring str(LoadString(IDS_ACTION_ERASEUNUSEDSPACE));
274            std::vector<wchar_t> buffer(str.length() + 1);
275            wcscpy_s(&buffer.front(), str.length() + 1, str.c_str());
276            mii.dwTypeData = &buffer.front();
277
278            if (isVistaOrLater)
279            {
280                SHSTOCKICONINFO sii;
281                ::ZeroMemory(&sii, sizeof(sii));
282                sii.cbSize = sizeof(sii);
283
284                static HMODULE shellAPI = LoadLibrary(L"Shell32.dll");
285                typedef HRESULT (__stdcall *pSHGetStockIconInfo)(SHSTOCKICONID siid, UINT uFlags,
286                    SHSTOCKICONINFO* psii);
287                pSHGetStockIconInfo SHGetStockIconInfo = reinterpret_cast<pSHGetStockIconInfo>(
288                    GetProcAddress(shellAPI, "SHGetStockIconInfo"));
289
290                unsigned dimensions = GetSystemMetrics(SM_CXSMICON);
291                if (SUCCEEDED(SHGetStockIconInfo(SIID_SHIELD, SHGSI_ICON | SHGSI_SMALLICON, &sii)))
292                {
293                    Handle<HICON> icon(sii.hIcon);
294                    static HBITMAP dib = NULL;
295
296                    if (dib == NULL)
297                    {
298                        dib = CreateDIB(dimensions, dimensions, NULL);
299                        Handle<HDC> hdc(CreateCompatibleDC(NULL));
300                        SelectObject(hdc, dib);
301
302                        DrawIconEx(hdc, 0, 0, icon, dimensions, dimensions, 0, NULL, DI_NORMAL);
303                        SelectObject(hdc, NULL);
304                    }
305
306                    mii.hbmpItem = dib;
307                    mii.fMask |= MIIM_BITMAP;
308                }
309            }
310           
311            InsertMenuItem(hSubmenu, ACTION_ERASE_UNUSED_SPACE, MF_BYPOSITION, &mii);
312            VerbMenuIndices.push_back(ACTION_ERASE_UNUSED_SPACE);
313        }
314        //-------------------------------------------------------------------------
315        if (applicableActions & ACTION_SECURE_MOVE)
316        {
317            InsertSeparator(hSubmenu);
318            InsertMenu(hSubmenu, ACTION_SECURE_MOVE, MF_BYPOSITION, uID++,
319                LoadString(IDS_ACTION_SECUREMOVE).c_str());         //Secure Move
320            VerbMenuIndices.push_back(ACTION_SECURE_MOVE);
321        }
322
323        //Insert the submenu into the Context menu provided by Explorer.
324        {
325            MENUITEMINFO mii = { sizeof(MENUITEMINFO) };
326            mii.wID = uID++;
327            mii.fMask = MIIM_SUBMENU | MIIM_STRING | MIIM_ID;
328            mii.hSubMenu = hSubmenu;
329            mii.dwTypeData = const_cast<wchar_t*>(MenuTitle);
330
331            //Set the bitmap for the registered item. Vista machines will be set using a DIB,
332            //older machines will be ownerdrawn.
333            if (isVistaOrLater)
334            {
335                mii.fMask |= MIIM_BITMAP;
336                Handle<HICON> icon(GetMenuIcon());
337                mii.hbmpItem = GetMenuBitmapFromIcon(icon);
338            }
339            else if (InvokeReason != INVOKEREASON_DRAGDROP)
340            {
341                mii.fMask |= MIIM_FTYPE;
342                mii.fType = MFT_OWNERDRAW;
343            }
344
345            MenuID = uMenuIndex++;
346            InsertMenuItem(hmenu, MenuID, TRUE, &mii);
347
348            //Disable the menu item - IF the user selected the recycle bin AND the
349            //recycle bin is empty
350            if (InvokeReason == INVOKEREASON_RECYCLEBIN)
351            {
352                SHQUERYRBINFO sqrbi;
353                ::ZeroMemory(&sqrbi, sizeof(sqrbi));
354                sqrbi.cbSize = sizeof(sqrbi);
355                if (SUCCEEDED(SHQueryRecycleBin(NULL, &sqrbi)))
356                {
357                    EnableMenuItem(hmenu, MenuID, MF_BYPOSITION |
358                        ((sqrbi.i64NumItems != 0) ? MF_ENABLED : MF_DISABLED));
359                }
360            }
361        }
362
363        return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, uID - uidFirstCmd);
364    }
365
366    HRESULT CCtxMenu::HandleMenuMsg(UINT uMsg, WPARAM wParam, LPARAM lParam)
367    {
368        return HandleMenuMsg2(uMsg, wParam, lParam, NULL);
369    }
370
371    HRESULT CCtxMenu::HandleMenuMsg2(UINT uMsg, WPARAM /*wParam*/, LPARAM lParam,
372                                     LRESULT* result)
373    {
374        //Skip if we aren't handling our own.
375        bool handleResult = false;
376        switch (uMsg)
377        {
378        case WM_MEASUREITEM:
379            {
380                MEASUREITEMSTRUCT* mis = reinterpret_cast<MEASUREITEMSTRUCT*>(lParam);
381                if (mis->CtlID == MenuID)
382                    handleResult = OnMeasureItem(mis->itemWidth, mis->itemHeight);
383                else
384                    handleResult = false;
385                break;
386            }
387
388        case WM_DRAWITEM:
389            {
390                DRAWITEMSTRUCT* dis = reinterpret_cast<DRAWITEMSTRUCT*>(lParam);
391                if (dis->CtlID == MenuID)
392                    handleResult = OnDrawItem(dis->hDC, dis->rcItem, dis->itemAction, dis->itemState);
393                else
394                    handleResult = false;
395            }
396        }
397
398        if (result)
399            *result = handleResult;
400        return S_OK;
401    }
402
403    bool CCtxMenu::OnMeasureItem(UINT& itemWidth, UINT& itemHeight)
404    {
405        LOGFONT logFont;
406        if (!SystemParametersInfo(SPI_GETICONTITLELOGFONT, sizeof(logFont), &logFont, 0))
407            return false;
408
409        //Measure the size of the text.
410        Handle<HDC> screenDC = GetDC(NULL);
411        Handle<HFONT> font = CreateFontIndirect(&logFont);
412        SelectObject(screenDC, font);
413        SIZE textSize;
414        if (!GetTextExtentPoint32(screenDC, MenuTitle, static_cast<DWORD>(wcslen(MenuTitle)), &textSize))
415            return false;
416
417        itemWidth = textSize.cx;
418        itemHeight = textSize.cy;
419
420        //Account for the size of the bitmap.
421        UINT iconWidth = GetSystemMetrics(SM_CXMENUCHECK);
422        itemWidth += iconWidth;
423        itemHeight = std::max(iconWidth, itemHeight);
424
425        //And remember the minimum size for menu items.
426        itemHeight = std::max((int)itemHeight, GetSystemMetrics(SM_CXMENUSIZE));
427        return true;
428    }
429
430    bool CCtxMenu::OnDrawItem(HDC hdc, RECT rect, UINT /*action*/, UINT state)
431    {
432        //Draw the background.
433        LOGBRUSH logBrush = { BS_SOLID,
434            (state & ODS_SELECTED) ?
435                GetSysColor(COLOR_HIGHLIGHT) : GetSysColor(COLOR_MENU),
436            0
437        };
438        Handle<HBRUSH> bgBrush = CreateBrushIndirect(&logBrush);
439        FillRect(hdc, &rect, bgBrush);
440
441        //Then the bitmap.
442        {
443            //Draw the icon with alpha and all first.
444            Handle<HICON> icon(GetMenuIcon());
445            int iconSize = GetSystemMetrics(SM_CXMENUCHECK);
446            int iconMargin = GetSystemMetrics(SM_CXEDGE);
447            DrawState(hdc, NULL, NULL, reinterpret_cast<LPARAM>(static_cast<HICON>(icon)),
448                NULL, rect.left + iconMargin, rect.top + (rect.bottom - rect.top - iconSize) / 2,
449                0, 0, DST_ICON | ((state & ODS_DISABLED) ? DSS_DISABLED : 0));
450           
451            //Move the rectangle's left bound to the text starting position
452            rect.left += iconMargin * 2 + iconSize;
453        }
454       
455        //Draw the text.
456        SetBkMode(hdc, TRANSPARENT);
457        LOGFONT logFont;
458        if (!SystemParametersInfo(SPI_GETICONTITLELOGFONT, sizeof(logFont), &logFont, 0))
459            return false;
460
461        SIZE textSize;
462        if (!GetTextExtentPoint32(hdc, MenuTitle, static_cast<DWORD>(wcslen(MenuTitle)), &textSize))
463            return false;
464
465        COLORREF oldColour = SetTextColor(hdc, 
466            (state & ODS_DISABLED) ? GetSysColor(COLOR_GRAYTEXT) :          //Disabled menu item
467            (state & ODS_SELECTED) ? GetSysColor(COLOR_HIGHLIGHTTEXT) :     //Highlighted menu item
468                GetSysColor(COLOR_MENUTEXT));                               //Normal menu item
469        UINT flags = DST_PREFIXTEXT;
470        if (state & ODS_NOACCEL)
471            flags |= DSS_HIDEPREFIX;
472        ::DrawState(hdc, NULL, NULL, reinterpret_cast<LPARAM>(MenuTitle), wcslen(MenuTitle),
473            rect.left, rect.top + (rect.bottom - rect.top - textSize.cy) / 2, textSize.cx, textSize.cy, flags);
474        SetTextColor(hdc, oldColour);
475        return true;
476    }
477
478    HRESULT CCtxMenu::GetCommandString(UINT_PTR idCmd, UINT uFlags, UINT* /*pwReserved*/,
479                                       LPSTR pszName, UINT cchMax)
480    {
481        //We only know how to handle help string requests.
482        if (!(uFlags & GCS_HELPTEXT))
483            return E_INVALIDARG;
484
485        //Get the command string for the given id
486        if (idCmd >= VerbMenuIndices.size())
487            return E_INVALIDARG;
488
489        std::wstring commandString;
490        switch (VerbMenuIndices[idCmd])
491        {
492        case ACTION_ERASE:
493            commandString = LoadString(IDS_HELPSTRING_ERASE);
494            break;
495        case ACTION_ERASE_ON_RESTART:
496            commandString = LoadString(IDS_HELPSTRING_ERASEONRESTART);
497            break;
498        case ACTION_ERASE_UNUSED_SPACE:
499            commandString = LoadString(IDS_HELPSTRING_ERASEUNUSEDSPACE);
500            break;
501        case ACTION_SECURE_MOVE:
502            commandString = LoadString(IDS_HELPSTRING_SECUREMOVE);
503            break;
504        default:
505            //We don't know what action this is: return E_INVALIDARG.
506            return E_INVALIDARG;
507        }
508
509        //Return the help string to Explorer.
510        if (uFlags & GCS_UNICODE)
511            wcscpy_s(reinterpret_cast<wchar_t*>(pszName), cchMax, commandString.c_str());
512        else
513        {
514            size_t convCount = 0;
515            wcstombs_s(&convCount, pszName, cchMax, commandString.c_str(), commandString.length());
516        }
517
518        return S_OK;
519    }
520
521    HRESULT CCtxMenu::InvokeCommand(LPCMINVOKECOMMANDINFO pCmdInfo)
522    {
523        //If lpVerb really points to a string, ignore this function call and bail out.
524        if (HIWORD(pCmdInfo->lpVerb) != 0)
525            return E_INVALIDARG;
526
527        //If the verb index refers to an item outside the bounds of our VerbMenuIndices
528        //vector, exit.
529        if (LOWORD(pCmdInfo->lpVerb) >= VerbMenuIndices.size())
530            return E_INVALIDARG;
531
532        //Build the command line
533        std::wstring commandAction;
534        std::wstring commandLine;
535        bool commandElevate = false;
536        HRESULT result = E_INVALIDARG;
537        switch (VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)])
538        {
539        case ACTION_ERASE_ON_RESTART:
540            commandLine += L"-s=restart ";
541
542        case ACTION_ERASE:
543            {
544                //Add Task command.
545                commandAction = L"addtask";
546
547                //See the invocation context: if it is executed from the recycle bin
548                //then the list of selected files will be empty.
549                if (InvokeReason == INVOKEREASON_RECYCLEBIN)
550                {
551                    commandLine += L"-r";
552                }
553                else
554                {
555                    //Okay, we were called from an item right-click; iterate over every
556                    //selected file
557                    for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
558                        i != SelectedFiles.end(); ++i)
559                    {
560                        //Check if the current item is a file or folder.
561                        std::wstring item(*i);
562                        if (item.length() > 2 && item[item.length() - 1] == '\\')
563                            item.erase(item.end() - 1);
564                        DWORD attributes = GetFileAttributes(item.c_str());
565
566                        //Escape the string
567                        std::wstring escapedItem(EscapeString(item));
568
569                        //Add the correct command line for the file type.
570                        if (attributes & FILE_ATTRIBUTE_DIRECTORY)
571                            commandLine += L"\"-d=" + escapedItem + L"\" ";
572                        else
573                            commandLine += L"\"-f=" + escapedItem + L"\" ";
574                    }
575                }
576
577                break;
578            }
579
580        case ACTION_ERASE_UNUSED_SPACE:
581            {
582                //We want to add a new task
583                commandAction = L"addtask";
584
585                //Erasing unused space requires elevation
586                commandElevate = true;
587
588                //Add every item onto the command line
589                for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
590                    i != SelectedFiles.end(); ++i)
591                {
592                    std::wstring escapedItem(EscapeString(*i));
593                    commandLine += L"\"-u=" + escapedItem + L",clusterTips\" ";
594                }
595               
596                break;
597            }
598
599        default:
600            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
601            {
602                MessageBox(pCmdInfo->hwnd, FormatString(LoadString(IDS_ERROR_UNKNOWNACTION),
603                    VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)]).c_str(),
604                    LoadString(IDS_ERASERSHELLEXT).c_str(), MB_OK | MB_ICONERROR);
605            }
606        }
607
608        try
609        {
610            RunEraser(commandAction, commandLine, commandElevate, pCmdInfo->hwnd,
611                pCmdInfo->nShow);
612        }
613        catch (const std::wstring& e)
614        {
615            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
616            {
617                MessageBox(pCmdInfo->hwnd, e.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(),
618                    MB_OK | MB_ICONERROR);
619            }
620        }
621
622        return result;
623    }
624
625    CCtxMenu::Actions CCtxMenu::GetApplicableActions()
626    {
627        unsigned result = 0;
628       
629        //First decide the actions which are applicable to the current invocation
630        //reason.
631        switch (InvokeReason)
632        {
633        case INVOKEREASON_RECYCLEBIN:
634            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART;
635            break;
636        case INVOKEREASON_FILEFOLDER:
637            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE;
638#if 0
639        case INVOKEREASON_DRAGDROP:
640            result |= ACTION_SECURE_MOVE;
641#endif
642        }
643
644        //Remove actions that don't apply to the current invocation reason.
645        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
646            i != SelectedFiles.end(); ++i)
647        {
648            //Remove trailing slashes if they are directories.
649            std::wstring item(*i);
650
651            //Check if the path is a path to a volume, if it is not, remove the
652            //erase unused space verb.
653            wchar_t volumeName[MAX_PATH];
654            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName,
655                sizeof(volumeName) / sizeof(volumeName[0])))
656            {
657                result &= ~ACTION_ERASE_UNUSED_SPACE;
658            }
659        }
660
661        return static_cast<Actions>(result);
662    }
663
664    std::wstring CCtxMenu::LoadString(UINT stringID)
665    {
666        //Convert the resource ID to the block and item IDs.
667        UINT stringBlockId = (stringID >> 4) + 1;
668        UINT stringItemId = stringID % 16;
669
670        //Obtain a pointer to the memory holding the string table.
671        HRSRC resourceHandle = FindResourceEx(theApp.m_hInstance,
672            RT_STRING, MAKEINTRESOURCE(stringBlockId), MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL));
673        if (!resourceHandle)
674            AfxMessageBox(FormatError().c_str());
675
676        DWORD sizeOfResource = SizeofResource(theApp.m_hInstance, resourceHandle);
677        HGLOBAL resourceBlock = LoadResource(theApp.m_hInstance, resourceHandle);
678        if (!sizeOfResource || !resourceBlock)
679            AfxMessageBox(FormatError().c_str());
680
681        wchar_t* stringTable = reinterpret_cast<wchar_t*>(LockResource(resourceBlock));
682
683        //Iterate over the string table. The string table is null-delimited with
684        //the first byte storing the length of the string entry.
685        for ( ; stringItemId != 0; --stringItemId)
686        {
687            if (*stringTable == L'\0')
688                ++stringTable;
689            else
690                stringTable += *stringTable + 1;
691        }
692
693        return std::wstring(stringTable + 1, *stringTable);
694    }
695
696    std::wstring CCtxMenu::EscapeString(const std::wstring& string)
697    {
698        //Escape the command line (= and , are special characters)
699        std::wstring escapedItem;
700        escapedItem.reserve(string.length());
701        for (std::wstring::const_iterator i = string.begin(); i != string.end(); ++i)
702        {
703            if (wcschr(L"\\=,", *i))
704                escapedItem += '\\';
705            escapedItem += *i;
706        }
707
708        return escapedItem;
709    }
710
711    std::wstring CCtxMenu::FormatString(const std::wstring& formatString, ...)
712    {
713        std::vector<wchar_t> formatStr(formatString.length() + 1);
714        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str());
715
716        std::vector<wchar_t> buffer(formatStr.size());
717        for ( ; ; )
718        {
719            buffer.resize(buffer.size() * 2);
720            va_list arguments;
721            va_start(arguments, formatString);
722            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(),
723                arguments);
724            va_end(arguments);
725
726            if (result > 0 && static_cast<unsigned>(result) < buffer.size())
727            {
728                break;
729            }
730        }
731
732        //Return the result as a wstring
733        std::wstring result;
734        if (buffer.size() > 0)
735            result = &buffer.front();
736        return result;
737    }
738
739    std::wstring CCtxMenu::FormatError(DWORD lastError)
740    {
741        if (lastError == static_cast<DWORD>(-1))
742            lastError = GetLastError();
743
744        LPTSTR messageBuffer = NULL;
745        if (FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0,
746            lastError, 0, reinterpret_cast<LPWSTR>(&messageBuffer), 0, NULL) == 0)
747        {
748            return L"";
749        }
750
751        std::wstring result(messageBuffer);
752        LocalFree(messageBuffer);
753        return result;
754    }
755
756    std::wstring CCtxMenu::GetHKeyPath(HKEY handle)
757    {
758        if (NtQueryKey == NULL)
759            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress(
760                LoadLibrary(L"Ntdll.dll"), "NtQueryKey"));
761
762        //Keep querying for the key information until enough buffer space has been allocated.
763        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION));
764        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL;
765        ULONG keyInfoSize = 0;
766
767        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW)
768        {
769            buffer.resize(buffer.size() + keyInfoSize);
770            ZeroMemory(&buffer.front(), buffer.size());
771            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(),
772                static_cast<ULONG>(buffer.size()), &keyInfoSize);
773        }
774
775        if (queryResult != STATUS_SUCCESS)
776            return std::wstring();
777
778        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>(
779            &buffer.front());
780        return keyInfo->Name;
781    }
782
783    bool CCtxMenu::IsUserAdmin()
784    {
785        SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
786        PSID AdministratorsGroup;
787        if (AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID,
788            DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup))
789        {
790            BOOL result = false;
791            if (!CheckTokenMembership(NULL, AdministratorsGroup, &result))
792                result = false;
793
794            FreeSid(AdministratorsGroup);
795            return result != FALSE;
796        }
797
798        return false;
799    }
800
801    void CCtxMenu::RunEraser(const std::wstring& action, const std::wstring& parameters,
802        bool elevated, HWND parent, int show)
803    {
804        //Get the path to this DLL so we can look for Eraser.exe
805        wchar_t fileName[MAX_PATH];
806        DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName,
807            sizeof(fileName) / sizeof(fileName[0]));
808        if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0]))
809            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
810       
811        //Trim to the last \, then append Eraser.exe
812        std::wstring eraserPath(fileName, fileNameLength);
813        std::wstring::size_type lastBackslash = eraserPath.rfind('\\');
814        if (lastBackslash == std::wstring::npos)
815            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
816
817        eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end());
818        if (eraserPath.empty())
819            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
820
821        eraserPath += L"Eraser.exe";
822
823        std::wstring finalParameters;
824        {
825            std::wostringstream parametersStrm;
826            parametersStrm << "\"" << action << L"\" -q " << parameters;
827            finalParameters = parametersStrm.str();
828        }
829
830        //If the process must be elevated we use ShellExecute with the runas verb
831        //to elevate the new process.
832        if (elevated && !IsUserAdmin())
833        {
834            int result = reinterpret_cast<int>(ShellExecute(parent, L"runas",
835                eraserPath.c_str(), finalParameters.c_str(), NULL, show));
836            if (result <= 32)
837                switch (result)
838                {
839                case SE_ERR_ACCESSDENIED:
840                    throw LoadString(IDS_ERROR_ACCESSDENIED);
841                default:
842                    throw LoadString(IDS_ERROR_UNKNOWN);
843                }
844        }
845
846        //If the process isn't to be elevated, we use CreateProcess so we can get
847        //read the output from the child process
848        else
849        {
850            //Create the process.
851            STARTUPINFO startupInfo;
852            ZeroMemory(&startupInfo, sizeof(startupInfo));
853            startupInfo.cb = sizeof(startupInfo);
854            startupInfo.dwFlags = STARTF_USESHOWWINDOW;
855            startupInfo.wShowWindow = static_cast<WORD>(show);
856            startupInfo.hStdInput = startupInfo.hStdOutput = startupInfo.hStdError =
857                INVALID_HANDLE_VALUE;
858
859            //Create handles for output redirection
860            Handle<HANDLE> readPipe;
861            HANDLE writePipe;
862            SECURITY_ATTRIBUTES security;
863            ZeroMemory(&security, sizeof(security));
864            security.nLength = sizeof(security);
865            security.lpSecurityDescriptor = NULL;
866            security.bInheritHandle = true;
867
868            if (CreatePipe(&static_cast<HANDLE&>(readPipe), &writePipe, &security, 0))
869            {
870                startupInfo.dwFlags |= STARTF_USESTDHANDLES;
871                startupInfo.hStdOutput = startupInfo.hStdError =
872                    writePipe;
873            }
874
875            PROCESS_INFORMATION processInfo;
876            ZeroMemory(&processInfo, sizeof(processInfo));
877            std::vector<wchar_t> buffer(eraserPath.length() + finalParameters.length() + 4);
878            wcscpy_s(&buffer.front(), buffer.size(), (L"\"" + eraserPath + L"\" " +
879                finalParameters).c_str());
880
881            if (!CreateProcess(NULL, &buffer.front(), NULL, NULL, true, CREATE_NO_WINDOW,
882                NULL, NULL, &startupInfo, &processInfo))
883            {
884                if (GetLastError() == ERROR_FILENAME_EXCED_RANGE)
885                    throw FormatString(LoadString(IDS_ERROR_TOO_MANY_FILES));
886                else
887                    throw FormatString(LoadString(IDS_ERROR_MISC), FormatError().c_str());
888            }
889
890            //Wait for the process to finish.
891            Handle<HANDLE> hProcess(processInfo.hProcess),
892                           hThread(processInfo.hThread);
893            CloseHandle(writePipe);
894            WaitForSingleObject(hProcess, static_cast<DWORD>(-1));
895            DWORD exitCode = 0;
896           
897            if (GetExitCodeProcess(processInfo.hProcess, &exitCode))
898                if (exitCode == ERROR_ACCESS_DENIED)
899                {
900                    //The spawned instance could not connect with the master instance
901                    //because it is running as an administrator. Spawn the new instance
902                    //again, this time as an administrator
903                    RunEraser(action, parameters, true, parent, show);
904                }
905                else if (exitCode)
906                {
907                    char buffer[8192];
908                    DWORD lastRead = 0;
909                    std::wstring output;
910
911                    while (ReadFile(readPipe, buffer, sizeof(buffer), &lastRead, NULL) && lastRead != 0)
912                    {
913                        size_t lastConvert = 0;
914                        wchar_t wBuffer[8192];
915                        if (!mbstowcs_s(&lastConvert, wBuffer, sizeof(wBuffer) / sizeof(wBuffer[0]),
916                            buffer, lastRead))
917                        {
918                            output += std::wstring(wBuffer, lastConvert);
919                        }
920                    }
921
922                    //Show the error message.
923                    throw output;
924                }
925        }
926    }
927
928    void CCtxMenu::InsertSeparator(HMENU menu)
929    {
930        MENUITEMINFO mii;
931        mii.cbSize = sizeof(MENUITEMINFO);
932        mii.fMask = MIIM_TYPE;
933        mii.fType = MF_SEPARATOR;
934        InsertMenuItem(menu, 0, false, &mii);
935    }
936
937    HICON CCtxMenu::GetMenuIcon()
938    {
939        int smIconSize = GetSystemMetrics(SM_CXMENUCHECK);
940        return static_cast<HICON>(LoadImage(theApp.m_hInstance, L"Eraser",
941            IMAGE_ICON, smIconSize, smIconSize, LR_DEFAULTCOLOR));
942    }
943
944    HBITMAP CCtxMenu::GetMenuBitmapFromIcon(HICON icon)
945    {
946        BITMAP bitmap;
947        ICONINFO iconInfo;
948        ZeroMemory(&bitmap, sizeof(bitmap));
949        ZeroMemory(&iconInfo, sizeof(iconInfo));
950
951        //Try to get the icon's size, bitmap and bit depth. We will try to convert
952        //the bitmap into a DIB for display on Vista if it contains an alpha channel.
953        if (!GetIconInfo(icon, &iconInfo))
954            return NULL;
955
956        Handle<HBITMAP> iconMask(iconInfo.hbmMask);
957        if (!GetObject(iconInfo.hbmColor, sizeof(BITMAP), &bitmap) ||
958            bitmap.bmBitsPixel < 32)
959            return iconInfo.hbmColor;
960
961        //Draw the icon onto the DIB which will preseve its alpha values
962        Handle<HDC> hdcDest = CreateCompatibleDC(NULL);
963        HBITMAP dib = CreateDIB(bitmap.bmWidth, bitmap.bmHeight, NULL);
964        SelectObject(hdcDest, dib);
965
966        Handle<HBITMAP> iconBitmap(iconInfo.hbmColor);
967        DrawIconEx(hdcDest, 0, 0, icon, bitmap.bmWidth, bitmap.bmHeight, 0, NULL, DI_NORMAL);
968        return dib;
969    }
970
971    HBITMAP CCtxMenu::CreateDIB(LONG width, LONG height, char** bitmapBits)
972    {
973        BITMAPINFO info;
974        ZeroMemory(&info, sizeof(info));
975        info.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
976        info.bmiHeader.biWidth = width;
977        info.bmiHeader.biHeight = height;
978        info.bmiHeader.biPlanes = 1;
979        info.bmiHeader.biBitCount = 32;
980
981        Handle<HDC> screenDC(GetDC(NULL));
982        return ::CreateDIBSection(screenDC, &info, DIB_RGB_COLORS,
983            reinterpret_cast<void**>(bitmapBits), NULL, 0);
984    }
985}
Note: See TracBrowser for help on using the repository browser.