source: branches/eraser6/6.0/ShellExt/CtxMenu.cpp @ 2948

Revision 2727, 29.2 KB checked in by lowjoel, 2 years ago (diff)

Only create a menu item for Eraser in the Explorer Context Menu if we've got submenu items to display.

  • Property svn:eol-style set to native
  • Property svn:keywords set to Id Rev
Line 
1/*
2 * $Id$
3 * Copyright 2008-2010 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 (!pDataObj || 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        //Check if the selected files is only one item long and if that item is the
222        //Start button (for Windows XP)
223        if (SelectedFiles.size() == 1)
224        {
225            wchar_t startMenuPath[MAX_PATH];
226            if (SUCCEEDED(SHGetFolderPath(NULL, CSIDL_STARTMENU, NULL,
227                SHGFP_TYPE_CURRENT, startMenuPath)))
228            {
229                if (SelectedFiles.front() == startMenuPath)
230                    //Yes, it is. Don't display the Eraser context menu for this.
231                    return E_INVALIDARG;
232            }
233        }
234
235        //Clean up.
236        GlobalUnlock(stg.hGlobal);
237        ReleaseStgMedium(&stg);
238        return hr;
239    }
240
241    HRESULT CCtxMenu::QueryContextMenu(HMENU hmenu, UINT uMenuIndex, UINT uidFirstCmd,
242                                       UINT /*uidLastCmd*/, UINT uFlags)
243    {
244        //First check if we're running on Vista or later
245        bool isVistaOrLater = false;
246        {
247            //Set the bitmap for the registered item. Vista machines will be set using a DIB,
248            //older machines will be ownerdrawn.
249            OSVERSIONINFO osvi;
250            ZeroMemory(&osvi, sizeof(osvi));
251            osvi.dwOSVersionInfoSize = sizeof(osvi);
252
253            isVistaOrLater = GetVersionEx(&osvi) && osvi.dwPlatformId == VER_PLATFORM_WIN32_NT &&
254                osvi.dwMajorVersion >= 6;
255        }
256
257        //If the flags include CMF_DEFAULTONLY then we shouldn't do anything.
258        if (uFlags & CMF_DEFAULTONLY)
259            return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, 0);
260
261        //If the flags include CMF_VERBSONLY then we shouldn't do anything as we do not
262        //want to operate on the target of a shortcut implicitly.
263        if (uFlags & CMF_VERBSONLY)
264            return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, 0);
265
266        //First, create and populate a submenu.
267        UINT uID = uidFirstCmd;
268        HMENU hSubmenu = CreatePopupMenu();
269
270        //Create the submenu, following the order defined in the CEraserLPVERB enum, creating
271        //only items which are applicable.
272        Actions applicableActions = GetApplicableActions();
273        VerbMenuIndices.clear();
274        if (applicableActions & ACTION_ERASE)
275        {
276            InsertMenu(hSubmenu, ACTION_ERASE, MF_BYPOSITION, uID++,
277                LoadString(IDS_ACTION_ERASE).c_str());              //Erase
278            VerbMenuIndices.push_back(ACTION_ERASE);
279        }
280        if (applicableActions & ACTION_ERASE_ON_RESTART)
281        {
282            InsertMenu(hSubmenu, ACTION_ERASE_ON_RESTART, MF_BYPOSITION, uID++,
283                LoadString(IDS_ACTION_ERASERESTART).c_str());       //Erase on Restart
284            VerbMenuIndices.push_back(ACTION_ERASE_ON_RESTART);
285        }
286        if (applicableActions & ACTION_ERASE_UNUSED_SPACE)
287        {
288            MENUITEMINFO mii = { sizeof(MENUITEMINFO) };
289            mii.wID = uID++;
290            mii.fMask = MIIM_STRING | MIIM_ID;
291
292            std::wstring str(LoadString(IDS_ACTION_ERASEUNUSEDSPACE));
293            std::vector<wchar_t> buffer(str.length() + 1);
294            wcscpy_s(&buffer.front(), str.length() + 1, str.c_str());
295            mii.dwTypeData = &buffer.front();
296
297            if (isVistaOrLater)
298            {
299                SHSTOCKICONINFO sii;
300                ::ZeroMemory(&sii, sizeof(sii));
301                sii.cbSize = sizeof(sii);
302
303                static HMODULE shellAPI = LoadLibrary(L"Shell32.dll");
304                typedef HRESULT (__stdcall *pSHGetStockIconInfo)(SHSTOCKICONID siid, UINT uFlags,
305                    SHSTOCKICONINFO* psii);
306                pSHGetStockIconInfo SHGetStockIconInfo = reinterpret_cast<pSHGetStockIconInfo>(
307                    GetProcAddress(shellAPI, "SHGetStockIconInfo"));
308
309                unsigned dimensions = GetSystemMetrics(SM_CXSMICON);
310                if (SUCCEEDED(SHGetStockIconInfo(SIID_SHIELD, SHGSI_ICON | SHGSI_SMALLICON, &sii)))
311                {
312                    Handle<HICON> icon(sii.hIcon);
313                    static HBITMAP dib = NULL;
314
315                    if (dib == NULL)
316                    {
317                        dib = CreateDIB(dimensions, dimensions, NULL);
318                        Handle<HDC> hdc(CreateCompatibleDC(NULL));
319                        SelectObject(hdc, dib);
320
321                        DrawIconEx(hdc, 0, 0, icon, dimensions, dimensions, 0, NULL, DI_NORMAL);
322                        SelectObject(hdc, NULL);
323                    }
324
325                    mii.hbmpItem = dib;
326                    mii.fMask |= MIIM_BITMAP;
327                }
328            }
329           
330            InsertMenuItem(hSubmenu, ACTION_ERASE_UNUSED_SPACE, MF_BYPOSITION, &mii);
331            VerbMenuIndices.push_back(ACTION_ERASE_UNUSED_SPACE);
332        }
333        //-------------------------------------------------------------------------
334        if (applicableActions & ACTION_SECURE_MOVE)
335        {
336            InsertSeparator(hSubmenu);
337            InsertMenu(hSubmenu, ACTION_SECURE_MOVE, MF_BYPOSITION, uID++,
338                LoadString(IDS_ACTION_SECUREMOVE).c_str());         //Secure Move
339            VerbMenuIndices.push_back(ACTION_SECURE_MOVE);
340        }
341
342        //Insert the submenu into the Context menu provided by Explorer.
343        if (GetMenuItemCount(hSubmenu) > 0)
344        {
345            MENUITEMINFO mii = { sizeof(MENUITEMINFO) };
346            mii.wID = uID++;
347            mii.fMask = MIIM_SUBMENU | MIIM_STRING | MIIM_ID | MIIM_BITMAP;
348            mii.hSubMenu = hSubmenu;
349            mii.dwTypeData = const_cast<wchar_t*>(MenuTitle);
350            MenuID = mii.wID;
351
352            //Set the bitmap for the registered item. Vista machines will be set using a DIB,
353            //older machines will be ownerdrawn.
354            if (isVistaOrLater)
355            {
356                Handle<HICON> icon(GetMenuIcon());
357                mii.hbmpItem = GetMenuBitmapFromIcon(icon);
358            }
359            else if (InvokeReason != INVOKEREASON_DRAGDROP)
360            {
361                mii.hbmpItem = HBMMENU_CALLBACK;
362            }
363
364            UINT menuIndex = uMenuIndex++;
365            InsertMenuItem(hmenu, menuIndex, TRUE, &mii);
366
367            //Disable the menu item - IF the user selected the recycle bin AND the
368            //recycle bin is empty
369            if (InvokeReason == INVOKEREASON_RECYCLEBIN)
370            {
371                SHQUERYRBINFO sqrbi;
372                ::ZeroMemory(&sqrbi, sizeof(sqrbi));
373                sqrbi.cbSize = sizeof(sqrbi);
374                if (SUCCEEDED(SHQueryRecycleBin(NULL, &sqrbi)))
375                {
376                    EnableMenuItem(hmenu, menuIndex, MF_BYPOSITION |
377                        ((sqrbi.i64NumItems != 0) ? MF_ENABLED : MF_DISABLED));
378                }
379            }
380        }
381
382        return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, uID - uidFirstCmd);
383    }
384
385    HRESULT CCtxMenu::HandleMenuMsg(UINT uMsg, WPARAM wParam, LPARAM lParam)
386    {
387        return HandleMenuMsg2(uMsg, wParam, lParam, NULL);
388    }
389
390    HRESULT CCtxMenu::HandleMenuMsg2(UINT uMsg, WPARAM /*wParam*/, LPARAM lParam,
391                                     LRESULT* result)
392    {
393        //Skip if we aren't handling our own.
394        bool handleResult = false;
395        switch (uMsg)
396        {
397        case WM_MEASUREITEM:
398            {
399                MEASUREITEMSTRUCT* mis = reinterpret_cast<MEASUREITEMSTRUCT*>(lParam);
400                if (mis->itemID == MenuID)
401                    handleResult = OnMeasureItem(mis->itemWidth, mis->itemHeight);
402                break;
403            }
404
405        case WM_DRAWITEM:
406            {
407                DRAWITEMSTRUCT* dis = reinterpret_cast<DRAWITEMSTRUCT*>(lParam);
408                if (dis->itemID == MenuID)
409                    handleResult = OnDrawItem(dis->hDC, dis->rcItem, dis->itemAction, dis->itemState);
410            }
411        }
412
413        if (result)
414            *result = handleResult;
415        return S_OK;
416    }
417
418    bool CCtxMenu::OnMeasureItem(UINT& itemWidth, UINT& itemHeight)
419    {
420        //Account for the size of the bitmap.
421        itemWidth = 0;
422        itemHeight = std::max<UINT>(GetSystemMetrics(SM_CYMENUCHECK), itemHeight);
423        return true;
424    }
425
426    bool CCtxMenu::OnDrawItem(HDC hdc, RECT rect, UINT /*action*/, UINT state)
427    {
428        //Get the icon and calculate its size.
429        Handle<HICON> icon(GetMenuIcon());
430        int iconSize = GetSystemMetrics(SM_CXMENUCHECK);
431        int iconMargin = GetSystemMetrics(SM_CXEDGE);
432
433        //Draw the bitmap.
434        DrawState(hdc, NULL, NULL, reinterpret_cast<LPARAM>(static_cast<HICON>(icon)),
435            NULL, rect.left - iconMargin - iconSize,
436            rect.top + (rect.bottom - rect.top - iconSize) / 2, 0, 0,
437            DST_ICON | ((state & ODS_DISABLED) ? DSS_DISABLED : 0));
438
439        return true;
440    }
441
442    HRESULT CCtxMenu::GetCommandString(UINT_PTR idCmd, UINT uFlags, UINT* /*pwReserved*/,
443                                       LPSTR pszName, UINT cchMax)
444    {
445        //We only know how to handle help string requests.
446        if (!(uFlags & GCS_HELPTEXT))
447            return E_INVALIDARG;
448
449        //Get the command string for the given id
450        if (idCmd >= VerbMenuIndices.size())
451            return E_INVALIDARG;
452
453        std::wstring commandString;
454        switch (VerbMenuIndices[idCmd])
455        {
456        case ACTION_ERASE:
457            commandString = LoadString(IDS_HELPSTRING_ERASE);
458            break;
459        case ACTION_ERASE_ON_RESTART:
460            commandString = LoadString(IDS_HELPSTRING_ERASEONRESTART);
461            break;
462        case ACTION_ERASE_UNUSED_SPACE:
463            commandString = LoadString(IDS_HELPSTRING_ERASEUNUSEDSPACE);
464            break;
465        case ACTION_SECURE_MOVE:
466            commandString = LoadString(IDS_HELPSTRING_SECUREMOVE);
467            break;
468        default:
469            //We don't know what action this is: return E_INVALIDARG.
470            return E_INVALIDARG;
471        }
472
473        //Return the help string to Explorer.
474        if (uFlags & GCS_UNICODE)
475            wcscpy_s(reinterpret_cast<wchar_t*>(pszName), cchMax, commandString.c_str());
476        else
477        {
478            size_t convCount = 0;
479            wcstombs_s(&convCount, pszName, cchMax, commandString.c_str(), commandString.length());
480        }
481
482        return S_OK;
483    }
484
485    HRESULT CCtxMenu::InvokeCommand(LPCMINVOKECOMMANDINFO pCmdInfo)
486    {
487        //If lpVerb really points to a string, ignore this function call and bail out.
488        if (HIWORD(pCmdInfo->lpVerb) != 0)
489            return E_INVALIDARG;
490
491        //If the verb index refers to an item outside the bounds of our VerbMenuIndices
492        //vector, exit.
493        if (LOWORD(pCmdInfo->lpVerb) >= VerbMenuIndices.size())
494            return E_INVALIDARG;
495
496        //Build the command line
497        std::wstring commandAction;
498        std::wstring commandLine;
499        bool commandElevate = false;
500        HRESULT result = E_INVALIDARG;
501        switch (VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)])
502        {
503        case ACTION_ERASE_ON_RESTART:
504            commandLine += L"-s=restart ";
505
506        case ACTION_ERASE:
507            {
508                //Add Task command.
509                commandAction = L"addtask";
510
511                //See the invocation context: if it is executed from the recycle bin
512                //then the list of selected files will be empty.
513                if (InvokeReason == INVOKEREASON_RECYCLEBIN)
514                {
515                    commandLine += L"-r";
516                }
517                else
518                {
519                    //Okay, we were called from an item right-click; iterate over every
520                    //selected file
521                    for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
522                        i != SelectedFiles.end(); ++i)
523                    {
524                        //Check if the current item is a file or folder.
525                        std::wstring item(*i);
526                        if (item.length() > 2 && item[item.length() - 1] == '\\')
527                            item.erase(item.end() - 1);
528                        DWORD attributes = GetFileAttributes(item.c_str());
529
530                        //Escape the string
531                        std::wstring escapedItem(EscapeString(item));
532
533                        //Add the correct command line for the file type.
534                        if (attributes & FILE_ATTRIBUTE_DIRECTORY)
535                            commandLine += L"\"-d=" + escapedItem + L"\" ";
536                        else
537                            commandLine += L"\"-f=" + escapedItem + L"\" ";
538                    }
539                }
540
541                break;
542            }
543
544        case ACTION_ERASE_UNUSED_SPACE:
545            {
546                //We want to add a new task
547                commandAction = L"addtask";
548
549                //Erasing unused space requires elevation
550                commandElevate = true;
551
552                //Add every item onto the command line
553                for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
554                    i != SelectedFiles.end(); ++i)
555                {
556                    std::wstring escapedItem(EscapeString(*i));
557                    commandLine += L"\"-u=" + escapedItem + L",clusterTips\" ";
558                }
559               
560                break;
561            }
562
563        default:
564            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
565            {
566                MessageBox(pCmdInfo->hwnd, FormatString(LoadString(IDS_ERROR_UNKNOWNACTION),
567                    VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)]).c_str(),
568                    LoadString(IDS_ERASERSHELLEXT).c_str(), MB_OK | MB_ICONERROR);
569            }
570        }
571
572        try
573        {
574            RunEraser(commandAction, commandLine, commandElevate, pCmdInfo->hwnd,
575                pCmdInfo->nShow);
576        }
577        catch (const std::wstring& e)
578        {
579            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
580            {
581                MessageBox(pCmdInfo->hwnd, e.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(),
582                    MB_OK | MB_ICONERROR);
583            }
584        }
585
586        return result;
587    }
588
589    CCtxMenu::Actions CCtxMenu::GetApplicableActions()
590    {
591        unsigned result = 0;
592       
593        //First decide the actions which are applicable to the current invocation
594        //reason.
595        switch (InvokeReason)
596        {
597        case INVOKEREASON_RECYCLEBIN:
598            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART;
599            break;
600        case INVOKEREASON_FILEFOLDER:
601            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE;
602#if 0
603        case INVOKEREASON_DRAGDROP:
604            result |= ACTION_SECURE_MOVE;
605#endif
606        }
607
608        //Remove actions that don't apply to the current invocation reason.
609        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
610            i != SelectedFiles.end(); ++i)
611        {
612            //Remove trailing slashes if they are directories.
613            std::wstring item(*i);
614
615            //Check if the path is a path to a volume, if it is not, remove the
616            //erase unused space verb.
617            wchar_t volumeName[MAX_PATH];
618            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName,
619                sizeof(volumeName) / sizeof(volumeName[0])))
620            {
621                result &= ~ACTION_ERASE_UNUSED_SPACE;
622            }
623        }
624
625        return static_cast<Actions>(result);
626    }
627
628    LCID LocaleNameToLCID(const std::wstring& localeName)
629    {
630        LCID result = LOCALE_USER_DEFAULT;
631        IMultiLanguage* multiLanguage = NULL;
632
633        //Create an instance of the IMultiLanguage interface
634        if (SUCCEEDED(CoCreateInstance(CLSID_CMultiLanguage, NULL, CLSCTX_ALL,
635                IID_IMultiLanguage, (void**)&multiLanguage)))
636        {
637            //Convert our locale name to a BString
638            BSTR localeNameBStr = SysAllocString(localeName.c_str());
639            if (localeNameBStr)
640                multiLanguage->GetLcidFromRfc1766(&result, localeNameBStr);
641            SysFreeString(localeNameBStr);
642
643            //Clean up
644            multiLanguage->Release();
645        }
646
647        return result;
648    }
649
650    std::wstring CCtxMenu::LoadString(UINT stringID)
651    {
652        //Convert the resource ID to the block and item IDs.
653        UINT stringBlockID = (stringID >> 4) + 1;
654        UINT stringItemID = stringID % 16;
655        WORD langID = LANG_USER_DEFAULT;
656        std::wstring localeName;
657
658        if (localeName.empty())
659        {
660            bool foundLanguage = false;
661            Handle<HKEY> eraserKey;
662            LONG openKeyResult = RegOpenKeyEx(HKEY_CURRENT_USER,
663                L"Software\\Eraser\\Eraser 6\\3460478d-ed1b-4ecc-96c9-2ca0e8500557\\", 0,
664                KEY_READ, &static_cast<HKEY&>(eraserKey));
665
666            if (openKeyResult != ERROR_FILE_NOT_FOUND)
667            {
668                //Check the value of the Language value.
669                std::vector<wchar_t> value(256);
670                while (!foundLanguage)
671                {
672                    DWORD valueType = 0;
673                    DWORD valueSize = value.size();
674                    DWORD error = RegQueryValueEx(eraserKey, L"Language", NULL, &valueType,
675                        reinterpret_cast<BYTE*>(&value.front()), &valueSize);
676
677                    if (error == ERROR_SUCCESS)
678                        foundLanguage = true;
679                    else if (error == ERROR_INSUFFICIENT_BUFFER)
680                        value.resize(value.size() * 2);
681                    else
682                        break;
683                }
684
685                if (foundLanguage)
686                {
687                    localeName.assign(value.begin(), value.end());
688                    langID = LANGIDFROMLCID(LocaleNameToLCID(localeName));
689                }
690            }
691        }
692
693        //Obtain a pointer to the memory holding the string table.
694        WORD langIDStack[] = { langID, PRIMARYLANGID(langID), LANG_USER_DEFAULT, LANG_ENGLISH };
695        HRSRC resourceHandle = NULL;
696        for (size_t i = 0; resourceHandle == NULL &&
697            i < sizeof(langIDStack) / sizeof(langIDStack[0]); ++i)
698        {
699            resourceHandle = FindResourceEx(theApp.m_hInstance,
700                RT_STRING, MAKEINTRESOURCE(stringBlockID), langIDStack[i]);
701            if (GetLastError() == ERROR_RESOURCE_LANG_NOT_FOUND)
702                continue;
703        }
704
705        if (resourceHandle == NULL)
706        {
707            AfxMessageBox(FormatError().c_str());
708            return std::wstring();
709        }
710
711        DWORD sizeOfResource = SizeofResource(theApp.m_hInstance, resourceHandle);
712        HGLOBAL resourceBlock = LoadResource(theApp.m_hInstance, resourceHandle);
713        if (!sizeOfResource || !resourceBlock)
714            AfxMessageBox(FormatError().c_str());
715
716        wchar_t* stringTable = reinterpret_cast<wchar_t*>(LockResource(resourceBlock));
717
718        //Iterate over the string table. The string table is null-delimited with
719        //the first byte storing the length of the string entry.
720        for ( ; stringItemID != 0; --stringItemID)
721        {
722            if (*stringTable == L'\0')
723                ++stringTable;
724            else
725                stringTable += *stringTable + 1;
726        }
727
728        return std::wstring(stringTable + 1, *stringTable);
729    }
730
731    std::wstring CCtxMenu::EscapeString(const std::wstring& string)
732    {
733        //Escape the command line (= and , are special characters)
734        std::wstring escapedItem;
735        escapedItem.reserve(string.length());
736        for (std::wstring::const_iterator i = string.begin(); i != string.end(); ++i)
737        {
738            if (wcschr(L"\\=,", *i))
739                escapedItem += '\\';
740            escapedItem += *i;
741        }
742
743        return escapedItem;
744    }
745
746    std::wstring CCtxMenu::FormatString(const std::wstring& formatString, ...)
747    {
748        std::vector<wchar_t> formatStr(formatString.length() + 1);
749        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str());
750
751        std::vector<wchar_t> buffer(formatStr.size());
752        for ( ; ; )
753        {
754            buffer.resize(buffer.size() * 2);
755            va_list arguments;
756            va_start(arguments, formatString);
757            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(),
758                arguments);
759            va_end(arguments);
760
761            if (result > 0 && static_cast<unsigned>(result) < buffer.size())
762            {
763                break;
764            }
765        }
766
767        //Return the result as a wstring
768        std::wstring result;
769        if (buffer.size() > 0)
770            result = &buffer.front();
771        return result;
772    }
773
774    std::wstring CCtxMenu::FormatError(DWORD lastError)
775    {
776        if (lastError == static_cast<DWORD>(-1))
777            lastError = GetLastError();
778
779        LPTSTR messageBuffer = NULL;
780        if (FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, 0,
781            lastError, 0, reinterpret_cast<LPWSTR>(&messageBuffer), 0, NULL) == 0)
782        {
783            return L"";
784        }
785
786        std::wstring result(messageBuffer);
787        LocalFree(messageBuffer);
788        return result;
789    }
790
791    std::wstring CCtxMenu::GetHKeyPath(HKEY handle)
792    {
793        if (NtQueryKey == NULL)
794            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress(
795                LoadLibrary(L"Ntdll.dll"), "NtQueryKey"));
796
797        //Keep querying for the key information until enough buffer space has been allocated.
798        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION));
799        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL;
800        ULONG keyInfoSize = 0;
801
802        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW)
803        {
804            buffer.resize(buffer.size() + keyInfoSize);
805            ZeroMemory(&buffer.front(), buffer.size());
806            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(),
807                static_cast<ULONG>(buffer.size()), &keyInfoSize);
808        }
809
810        if (queryResult != STATUS_SUCCESS)
811            return std::wstring();
812
813        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>(
814            &buffer.front());
815        return keyInfo->Name;
816    }
817
818    bool CCtxMenu::IsUserAdmin()
819    {
820        SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
821        PSID AdministratorsGroup;
822        if (AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID,
823            DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup))
824        {
825            BOOL result = false;
826            if (!CheckTokenMembership(NULL, AdministratorsGroup, &result))
827                result = false;
828
829            FreeSid(AdministratorsGroup);
830            return result != FALSE;
831        }
832
833        return false;
834    }
835
836    void CCtxMenu::RunEraser(const std::wstring& action, const std::wstring& parameters,
837        bool elevated, HWND parent, int show)
838    {
839        //Get the path to this DLL so we can look for Eraser.exe
840        wchar_t fileName[MAX_PATH];
841        DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName,
842            sizeof(fileName) / sizeof(fileName[0]));
843        if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0]))
844            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
845       
846        //Trim to the last \, then append Eraser.exe
847        std::wstring eraserPath(fileName, fileNameLength);
848        std::wstring::size_type lastBackslash = eraserPath.rfind('\\');
849        if (lastBackslash == std::wstring::npos)
850            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
851
852        eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end());
853        if (eraserPath.empty())
854            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
855
856        eraserPath += L"Eraser.exe";
857
858        std::wstring finalParameters;
859        {
860            std::wostringstream parametersStrm;
861            parametersStrm << "\"" << action << L"\" -q " << parameters;
862            finalParameters = parametersStrm.str();
863        }
864
865        //If the process must be elevated we use ShellExecute with the runas verb
866        //to elevate the new process.
867        if (elevated && !IsUserAdmin())
868        {
869            int result = reinterpret_cast<int>(ShellExecute(parent, L"runas",
870                eraserPath.c_str(), finalParameters.c_str(), NULL, show));
871            if (result <= 32)
872                switch (result)
873                {
874                case SE_ERR_ACCESSDENIED:
875                    throw LoadString(IDS_ERROR_ACCESSDENIED);
876                default:
877                    throw LoadString(IDS_ERROR_UNKNOWN);
878                }
879        }
880
881        //If the process isn't to be elevated, we use CreateProcess so we can get
882        //read the output from the child process
883        else
884        {
885            //Create the process.
886            STARTUPINFO startupInfo;
887            ZeroMemory(&startupInfo, sizeof(startupInfo));
888            startupInfo.cb = sizeof(startupInfo);
889            startupInfo.dwFlags = STARTF_USESHOWWINDOW;
890            startupInfo.wShowWindow = static_cast<WORD>(show);
891            startupInfo.hStdInput = startupInfo.hStdOutput = startupInfo.hStdError =
892                INVALID_HANDLE_VALUE;
893
894            //Create handles for output redirection
895            Handle<HANDLE> readPipe;
896            HANDLE writePipe;
897            SECURITY_ATTRIBUTES security;
898            ZeroMemory(&security, sizeof(security));
899            security.nLength = sizeof(security);
900            security.lpSecurityDescriptor = NULL;
901            security.bInheritHandle = true;
902
903            if (CreatePipe(&static_cast<HANDLE&>(readPipe), &writePipe, &security, 0))
904            {
905                startupInfo.dwFlags |= STARTF_USESTDHANDLES;
906                startupInfo.hStdOutput = startupInfo.hStdError =
907                    writePipe;
908            }
909
910            PROCESS_INFORMATION processInfo;
911            ZeroMemory(&processInfo, sizeof(processInfo));
912            std::vector<wchar_t> buffer(eraserPath.length() + finalParameters.length() + 4);
913            wcscpy_s(&buffer.front(), buffer.size(), (L"\"" + eraserPath + L"\" " +
914                finalParameters).c_str());
915
916            if (!CreateProcess(NULL, &buffer.front(), NULL, NULL, true, CREATE_NO_WINDOW,
917                NULL, NULL, &startupInfo, &processInfo))
918            {
919                //Why did we fail? Is it because we have too many files
920                if (GetLastError() == ERROR_FILENAME_EXCED_RANGE)
921                    throw FormatString(LoadString(IDS_ERROR_TOO_MANY_FILES));
922               
923                //Or if elevation is required for this operation
924                else if (GetLastError() == ERROR_ELEVATION_REQUIRED)
925                    return RunEraser(action, parameters, true, parent, show);
926
927                //Or otherwise?
928                else
929                    throw FormatString(LoadString(IDS_ERROR_MISC), FormatError().c_str());
930            }
931
932            //Wait for the process to finish.
933            Handle<HANDLE> hProcess(processInfo.hProcess),
934                           hThread(processInfo.hThread);
935            CloseHandle(writePipe);
936            WaitForSingleObject(hProcess, static_cast<DWORD>(-1));
937            DWORD exitCode = 0;
938           
939            if (GetExitCodeProcess(processInfo.hProcess, &exitCode))
940                if (exitCode == ERROR_ACCESS_DENIED)
941                {
942                    //The spawned instance could not connect with the master instance
943                    //because it is running as an administrator. Spawn the new instance
944                    //again, this time as an administrator
945                    RunEraser(action, parameters, true, parent, show);
946                }
947                else if (exitCode)
948                {
949                    char buffer[8192];
950                    DWORD lastRead = 0;
951                    std::wstring output;
952
953                    while (ReadFile(readPipe, buffer, sizeof(buffer), &lastRead, NULL) && lastRead != 0)
954                    {
955                        size_t lastConvert = 0;
956                        wchar_t wBuffer[8192];
957                        if (!mbstowcs_s(&lastConvert, wBuffer, sizeof(wBuffer) / sizeof(wBuffer[0]),
958                            buffer, lastRead))
959                        {
960                            output += std::wstring(wBuffer, lastConvert);
961                        }
962                    }
963
964                    //Show the error message.
965                    throw output;
966                }
967        }
968    }
969
970    void CCtxMenu::InsertSeparator(HMENU menu)
971    {
972        MENUITEMINFO mii;
973        mii.cbSize = sizeof(MENUITEMINFO);
974        mii.fMask = MIIM_TYPE;
975        mii.fType = MF_SEPARATOR;
976        InsertMenuItem(menu, 0, false, &mii);
977    }
978
979    HICON CCtxMenu::GetMenuIcon()
980    {
981        int smIconSize = GetSystemMetrics(SM_CXMENUCHECK);
982        return static_cast<HICON>(LoadImage(theApp.m_hInstance, L"Eraser",
983            IMAGE_ICON, smIconSize, smIconSize, LR_DEFAULTCOLOR));
984    }
985
986    HBITMAP CCtxMenu::GetMenuBitmapFromIcon(HICON icon)
987    {
988        BITMAP bitmap;
989        ICONINFO iconInfo;
990        ZeroMemory(&bitmap, sizeof(bitmap));
991        ZeroMemory(&iconInfo, sizeof(iconInfo));
992
993        //Try to get the icon's size, bitmap and bit depth. We will try to convert
994        //the bitmap into a DIB for display on Vista if it contains an alpha channel.
995        if (!GetIconInfo(icon, &iconInfo))
996            return NULL;
997
998        Handle<HBITMAP> iconMask(iconInfo.hbmMask);
999        if (!GetObject(iconInfo.hbmColor, sizeof(BITMAP), &bitmap) ||
1000            bitmap.bmBitsPixel < 32)
1001            return iconInfo.hbmColor;
1002
1003        //Draw the icon onto the DIB which will preseve its alpha values
1004        Handle<HDC> hdcDest = CreateCompatibleDC(NULL);
1005        HBITMAP dib = CreateDIB(bitmap.bmWidth, bitmap.bmHeight, NULL);
1006        SelectObject(hdcDest, dib);
1007
1008        Handle<HBITMAP> iconBitmap(iconInfo.hbmColor);
1009        DrawIconEx(hdcDest, 0, 0, icon, bitmap.bmWidth, bitmap.bmHeight, 0, NULL, DI_NORMAL);
1010        return dib;
1011    }
1012
1013    HBITMAP CCtxMenu::CreateDIB(LONG width, LONG height, char** bitmapBits)
1014    {
1015        BITMAPINFO info;
1016        ZeroMemory(&info, sizeof(info));
1017        info.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
1018        info.bmiHeader.biWidth = width;
1019        info.bmiHeader.biHeight = height;
1020        info.bmiHeader.biPlanes = 1;
1021        info.bmiHeader.biBitCount = 32;
1022
1023        Handle<HDC> screenDC(GetDC(NULL));
1024        return ::CreateDIBSection(screenDC, &info, DIB_RGB_COLORS,
1025            reinterpret_cast<void**>(bitmapBits), NULL, 0);
1026    }
1027}
Note: See TracBrowser for help on using the repository browser.