source: trunk/eraser/Eraser.Shell/CtxMenu.cpp @ 2060

Revision 2060, 28.9 KB checked in by lowjoel, 5 years ago (diff)

Removed debug code, oops.

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