diff --git a/Scylla/OptionsGui.cpp b/Scylla/OptionsGui.cpp index 8bab408..fb035ca 100644 --- a/Scylla/OptionsGui.cpp +++ b/Scylla/OptionsGui.cpp @@ -1,89 +1,89 @@ #include "OptionsGui.h" #include "ConfigurationHolder.h" BOOL OptionsGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) { loadOptions(); DoDataExchange(DDX_LOAD); // show settings EditSectionName.LimitText(IMAGE_SIZEOF_SHORT_NAME); CenterWindow(); return TRUE; } void OptionsGui::OnOK(UINT uNotifyCode, int nID, CWindow wndCtl) { DoDataExchange(DDX_SAVE); saveOptions(); ConfigurationHolder::saveConfiguration(); EndDialog(0); } void OptionsGui::OnCancel(UINT uNotifyCode, int nID, CWindow wndCtl) { EndDialog(0); } void OptionsGui::saveOptions() { std::map::iterator mapIter; for (mapIter = ConfigurationHolder::getConfigList().begin() ; mapIter != ConfigurationHolder::getConfigList().end(); mapIter++) { switch(mapIter->first) { case USE_PE_HEADER_FROM_DISK: usePEHeaderFromDisk ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case DEBUG_PRIVILEGE: debugPrivilege ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case CREATE_BACKUP: createBackup ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case DLL_INJECTION_AUTO_UNLOAD: dllInjectionAutoUnload ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case UPDATE_HEADER_CHECKSUM: updateHeaderChecksum ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case IAT_SECTION_NAME: - wcscpy(mapIter->second.valueString, iatSectionName); + wcscpy_s(mapIter->second.valueString, _countof(mapIter->second.valueString), iatSectionName); break; } } } void OptionsGui::loadOptions() { std::map::iterator mapIter; for (mapIter = ConfigurationHolder::getConfigList().begin() ; mapIter != ConfigurationHolder::getConfigList().end(); mapIter++) { switch(mapIter->first) { case USE_PE_HEADER_FROM_DISK: usePEHeaderFromDisk = mapIter->second.isTrue(); break; case DEBUG_PRIVILEGE: debugPrivilege = mapIter->second.isTrue(); break; case CREATE_BACKUP: createBackup = mapIter->second.isTrue(); break; case DLL_INJECTION_AUTO_UNLOAD: dllInjectionAutoUnload = mapIter->second.isTrue(); break; case UPDATE_HEADER_CHECKSUM: updateHeaderChecksum = mapIter->second.isTrue(); break; case IAT_SECTION_NAME: - wcsncpy(iatSectionName, mapIter->second.valueString, _countof(iatSectionName)-1); + wcsncpy_s(iatSectionName, _countof(iatSectionName), mapIter->second.valueString, _countof(iatSectionName)-1); iatSectionName[_countof(iatSectionName)-1] = L'\0'; break; } } } diff --git a/Scylla/PickDllGui.cpp b/Scylla/PickDllGui.cpp index 9942782..acb1b21 100644 --- a/Scylla/PickDllGui.cpp +++ b/Scylla/PickDllGui.cpp @@ -1,145 +1,145 @@ #include "PickDllGui.h" PickDllGui::PickDllGui(std::vector &moduleList) : moduleList(moduleList) { selectedModule = 0; prevColumn = -1; ascending = true; } BOOL PickDllGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) { DoDataExchange(); // attach controls DlgResize_Init(true, true); addColumnsToModuleList(ListDLLSelect); displayModuleList(ListDLLSelect); CenterWindow(); return TRUE; } LRESULT PickDllGui::OnListDllColumnClicked(NMHDR* pnmh) { NMLISTVIEW* list = (NMLISTVIEW*)pnmh; int column = list->iSubItem; if(column == prevColumn) { ascending = !ascending; } else { prevColumn = column; ascending = true; } // lo-byte: column, hi-byte: sort-order ListDLLSelect.SortItems(&listviewCompareFunc, MAKEWORD(column, ascending)); return 0; } LRESULT PickDllGui::OnListDllDoubleClick(NMHDR* pnmh) { NMITEMACTIVATE* ia = (NMITEMACTIVATE*)pnmh; LVHITTESTINFO hti; hti.pt = ia->ptAction; int clicked = ListDLLSelect.HitTest(&hti); if(clicked != -1) { selectedModule = (ModuleInfo *)ListDLLSelect.GetItemData(clicked); EndDialog(1); } return 0; } void PickDllGui::OnOK(UINT uNotifyCode, int nID, CWindow wndCtl) { int index = ListDLLSelect.GetSelectionMark(); if (index != -1) { selectedModule = (ModuleInfo *)ListDLLSelect.GetItemData(index); EndDialog(1); } } void PickDllGui::OnCancel(UINT uNotifyCode, int nID, CWindow wndCtl) { EndDialog(0); } void PickDllGui::addColumnsToModuleList(CListViewCtrl& list) { list.SetExtendedListViewStyle(LVS_EX_FULLROWSELECT, LVS_EX_FULLROWSELECT); list.InsertColumn(COL_NAME, L"Name", LVCFMT_LEFT); list.InsertColumn(COL_IMAGEBASE, L"ImageBase", LVCFMT_CENTER); list.InsertColumn(COL_IMAGESIZE, L"ImageSize", LVCFMT_CENTER); list.InsertColumn(COL_PATH, L"Path", LVCFMT_LEFT); } void PickDllGui::displayModuleList(CListViewCtrl& list) { WCHAR temp[20]; list.DeleteAllItems(); std::vector::const_iterator iter; int count = 0; for( iter = moduleList.begin(); iter != moduleList.end(); iter++ , count++) { list.InsertItem(count, iter->getFilename()); #ifdef _WIN64 swprintf_s(temp, _countof(temp), L"%016I64X", iter->modBaseAddr); #else swprintf_s(temp, _countof(temp), L"%08X", iter->modBaseAddr); #endif list.SetItemText(count, COL_IMAGEBASE, temp); swprintf_s(temp, _countof(temp),L"%08X",iter->modBaseSize); list.SetItemText(count, COL_IMAGESIZE, temp); list.SetItemText(count, COL_PATH, iter->fullPath); list.SetItemData(count, (DWORD_PTR)&(*iter)); } list.SetColumnWidth(COL_NAME, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_IMAGEBASE, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_IMAGESIZE, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_PATH, LVSCW_AUTOSIZE_USEHEADER); } // lParamSort - lo-byte: column, hi-byte: sort-order int PickDllGui::listviewCompareFunc(LPARAM lParam1, LPARAM lParam2, LPARAM lParamSort) { const ModuleInfo * module1 = (ModuleInfo *)lParam1; const ModuleInfo * module2 = (ModuleInfo *)lParam2; int column = LOBYTE(lParamSort); - bool ascending = HIBYTE(lParamSort) == true; + bool ascending = HIBYTE(lParamSort) != 0; int diff = 0; switch(column) { case COL_NAME: diff = _wcsicmp(module1->getFilename(), module2->getFilename()); break; case COL_IMAGEBASE: diff = module1->modBaseAddr < module2->modBaseAddr ? -1 : 1; break; case COL_IMAGESIZE: diff = module1->modBaseSize < module2->modBaseSize ? -1 : 1; break; case COL_PATH: diff = _wcsicmp(module1->fullPath, module2->fullPath); break; } return ascending ? diff : -diff; } diff --git a/Scylla/multitree.h b/Scylla/multitree.h index 2cd5df5..d7cb27e 100644 --- a/Scylla/multitree.h +++ b/Scylla/multitree.h @@ -1,580 +1,580 @@ #pragma once ///////////////////////////////////////////////////////////////////////////// // MultiSelectTree - A tree control with multi-select capabilities // // Written by Bjarke Viksoe (bjarke@viksoe.dk) // Copyright (c) 2005 Bjarke Viksoe. // // Add the following macro to the parent's message map: // REFLECT_NOTIFICATIONS() // // This code may be used in compiled form in any way you desire. This // source file may be redistributed by any means PROVIDING it is // not sold for profit without the authors written consent, and // providing that this notice and the authors name is included. // // This file is provided "as is" with no expressed or implied warranty. // The author accepts no liability if it causes any damage to you or your // computer whatsoever. It's free, so don't hassle me about it. // // Beware of bugs. // #ifndef __cplusplus #error WTL requires C++ compilation (use a .cpp suffix) #endif #ifndef __ATLMISC_H__ #error multitree.h requires atlmisc.h to be included first #endif #ifndef __ATLCRACK_H__ #error multitree.h requires atlcrack.h to be included first #endif #ifndef __ATLCTRLS_H__ #error multitree.h requires atlctrls.h to be included first #endif // Extended MultiSelectTree styles static const DWORD MTVS_EX_NOMARQUEE = 0x00000001; // New control notifications static const UINT TVN_ITEMSELECTING = 0x0001; static const UINT TVN_ITEMSELECTED = 0x0002; static bool operator==(CTreeItem ti1, CTreeItem ti2) { return ti1.m_hTreeItem == ti2.m_hTreeItem && ti1.m_pTreeView == ti2.m_pTreeView; } template< class T, class TBase = CTreeViewCtrlEx, class TWinTraits = CControlWinTraits > class ATL_NO_VTABLE CMultiSelectTreeViewImpl : public CWindowImpl< T, TBase, TWinTraits >, public CCustomDraw< T > { public: DECLARE_WND_SUPERCLASS(NULL, TBase::GetWndClassName()) DWORD m_dwExStyle; // Additional styles CTreeItem m_hExtSelStart; // Item where SHIFT was last pressed bool m_bMarquee; // Are we drawing rubberband? CPoint m_ptDragStart; // Point where rubberband started CPoint m_ptDragOld; // Last mousepos of rubberband CSimpleMap m_aData; CMultiSelectTreeViewImpl() : m_dwExStyle(0), m_bMarquee(false) { } // Operations BOOL SubclassWindow(HWND hWnd) { ATLASSERT(m_hWnd == NULL); ATLASSERT(::IsWindow(hWnd)); BOOL bRet = CWindowImpl< T, TBase, TWinTraits >::SubclassWindow(hWnd); if( bRet ) _Init(); return bRet; } DWORD SetExtendedTreeStyle(DWORD dwStyle) { ATLASSERT(!m_ctrlTree.IsWindow()); // Before control is created, please! DWORD dwOldStyle = m_dwTreeStyle; m_dwTreeStyle = dwStyle; return dwOldStyle; } void SelectItem(HTREEITEM hItem, BOOL bSelect) { ATLASSERT(::IsWindow(m_hWnd)); _SelectItem(hItem, bSelect == TRUE); if( bSelect ) TBase::SelectItem(hItem); } void SelectAllItems(BOOL bSelect) { ATLASSERT(::IsWindow(m_hWnd)); for( int i = 0; i < m_aData.GetSize(); i++ ) { - _SelectItem(i, bSelect); + _SelectItem(i, bSelect == TRUE); } } CTreeItem GetSelectedItem() const { ATLASSERT(false); // Not usable! return GetFirstSelectedItem(); } CTreeItem GetFocusItem() const { return TBase::GetSelectedItem(); } UINT GetItemState(HTREEITEM hItem, UINT nStateMask) const { UINT nRes = TBase::GetItemState(hItem, nStateMask); if( (nStateMask & TVIS_SELECTED) != 0 ) { int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)); if( iIndex >= 0 ) { nRes &= ~TVIS_SELECTED; if( m_aData.GetValueAt(iIndex) ) nRes |= TVIS_SELECTED; } } return nRes; } CTreeItem GetFirstSelectedItem() const { if( m_aData.GetSize() == 0 ) return NULL; for( int i = 0; i < m_aData.GetSize(); i++ ) { if( m_aData.GetValueAt(i) ) return m_aData.GetKeyAt(i); } return NULL; } CTreeItem GetNextSelectedItem(HTREEITEM hItem) const { int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)); if( iIndex < 0 ) return NULL; for( int i = iIndex + 1; i < m_aData.GetSize(); i++ ) { if( m_aData.GetValueAt(i) ) return m_aData.GetKeyAt(i); } return NULL; } int GetSelectedCount() const { int nCount = 0; for( int i = 0; i < m_aData.GetSize(); i++ ) { if( m_aData.GetValueAt(i) ) nCount++; } return nCount; } // Implementation void _Init() { ATLASSERT(::IsWindow(m_hWnd)); ModifyStyle(TVS_SHOWSELALWAYS, 0); } void _SelectItem(int iIndex, bool bSelect, int action = TVC_UNKNOWN) { if( iIndex < 0 ) return; bool bSelected = m_aData.GetValueAt(iIndex); // Don't change if state is already updated (avoids flicker) if( bSelected == bSelect ) return; CTreeItem cItem = m_aData.GetKeyAt(iIndex); CWindow parent = GetParent(); // Send notifications NMTREEVIEW nmtv = { 0 }; nmtv.hdr.code = TVN_ITEMSELECTING; nmtv.hdr.hwndFrom = m_hWnd; nmtv.hdr.idFrom = GetDlgCtrlID(); nmtv.action = action; nmtv.itemNew.hItem = cItem; nmtv.itemNew.lParam = GetItemData(cItem); nmtv.itemNew.state = bSelect ? TVIS_SELECTED : 0; nmtv.itemNew.stateMask = TVIS_SELECTED; if( parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv) != 0 ) return; // Change state m_aData.SetAtIndex(iIndex, cItem, bSelect); // Repaint item CRect rcItem; if( GetItemRect(cItem, &rcItem, FALSE) ) InvalidateRect(&rcItem, TRUE); // More notifications nmtv.hdr.code = TVN_ITEMSELECTED; parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv); } void _SelectItem(HTREEITEM hItem, bool bSelect, int action = TVC_UNKNOWN) { _SelectItem(m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)), bSelect, action); } void _SelectTree(HTREEITEM hItem, HTREEITEM hGoal, int action) { if( !_SelectTreeSub(hItem, hGoal, action) ) return; hItem = GetParentItem(hItem); while( (hItem = GetNextSiblingItem(hItem)) != NULL ) { if( !_SelectTreeSub(hItem, hGoal, action) ) return; } } bool _SelectTreeSub(HTREEITEM hItem, HTREEITEM hGoal, int action) { while( hItem != NULL ) { _SelectItem(hItem, true, action); if( hItem == hGoal ) return false; if( (TBase::GetItemState(hItem, TVIS_EXPANDED) & TVIS_EXPANDED) != 0 ) { if( !_SelectTreeSub(GetChildItem(hItem), hGoal, action) ) return false; } hItem = GetNextSiblingItem(hItem); } return true; } void _SelectBox(CRect rc) { CTreeItem hItem = GetFirstVisibleItem(); while( hItem != NULL ) { CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this); int i = m_aData.FindKey(cItem); if(i >= 0 && !m_aData.GetValueAt(i)) // ignore already selected { CRect rcItem, rcTemp; GetItemRect(hItem, &rcItem, TRUE); _SelectItem(hItem, rcTemp.IntersectRect(&rcItem, &rc) == TRUE, TVC_BYMOUSE); } hItem = GetNextVisibleItem(hItem); } } void _DrawDragRect(CPoint pt) { CClientDC dc = m_hWnd; CSize szFrame(1, 1); CRect rect(m_ptDragStart, pt); rect.NormalizeRect(); CBrush brush = CDCHandle::GetHalftoneBrush(); if( !brush.IsNull() ) { CBrushHandle hOldBrush = dc.SelectBrush(brush); dc.PatBlt(rect.left, rect.top, rect.Width(), szFrame.cy, PATINVERT); dc.PatBlt(rect.left, rect.bottom - szFrame.cy, rect.Width(), szFrame.cy, PATINVERT); dc.PatBlt(rect.left, rect.top + szFrame.cy, szFrame.cx, rect.Height() - (szFrame.cy * 2), PATINVERT); dc.PatBlt(rect.right - szFrame.cx, rect.top + szFrame.cy, szFrame.cx, rect.Height() - (szFrame.cy * 2), PATINVERT); dc.SelectBrush(hOldBrush); } } // Message map and handlers BEGIN_MSG_MAP_EX(CMultiSelectTreeViewImpl) - MSG_WM_CREATE(OnCreate) + MESSAGE_HANDLER_EX(WM_CREATE, OnCreate) MSG_WM_DESTROY(OnDestroy) MSG_WM_KEYDOWN(OnKeyDown) MSG_WM_KEYUP(OnKeyUp) MSG_WM_CHAR(OnChar) MSG_WM_SETFOCUS(OnSetFocus) MSG_WM_LBUTTONDOWN(OnLButtonDown) MSG_WM_LBUTTONUP(OnLButtonUp) MSG_WM_MOUSEMOVE(OnMouseMove) MSG_WM_CAPTURECHANGED(OnCaptureChanged) MESSAGE_HANDLER_EX(TVM_INSERTITEM, OnInsertItem) REFLECTED_NOTIFY_CODE_HANDLER_EX(TVN_DELETEITEM, OnDeleteItem) CHAIN_MSG_MAP_ALT( CCustomDraw< T >, 1 ) END_MSG_MAP() - int OnCreate(CREATESTRUCT* lpCreateStruct) + LRESULT OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam) { LRESULT lRes = DefWindowProc(); _Init(); return lRes; } void OnDestroy() { m_aData.RemoveAll(); SetMsgHandled(FALSE); } void OnKeyDown(UINT nChar, UINT nRepCnt, UINT nFlags) { if( nChar == VK_SHIFT ) m_hExtSelStart = GetFocusItem(); if( ::GetAsyncKeyState(VK_SHIFT) < 0 && m_hExtSelStart == GetFocusItem() ) { switch( nChar ) { case VK_UP: case VK_DOWN: case VK_HOME: case VK_END: case VK_NEXT: case VK_PRIOR: for( int i = 0; i < m_aData.GetSize(); i++ ) { _SelectItem(i, false, TVC_BYKEYBOARD); } } } SetMsgHandled(FALSE); } void OnKeyUp(UINT nChar, UINT nRepCnt, UINT nFlags) { if( ::GetAsyncKeyState(VK_SHIFT) < 0 ) { switch( nChar ) { case VK_UP: case VK_DOWN: case VK_HOME: case VK_END: case VK_NEXT: case VK_PRIOR: HTREEITEM hItem = GetFocusItem(); // Is current or first-shift-item the upper item? CRect rcItem1, rcItem2; GetItemRect(m_hExtSelStart, &rcItem1, TRUE); GetItemRect(hItem, &rcItem2, TRUE); // Select from current item to item where SHIFT was pressed if( rcItem1.top > rcItem2.top ) _SelectTree(hItem, m_hExtSelStart, TVC_BYKEYBOARD); else _SelectTree(m_hExtSelStart, hItem, TVC_BYKEYBOARD); _SelectItem(hItem, true, TVC_BYKEYBOARD); } } SetMsgHandled(FALSE); } void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags) { if( nChar == VK_SPACE ) { HTREEITEM hItem = GetFocusItem(); _SelectItem(hItem, (GetItemState(hItem, TVIS_SELECTED) & TVIS_SELECTED) == 0, TVC_BYKEYBOARD); return; } SetMsgHandled(FALSE); } void OnSetFocus(CWindow wndOld) { DefWindowProc(); // FIX: We really need the focus-rectangle in this control since it // improves the navigation a lot. So let's ask Windows to display it. SendMessage(WM_UPDATEUISTATE, MAKEWPARAM(UIS_CLEAR, UISF_HIDEFOCUS)); } void OnLButtonDown(UINT nFlags, CPoint point) { SetMsgHandled(FALSE); // Hit-test and figure out where we're clicking... TVHITTESTINFO hti = { 0 }; hti.pt = point; HTREEITEM hItem = HitTest(&hti); if( (hItem == NULL || (hti.flags & TVHT_ONITEMRIGHT) != 0) ) { if( (m_dwExStyle & MTVS_EX_NOMARQUEE) == 0 && ::DragDetect(m_hWnd, point) ) { // Great we're dragging a rubber-band // Clear selection of CTRL is not down if( ::GetAsyncKeyState(VK_CONTROL) >= 0 ) { for( int i = 0; i < m_aData.GetSize(); i++ ) { _SelectItem(i, false, TVC_BYMOUSE); } UpdateWindow(); } // Now start drawing the rubber-band... SetCapture(); m_ptDragStart = m_ptDragOld = point; _DrawDragRect(point); m_bMarquee = true; SetMsgHandled(TRUE); return; } } if( hItem == NULL ) return; if( (hti.flags & TVHT_ONITEMBUTTON) != 0 ) return; // Great, let's do an advanced selection if( (hti.flags & TVHT_ONITEMRIGHT) != 0 ) { for( int i = 0; i < m_aData.GetSize(); i++ ) { _SelectItem(i, false, TVC_BYMOUSE); } return; } int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)); if( iIndex < 0 ) return; // Simulate drag'n'drop? if( m_aData.GetValueAt(iIndex) && (GetStyle() & TVS_DISABLEDRAGDROP) == 0 && ::DragDetect(m_hWnd, point) ) { NMTREEVIEW nmtv = { 0 }; nmtv.hdr.code = TVN_BEGINDRAG; nmtv.hdr.hwndFrom = m_hWnd; nmtv.hdr.idFrom = GetDlgCtrlID(); nmtv.itemNew.hItem = hItem; nmtv.itemNew.lParam = GetItemData(hItem); CWindow parent = GetParent(); parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv); } bool bSelected = m_aData.GetValueAt(iIndex); if( ::GetAsyncKeyState(VK_SHIFT) < 0 ) { // Is current or first-shift-item the upper item? CRect rcItem1, rcItem2; GetItemRect(m_hExtSelStart, &rcItem1, TRUE); GetItemRect(hItem, &rcItem2, TRUE); // Select from current item to item where SHIFT was pressed if( rcItem1.top > rcItem2.top ) _SelectTree(hItem, m_hExtSelStart, TVC_BYMOUSE); else _SelectTree(m_hExtSelStart, hItem, TVC_BYMOUSE); } else if( ::GetAsyncKeyState(VK_CONTROL) < 0 ) { // Just toggle item _SelectItem(iIndex, !bSelected, TVC_BYMOUSE); } else { // Remove current selection and replace it with clicked item for( int i = 0; i < m_aData.GetSize(); i++ ) { _SelectItem(i, i == iIndex, TVC_BYMOUSE); } } } void OnLButtonUp(UINT nFlags, CPoint point) { if( m_bMarquee ) ReleaseCapture(); SetMsgHandled(FALSE); } void OnMouseMove(UINT nFlags, CPoint point) { if( m_bMarquee ) { CRect rc(m_ptDragStart, point); _DrawDragRect(m_ptDragOld); rc.NormalizeRect(); _SelectBox(rc); UpdateWindow(); _DrawDragRect(point); m_ptDragOld = point; } SetMsgHandled(FALSE); } void OnCaptureChanged(CWindow wnd) { if( m_bMarquee ) { _DrawDragRect(m_ptDragOld); m_bMarquee = false; } SetMsgHandled(FALSE); } LRESULT OnInsertItem(UINT uMsg, WPARAM wParam, LPARAM lParam) { HTREEITEM hItem = (HTREEITEM) DefWindowProc(uMsg, wParam, lParam); if( hItem == NULL ) return (LRESULT) hItem; CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this); // We manage a bit of extra information for each item. We'll store // this in an ATL::CSimpleMap. Not a particular speedy structure for lookups. // Don't keep too many items around in the tree! m_aData.Add(cItem, false); return (LRESULT) hItem; } LRESULT OnDeleteItem(NMHDR* pnmh) { const NMTREEVIEW* lpNMTV = (NMTREEVIEW*) pnmh; m_aData.Remove(CTreeItem(lpNMTV->itemNew.hItem, (CTreeViewCtrlEx*)this)); return 0; } // Custom Draw DWORD OnPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* /*lpNMCustomDraw*/) { return CDRF_NOTIFYITEMDRAW; // We need per-item notifications } DWORD OnItemPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* lpNMCustomDraw) { NMTVCUSTOMDRAW* lpTVCD = (NMTVCUSTOMDRAW*) lpNMCustomDraw; HTREEITEM hItem = (HTREEITEM) lpTVCD->nmcd.dwItemSpec; int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)); if( iIndex >= 0 ) { bool bSelected = m_aData.GetValueAt(iIndex); // Trick TreeView into displaying correct selection colors if( bSelected ) { lpTVCD->clrText = ::GetSysColor(COLOR_HIGHLIGHTTEXT); lpTVCD->clrTextBk = ::GetSysColor(COLOR_HIGHLIGHT); } else { // Special case of tree-item actually have selection, but our // state says it is currently not selected (CTRL+click on same item twice). if( (lpTVCD->nmcd.uItemState & CDIS_SELECTED) != 0 ) { COLORREF clrText = GetTextColor(); if( clrText == CLR_NONE ) clrText = ::GetSysColor(COLOR_WINDOWTEXT); COLORREF clrBack = GetBkColor(); if( clrBack == CLR_NONE ) clrBack = ::GetSysColor(COLOR_WINDOW); //CDCHandle dc = lpTVCD->nmcd.hdc; //dc.SetTextColor(clrText); //dc.SetBkColor(clrBack); lpTVCD->clrText = clrText; lpTVCD->clrTextBk = clrBack; } } return CDRF_NEWFONT; } return CDRF_DODEFAULT; } }; class CMultiSelectTreeViewCtrl : public CMultiSelectTreeViewImpl > { public: DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName()) };