source: branches/eraser6/ShellExt/CtxMenu.cpp @ 789

Revision 789, 23.5 KB checked in by lowjoel, 6 years ago (diff)

-Replace all functions taking a string parameter to take a const reference to reduce copying
-Factor out the call to launch Eraser, allowing for the spawned instance to run elevated (though it is currently unimplemented)

  • Property svn:eol-style set to native
  • Property svn:keywords set to Id Rev
Line 
1/*
2 * $Id$
3 * Copyright 2008 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<HANDLE>::~Handle()
91{
92    CloseHandle(Object);
93}
94
95Handle<HKEY>::~Handle()
96{
97    CloseHandle(Object);
98}
99
100namespace Eraser {
101    HRESULT CCtxMenu::FinalConstruct()
102    {
103        //Initialise member variables.
104        MenuID = 0;
105        std::wstring menuTitle(LoadString(IDS_ERASER));
106        MenuTitle = new wchar_t[menuTitle.length() + 1];
107        wcscpy_s(MenuTitle, menuTitle.length() + 1, menuTitle.c_str());
108
109        //Check if the shell extension has been disabled.
110        Handle<HKEY> eraserKey;
111        LONG openKeyResult = RegOpenKeyEx(HKEY_CURRENT_USER,
112            L"Software\\Eraser\\Eraser 6\\3460478d-ed1b-4ecc-96c9-2ca0e8500557\\", 0,
113            KEY_READ, &static_cast<HKEY&>(eraserKey));
114
115        switch (openKeyResult)
116        {
117        case ERROR_FILE_NOT_FOUND:
118            //No settings defined: we default to enabling the shell extension.
119            return S_OK;
120
121        case ERROR_SUCCESS:
122            break;
123           
124        default:
125            return E_FAIL;
126        }
127
128        //Check the value of the IntegrateWithShell value.
129        DWORD valueType;       
130        BYTE valueBuffer[512];
131        DWORD valueBufferSize = sizeof(valueBuffer);
132        DWORD error = RegQueryValueEx(eraserKey, L"IntegrateWithShell", NULL, &valueType,
133            valueBuffer, &valueBufferSize);
134        if (error == ERROR_SUCCESS)
135        {
136            if (!*reinterpret_cast<char*>(valueBuffer + 51))
137                return E_FAIL;
138        }
139
140        return S_OK;
141    }
142
143    HRESULT CCtxMenu::FinalRelease()
144    {
145        delete[] MenuTitle;
146        return S_OK;
147    }
148
149    HRESULT CCtxMenu::Initialize(LPCITEMIDLIST pidlFolder, LPDATAOBJECT pDataObj,
150                                 HKEY hProgID)
151    {
152        //Determine where the shell extension was invoked from.
153        if (GetHKeyPath(hProgID) == L"{645FF040-5081-101B-9F08-00AA002F954E}")
154        {
155            InvokeReason = INVOKEREASON_RECYCLEBIN;
156
157            //We can't do much other processing: the LPDATAOBJECT parameter contains
158            //data that is a private type so we don't know how to query for it.
159            return S_OK;
160        }
161
162        //Check pidlFolder for the drop path, if it exists. This is for drag-and-drop
163        //context menus.
164        else if (pidlFolder != NULL)
165        {
166            InvokeReason = INVOKEREASON_DRAGDROP;
167
168            //Translate the drop path to a location on the filesystem.
169            wchar_t dropTargetPath[MAX_PATH];
170            if (!SHGetPathFromIDList(pidlFolder, dropTargetPath))
171                return E_FAIL;
172
173            DragDropDestinationDirectory = dropTargetPath;
174        }
175
176        //Okay, everything else is a simple context menu for a set of selected files/
177        //folders/drives.
178        else
179            InvokeReason = INVOKEREASON_FILEFOLDER;
180
181        //Look for CF_HDROP data in the data object.
182        FORMATETC fmt = { CF_HDROP, NULL, DVASPECT_CONTENT, -1, TYMED_HGLOBAL };
183        STGMEDIUM stg = { TYMED_HGLOBAL };
184        if (FAILED(pDataObj->GetData(&fmt, &stg)))
185            //Nope! Return an "invalid argument" error back to Explorer.
186            return E_INVALIDARG;
187
188        //Get a pointer to the actual data.
189        HDROP hDrop = static_cast<HDROP>(GlobalLock(stg.hGlobal));
190        if (hDrop == NULL)
191            return E_INVALIDARG;
192
193        //Sanity check - make sure there is at least one filename.
194        UINT uNumFiles = DragQueryFile(hDrop, 0xFFFFFFFF, NULL, 0);
195        if (!uNumFiles)
196        {
197            GlobalUnlock(stg.hGlobal);
198            ReleaseStgMedium(&stg);
199            return E_INVALIDARG;
200        }
201
202        //Collect all the files which have been selected.
203        HRESULT hr = S_OK;
204        WCHAR buffer[MAX_PATH] = {0};
205        for (UINT i = 0; i < uNumFiles; i++)
206        {
207            UINT charsWritten = DragQueryFile(hDrop, i, buffer, sizeof(buffer) / sizeof(buffer[0]));
208            if (!charsWritten)
209            {
210                hr = E_INVALIDARG;
211                continue;
212            }
213
214            SelectedFiles.push_back(std::wstring(buffer, charsWritten));
215        }
216
217        //Clean up.
218        GlobalUnlock(stg.hGlobal);
219        ReleaseStgMedium(&stg);
220        return hr;
221    }
222
223    HRESULT CCtxMenu::QueryContextMenu(HMENU hmenu, UINT uMenuIndex, UINT uidFirstCmd,
224                                       UINT /*uidLastCmd*/, UINT uFlags)
225    {
226        //If the flags include CMF_DEFAULTONLY then we shouldn't do anything.
227        if (uFlags & CMF_DEFAULTONLY)
228            return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, 0);
229
230        //First, create and populate a submenu.
231        UINT uID = uidFirstCmd;
232        HMENU hSubmenu = CreatePopupMenu();
233
234        //Create the submenu, following the order defined in the CEraserLPVERB enum, creating
235        //only items which are applicable.
236        Actions applicableActions = GetApplicableActions();
237        VerbMenuIndices.clear();
238        if (applicableActions & ACTION_ERASE)
239        {
240            InsertMenu(hSubmenu, ACTION_ERASE, MF_BYPOSITION, uID++,
241                LoadString(IDS_ACTION_ERASE).c_str());              //Erase
242            VerbMenuIndices.push_back(ACTION_ERASE);
243        }
244        if (applicableActions & ACTION_ERASE_ON_RESTART)
245        {
246            InsertMenu(hSubmenu, ACTION_ERASE_ON_RESTART, MF_BYPOSITION, uID++,
247                LoadString(IDS_ACTION_ERASERESTART).c_str());       //Erase on Restart
248            VerbMenuIndices.push_back(ACTION_ERASE_ON_RESTART);
249        }
250        if (applicableActions & ACTION_ERASE_UNUSED_SPACE)
251        {
252            InsertMenu(hSubmenu, ACTION_ERASE_UNUSED_SPACE, MF_BYPOSITION, uID++,
253                LoadString(IDS_ACTION_ERASEUNUSEDSPACE).c_str());   //Erase Unused Space
254            VerbMenuIndices.push_back(ACTION_ERASE_UNUSED_SPACE);
255        }
256        //-------------------------------------------------------------------------
257        if (applicableActions & ACTION_SECURE_MOVE)
258        {
259            if (uID - uidFirstCmd > 0)
260            {
261                std::auto_ptr<MENUITEMINFO> separator(GetSeparator());
262                InsertMenuItem(hSubmenu, 0, FALSE, separator.get());
263            }
264
265            InsertMenu(hSubmenu, ACTION_SECURE_MOVE, MF_BYPOSITION, uID++,
266                LoadString(IDS_ACTION_SECUREMOVE).c_str());         //Secure Move
267            VerbMenuIndices.push_back(ACTION_SECURE_MOVE);
268        }
269
270        //Insert the submenu into the Context menu provided by Explorer.
271        {
272            MENUITEMINFO mii = { sizeof(MENUITEMINFO) };
273            mii.wID = uID++;
274            mii.fMask = MIIM_SUBMENU | MIIM_STRING | MIIM_ID;
275            mii.hSubMenu = hSubmenu;
276            mii.dwTypeData = const_cast<wchar_t*>(MenuTitle);
277
278            //Set the bitmap for the registered item. Vista machines will be set using a DIB,
279            //older machines will be ownerdrawn.
280            OSVERSIONINFO osvi;
281            ZeroMemory(&osvi, sizeof(osvi));
282            osvi.dwOSVersionInfoSize = sizeof(osvi);
283
284            if (GetVersionEx(&osvi) && osvi.dwPlatformId == VER_PLATFORM_WIN32_NT &&
285                osvi.dwMajorVersion >= 6)
286            {
287                mii.fMask |= MIIM_CHECKMARKS;
288                mii.hbmpUnchecked = GetMenuBitmap();
289            }
290            else if (InvokeReason != INVOKEREASON_DRAGDROP)
291            {
292                mii.fMask |= MIIM_FTYPE;
293                mii.fType = MFT_OWNERDRAW;
294            }
295
296            MenuID = uMenuIndex++;
297            InsertMenuItem(hmenu, MenuID, TRUE, &mii);
298        }
299
300        return MAKE_HRESULT(SEVERITY_SUCCESS, FACILITY_NULL, uID - uidFirstCmd);
301    }
302
303    HRESULT CCtxMenu::HandleMenuMsg(UINT uMsg, WPARAM wParam, LPARAM lParam)
304    {
305        return HandleMenuMsg2(uMsg, wParam, lParam, NULL);
306    }
307
308    HRESULT CCtxMenu::HandleMenuMsg2(UINT uMsg, WPARAM /*wParam*/, LPARAM lParam,
309                                     LRESULT* result)
310    {
311        //Skip if we aren't handling our own.
312        bool handleResult = false;
313        switch (uMsg)
314        {
315        case WM_MEASUREITEM:
316            {
317                MEASUREITEMSTRUCT* mis = reinterpret_cast<MEASUREITEMSTRUCT*>(lParam);
318                if (mis->CtlID == MenuID)
319                    handleResult = OnMeasureItem(mis->itemWidth, mis->itemHeight);
320                else
321                    handleResult = false;
322                break;
323            }
324
325        case WM_DRAWITEM:
326            {
327                DRAWITEMSTRUCT* dis = reinterpret_cast<DRAWITEMSTRUCT*>(lParam);
328                if (dis->CtlID == MenuID)
329                    handleResult = OnDrawItem(dis->hDC, dis->rcItem, dis->itemAction, dis->itemState);
330                else
331                    handleResult = false;
332            }
333        }
334
335        if (result)
336            *result = handleResult;
337        return S_OK;
338    }
339
340    bool CCtxMenu::OnMeasureItem(UINT& itemWidth, UINT& itemHeight)
341    {
342        LOGFONT logFont;
343        if (!SystemParametersInfo(SPI_GETICONTITLELOGFONT, sizeof(logFont), &logFont, 0))
344            return false;
345
346        //Measure the size of the text.
347        Handle<HDC> screenDC = GetDC(NULL);
348        Handle<HFONT> font = CreateFontIndirect(&logFont);
349        SelectObject(screenDC, font);
350        SIZE textSize;
351        if (!GetTextExtentPoint32(screenDC, MenuTitle, static_cast<DWORD>(wcslen(MenuTitle)), &textSize))
352            return false;
353
354        itemWidth = textSize.cx;
355        itemHeight = textSize.cy;
356
357        //Account for the size of the bitmap.
358        UINT iconWidth = GetSystemMetrics(SM_CXMENUCHECK);
359        itemWidth += iconWidth;
360        itemHeight = std::max(iconWidth, itemHeight);
361
362        //And remember the minimum size for menu items.
363        itemHeight = std::max((int)itemHeight, GetSystemMetrics(SM_CXMENUSIZE));
364        return true;
365    }
366
367    bool CCtxMenu::OnDrawItem(HDC hdc, RECT rect, UINT /*action*/, UINT state)
368    {
369        //Draw the background.
370        LOGBRUSH logBrush = { BS_SOLID,
371            (state & ODS_SELECTED) ?
372                GetSysColor(COLOR_HIGHLIGHT) : GetSysColor(COLOR_MENU),
373            0
374        };
375        Handle<HBRUSH> bgBrush = CreateBrushIndirect(&logBrush);
376        FillRect(hdc, &rect, bgBrush);
377
378        //Then the bitmap.
379        {
380            //Draw the icon with alpha and all first.
381            Handle<HICON> icon(GetMenuIcon());
382            int iconSize = GetSystemMetrics(SM_CXMENUCHECK);
383            int iconMargin = GetSystemMetrics(SM_CXEDGE);
384            DrawIconEx(hdc, rect.left + iconMargin, rect.top + (rect.bottom - rect.top - iconSize) / 2,
385                icon, 0, 0, 0, bgBrush, DI_NORMAL);
386
387            //Move the rectangle's left bound to the text starting position
388            rect.left += iconMargin * 2 + iconSize;
389        }
390       
391        //Draw the text.
392        SetBkMode(hdc, TRANSPARENT);
393        LOGFONT logFont;
394        if (!SystemParametersInfo(SPI_GETICONTITLELOGFONT, sizeof(logFont), &logFont, 0))
395            return false;
396
397        SIZE textSize;
398        if (!GetTextExtentPoint32(hdc, MenuTitle, static_cast<DWORD>(wcslen(MenuTitle)), &textSize))
399            return false;
400
401        COLORREF oldColour = SetTextColor(hdc, (state & ODS_SELECTED) ?
402            GetSysColor(COLOR_HIGHLIGHTTEXT) : GetSysColor(COLOR_MENUTEXT));
403        UINT flags = DST_PREFIXTEXT;
404        if (state & ODS_NOACCEL)
405            flags |= DSS_HIDEPREFIX;
406        ::DrawState(hdc, NULL, NULL, reinterpret_cast<LPARAM>(MenuTitle), wcslen(MenuTitle),
407            rect.left, rect.top + (rect.bottom - rect.top - textSize.cy) / 2, textSize.cx, textSize.cy, flags);
408        SetTextColor(hdc, oldColour);
409        return true;
410    }
411
412    HRESULT CCtxMenu::GetCommandString(UINT_PTR idCmd, UINT uFlags, UINT* /*pwReserved*/,
413                                       LPSTR pszName, UINT cchMax)
414    {
415        //We only know how to handle help string requests.
416        if (!(uFlags & GCS_HELPTEXT))
417            return E_INVALIDARG;
418
419        //Get the command string for the given id
420        if (idCmd >= VerbMenuIndices.size())
421            return E_INVALIDARG;
422
423        std::wstring commandString;
424        switch (VerbMenuIndices[idCmd])
425        {
426        case ACTION_ERASE:
427            commandString = LoadString(IDS_HELPSTRING_ERASE);
428            break;
429        case ACTION_ERASE_ON_RESTART:
430            commandString = LoadString(IDS_HELPSTRING_ERASEONRESTART);
431            break;
432        case ACTION_ERASE_UNUSED_SPACE:
433            commandString = LoadString(IDS_HELPSTRING_ERASEUNUSEDSPACE);
434            break;
435        case ACTION_SECURE_MOVE:
436            commandString = LoadString(IDS_HELPSTRING_SECUREMOVE);
437            break;
438        default:
439            //We don't know what action this is: return E_INVALIDARG.
440            return E_INVALIDARG;
441        }
442
443        //Return the help string to Explorer.
444        if (uFlags & GCS_UNICODE)
445            wcscpy_s(reinterpret_cast<wchar_t*>(pszName), cchMax, commandString.c_str());
446        else
447        {
448            size_t convCount = 0;
449            wcstombs_s(&convCount, pszName, cchMax, commandString.c_str(), commandString.length());
450        }
451
452        return S_OK;
453    }
454
455    HRESULT CCtxMenu::InvokeCommand(LPCMINVOKECOMMANDINFO pCmdInfo)
456    {
457        //If lpVerb really points to a string, ignore this function call and bail out.
458        if (HIWORD(pCmdInfo->lpVerb) != 0)
459            return E_INVALIDARG;
460
461        //If the verb index refers to an item outside the bounds of our VerbMenuIndices
462        //vector, exit.
463        if (LOWORD(pCmdInfo->lpVerb) >= VerbMenuIndices.size())
464            return E_INVALIDARG;
465
466        //Build the command line
467        std::wstring commandAction;
468        std::wstring commandLine;
469        bool commandElevate = false;
470        HRESULT result = E_INVALIDARG;
471        switch (VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)])
472        {
473        case ACTION_ERASE_ON_RESTART:
474            commandLine += L"-s=restart ";
475
476        case ACTION_ERASE:
477            {
478                //Add Task command.
479                commandAction = L"addtask";
480
481                //See the invocation context: if it is executed from the recycle bin
482                //then the list of selected files will be empty.
483                if (InvokeReason == INVOKEREASON_RECYCLEBIN)
484                {
485                    commandLine += L"-r";
486                }
487                else
488                {
489                    //Okay, we were called from an item right-click; iterate over every
490                    //selected file
491                    for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
492                        i != SelectedFiles.end(); ++i)
493                    {
494                        //Check if the current item is a file or folder.
495                        std::wstring item(*i);
496                        if (item.length() > 3 && item[item.length() - 1] == '\\')
497                            item.erase(item.end() - 1);
498                        DWORD attributes = GetFileAttributes(item.c_str());
499
500                        //Add the correct command line for the file type.
501                        if (attributes & FILE_ATTRIBUTE_DIRECTORY)
502                            commandLine += L"\"-d=" + item + L"\" ";
503                        else
504                            commandLine += L"\"" + item + L"\" ";
505                    }
506                }
507
508                break;
509            }
510
511        case ACTION_ERASE_UNUSED_SPACE:
512            {
513                //We want to add a new task
514                commandAction = L"addtask";
515
516                //Erasing unused space requires elevation
517                commandElevate = true;
518
519                //Add every item onto the command line
520                for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
521                    i != SelectedFiles.end(); ++i)
522                {
523                    commandLine += L"\"-u=" + *i + L",clusterTips\" ";
524                }
525               
526                break;
527            }
528
529        default:
530            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
531            {
532                MessageBox(pCmdInfo->hwnd, FormatString(LoadString(IDS_ERROR_UNKNOWNACTION),
533                    VerbMenuIndices[LOWORD(pCmdInfo->lpVerb)]).c_str(),
534                    LoadString(IDS_ERASERSHELLEXT).c_str(), MB_OK | MB_ICONERROR);
535            }
536        }
537
538        try
539        {
540            RunEraser(commandAction, commandLine, commandElevate, pCmdInfo->hwnd,
541                pCmdInfo->nShow);
542        }
543        catch (const std::wstring& e)
544        {
545            if (!(pCmdInfo->fMask & CMIC_MASK_FLAG_NO_UI))
546            {
547                MessageBox(pCmdInfo->hwnd, e.c_str(), LoadString(IDS_ERASERSHELLEXT).c_str(),
548                    MB_OK | MB_ICONERROR);
549            }
550        }
551
552        return result;
553    }
554
555    CCtxMenu::Actions CCtxMenu::GetApplicableActions()
556    {
557        unsigned result = 0;
558       
559        //First decide the actions which are applicable to the current invocation
560        //reason.
561        switch (InvokeReason)
562        {
563        case INVOKEREASON_RECYCLEBIN:
564            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART;
565            break;
566        case INVOKEREASON_FILEFOLDER:
567            result |= ACTION_ERASE | ACTION_ERASE_ON_RESTART | ACTION_ERASE_UNUSED_SPACE;
568        case INVOKEREASON_DRAGDROP:
569            result |= ACTION_SECURE_MOVE;
570        }
571
572        //Remove actions that don't apply to the current invocation reason.
573        for (std::list<std::wstring>::const_iterator i = SelectedFiles.begin();
574            i != SelectedFiles.end(); ++i)
575        {
576            //Remove trailing slashes if they are directories.
577            std::wstring item(*i);
578
579            //Check if the path is a path to a volume, if it is not, remove the
580            //erase unused space verb.
581            wchar_t volumeName[MAX_PATH];
582            if (!GetVolumeNameForVolumeMountPoint(item.c_str(), volumeName,
583                sizeof(volumeName) / sizeof(volumeName[0])))
584            {
585                result &= ~ACTION_ERASE_UNUSED_SPACE;
586            }
587        }
588
589        return static_cast<Actions>(result);
590    }
591
592    std::wstring CCtxMenu::LoadString(UINT stringID)
593    {
594        //Get a pointer to the buffer containing the string (we're copying it anyway)
595        wchar_t* buffer = NULL;
596        DWORD lastCount = ::LoadString(theApp.m_hInstance, stringID,
597            reinterpret_cast<wchar_t*>(&buffer), 0);
598
599        if (lastCount > 0)
600            return std::wstring(buffer, lastCount);
601        return std::wstring();
602    }
603
604    std::wstring CCtxMenu::FormatString(const std::wstring& formatString, ...)
605    {
606        std::vector<wchar_t> formatStr(formatString.length() + 1);
607        wcscpy_s(&formatStr.front(), formatStr.size(), formatString.c_str());
608
609        std::vector<wchar_t> buffer(formatStr.size());
610        for ( ; ; )
611        {
612            buffer.resize(buffer.size() * 2);
613            va_list arguments;
614            va_start(arguments, formatString);
615            int result = vswprintf_s(&buffer.front(), buffer.size(), &formatStr.front(),
616                arguments);
617            va_end(arguments);
618
619            if (result > 0 && static_cast<unsigned>(result) < buffer.size())
620            {
621                break;
622            }
623        }
624
625        //Return the result as a wstring
626        std::wstring result;
627        if (buffer.size() > 0)
628            result = &buffer.front();
629        return result;
630    }
631
632    std::wstring CCtxMenu::GetHKeyPath(HKEY handle)
633    {
634        if (NtQueryKey == NULL)
635            NtQueryKey = reinterpret_cast<pNtQueryKey>(GetProcAddress(
636                LoadLibrary(L"Ntdll.dll"), "NtQueryKey"));
637
638        //Keep querying for the key information until enough buffer space has been allocated.
639        std::vector<char> buffer(sizeof(KEY_NODE_INFORMATION));
640        NTSTATUS queryResult = STATUS_BUFFER_TOO_SMALL;
641        ULONG keyInfoSize = 0;
642
643        while (queryResult == STATUS_BUFFER_TOO_SMALL || queryResult == STATUS_BUFFER_OVERFLOW)
644        {
645            buffer.resize(buffer.size() + keyInfoSize);
646            ZeroMemory(&buffer.front(), buffer.size());
647            queryResult = NtQueryKey(handle, KeyNodeInformation, &buffer.front(),
648                static_cast<ULONG>(buffer.size()), &keyInfoSize);
649        }
650
651        if (queryResult != STATUS_SUCCESS)
652            return std::wstring();
653
654        KEY_NODE_INFORMATION* keyInfo = reinterpret_cast<KEY_NODE_INFORMATION*>(
655            &buffer.front());
656        return keyInfo->Name;
657    }
658
659    void CCtxMenu::RunEraser(const std::wstring& action, const std::wstring& parameters,
660        bool elevated, HWND parent, int show)
661    {
662        //Get the path to this DLL so we can look for Eraser.exe
663        wchar_t fileName[MAX_PATH];
664        DWORD fileNameLength = GetModuleFileName(theApp.m_hInstance, fileName,
665            sizeof(fileName) / sizeof(fileName[0]));
666        if (!fileNameLength || fileNameLength >= sizeof(fileName) / sizeof(fileName[0]))
667            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
668       
669        //Trim to the last \, then append Eraser.exe
670        std::wstring eraserPath(fileName, fileNameLength);
671        std::wstring::size_type lastBackslash = eraserPath.rfind('\\');
672        if (lastBackslash == std::wstring::npos)
673            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
674
675        eraserPath.erase(eraserPath.begin() + lastBackslash + 1, eraserPath.end());
676        if (eraserPath.empty())
677            throw LoadString(IDS_ERROR_CANNOTFINDERASER);
678
679        eraserPath += L"Eraser.exe";
680
681        std::wostringstream finalCmdLine;
682        finalCmdLine << L"\"" << eraserPath << L"\" \"" << action << L"\" -q "
683                     << parameters;
684        std::wstring cmdLine(finalCmdLine.str());
685        std::vector<wchar_t> buffer(cmdLine.length() + 1);
686        wcscpy_s(&buffer.front(), cmdLine.length() + 1, cmdLine.c_str());
687
688#if 0
689        //If the process must be elevated we use ShellExecute with the runas verb
690        //to elevate the new process.
691        if (elevated)
692        {
693            int result = reinterpret_cast<int>(ShellExecute(parent, L"runas",
694                eraserPath.c_str(), &buffer.front(), NULL, show));
695            if (result <= 32)
696                throw LoadString(IDS_ERROR_UNKNOWN);
697        }
698
699        //If the process isn't to be elevated, we use CreateProcess so we can get
700        //read the output from the child process
701        else
702#endif
703        {
704            //Create the process.
705            STARTUPINFO startupInfo;
706            ZeroMemory(&startupInfo, sizeof(startupInfo));
707            startupInfo.cb = sizeof(startupInfo);
708            startupInfo.dwFlags = STARTF_USESHOWWINDOW;
709            startupInfo.wShowWindow = static_cast<WORD>(show);
710            startupInfo.hStdInput = startupInfo.hStdOutput = startupInfo.hStdError =
711                INVALID_HANDLE_VALUE;
712
713            //Create handles for output redirection
714            Handle<HANDLE> readPipe;
715            HANDLE writePipe;
716            SECURITY_ATTRIBUTES security;
717            ZeroMemory(&security, sizeof(security));
718            security.nLength = sizeof(security);
719            security.lpSecurityDescriptor = NULL;
720            security.bInheritHandle = true;
721
722            if (CreatePipe(&static_cast<HANDLE&>(readPipe), &writePipe, &security, 0))
723            {
724                startupInfo.dwFlags |= STARTF_USESTDHANDLES;
725                startupInfo.hStdOutput = startupInfo.hStdError =
726                    writePipe;
727            }
728
729            PROCESS_INFORMATION processInfo;
730            ZeroMemory(&processInfo, sizeof(processInfo));
731            if (!CreateProcess(NULL, &buffer.front(), NULL, NULL, true, CREATE_NO_WINDOW,
732                NULL, NULL, &startupInfo, &processInfo))
733            {
734                throw LoadString(IDS_ERROR_CANNOTFINDERASER);
735            }
736
737            //Wait for the process to finish.
738            Handle<HANDLE> hProcess(processInfo.hProcess),
739                           hThread(processInfo.hThread);
740            CloseHandle(writePipe);
741            WaitForSingleObject(hProcess, static_cast<DWORD>(-1));
742            DWORD exitCode = 0;
743           
744            if (GetExitCodeProcess(processInfo.hProcess, &exitCode) && exitCode)
745            {
746                char buffer[8192];
747                DWORD lastRead = 0;
748                std::wstring output;
749
750                while (ReadFile(readPipe, buffer, sizeof(buffer), &lastRead, NULL) && lastRead != 0)
751                {
752                    size_t lastConvert = 0;
753                    wchar_t wBuffer[8192];
754                    if (!mbstowcs_s(&lastConvert, wBuffer, sizeof(wBuffer) / sizeof(wBuffer[0]),
755                        buffer, lastRead))
756                    {
757                        output += std::wstring(wBuffer, lastConvert);
758                    }
759                }
760
761                //Show the error message.
762                throw output;
763            }
764        }
765    }
766
767    MENUITEMINFO* CCtxMenu::GetSeparator()
768    {
769        MENUITEMINFO* mii = new MENUITEMINFO();
770        mii->cbSize = sizeof(MENUITEMINFO);
771        mii->fMask = MIIM_TYPE;
772        mii->fType = MF_SEPARATOR;
773        return mii;
774    }
775
776    HICON CCtxMenu::GetMenuIcon()
777    {
778        int smIconSize = GetSystemMetrics(SM_CXMENUCHECK);
779        return static_cast<HICON>(LoadImage(theApp.m_hInstance, L"Eraser",
780            IMAGE_ICON, smIconSize, smIconSize, LR_LOADTRANSPARENT));
781    }
782
783    HBITMAP CCtxMenu::GetMenuBitmap()
784    {
785        BITMAP bitmap;
786        ICONINFO iconInfo;
787        ZeroMemory(&bitmap, sizeof(bitmap));
788        ZeroMemory(&iconInfo, sizeof(iconInfo));
789        Handle<HICON> icon(GetMenuIcon());
790
791        //Try to get the icon's size, bitmap and bit depth. We will try to convert
792        //the bitmap into a DIB for display on Vista if it contains an alpha channel.
793        if (!GetIconInfo(icon, &iconInfo))
794            return NULL;
795
796        Handle<HBITMAP> iconMask(iconInfo.hbmMask);
797        if (!GetObject(iconInfo.hbmColor, sizeof(BITMAP), &bitmap) ||
798            bitmap.bmBitsPixel < 32)
799            return iconInfo.hbmColor;
800
801        //Try converting the DDB into a DIB.
802        DIBSECTION dibSection;
803        HBITMAP dib = CreateDIB(bitmap.bmWidth, bitmap.bmHeight, NULL);
804        if (!GetObject(dib, sizeof(dibSection), &dibSection) ||
805            !GetDIBits(CreateCompatibleDC(NULL), iconInfo.hbmColor, 0, bitmap.bmHeight,
806                dibSection.dsBm.bmBits, reinterpret_cast<BITMAPINFO*>(&dibSection.dsBmih),
807                DIB_RGB_COLORS))
808        {
809            return iconInfo.hbmColor;
810        }
811
812        return dib;
813    }
814
815    HBITMAP CCtxMenu::CreateDIB(LONG width, LONG height, char** bitmapBits)
816    {
817        BITMAPINFO info;
818        ZeroMemory(&info, sizeof(info));
819        info.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
820        info.bmiHeader.biWidth = width;
821        info.bmiHeader.biHeight = height;
822        info.bmiHeader.biPlanes = 1;
823        info.bmiHeader.biBitCount = 32;
824        info.bmiHeader.biCompression = BI_RGB;
825        char* dibBitmapBits = NULL;
826        return ::CreateDIBSection(0, &info, DIB_RGB_COLORS,
827            reinterpret_cast<void**>(bitmapBits ? bitmapBits : &dibBitmapBits), NULL, 0);
828    }
829}
Note: See TracBrowser for help on using the repository browser.