Skip to content

Commit 613a832

Browse files
committed
OnlineChecks: Add VirusTotal column to process modules tab
1 parent 8a34d16 commit 613a832

File tree

3 files changed

+204
-80
lines changed

3 files changed

+204
-80
lines changed

plugins/OnlineChecks/main.c

Lines changed: 193 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@ PH_CALLBACK_REGISTRATION ProcessMenuInitializingCallbackRegistration;
3434
PH_CALLBACK_REGISTRATION ModuleMenuInitializingCallbackRegistration;
3535
PH_CALLBACK_REGISTRATION TreeNewMessageCallbackRegistration;
3636
PH_CALLBACK_REGISTRATION ProcessTreeNewInitializingCallbackRegistration;
37+
PH_CALLBACK_REGISTRATION ModulesTreeNewInitializingCallbackRegistration;
3738

3839
BOOLEAN VirusTotalScanningEnabled = FALSE;
39-
HWND ProcessTreeNewHandle = NULL;
4040
ULONG ProcessesUpdatedCount = 0;
4141
LIST_ENTRY ProcessListHead = { &ProcessListHead, &ProcessListHead };
4242
PH_QUEUED_LOCK ProcessesListLock = PH_QUEUED_LOCK_INIT;
@@ -47,28 +47,45 @@ VOID ProcessesUpdatedCallback(
4747
)
4848
{
4949
#ifdef VIRUSTOTAL_API
50+
static ULONG ProcessesUpdatedCount = 0;
5051
PLIST_ENTRY listEntry;
5152

5253
if (!VirusTotalScanningEnabled)
5354
return;
5455

56+
ProcessesUpdatedCount++;
57+
58+
if (ProcessesUpdatedCount < 2)
59+
return;
60+
5561
listEntry = ProcessListHead.Flink;
5662

5763
while (listEntry != &ProcessListHead)
5864
{
5965
PPROCESS_EXTENSION extension;
60-
PPH_PROCESS_ITEM processItem;
6166

6267
extension = CONTAINING_RECORD(listEntry, PROCESS_EXTENSION, ListEntry);
63-
processItem = extension->ProcessItem;
6468

65-
if (!PH_IS_FAKE_PROCESS_ID(processItem->ProcessId) && !PhIsNullOrEmptyString(processItem->FileName))
69+
PPH_STRING filePath = NULL;
70+
PPH_PROCESS_ITEM processItem = extension->ProcessItem;
71+
PPH_MODULE_ITEM moduleItem = extension->ModuleItem;
72+
73+
if (processItem)
74+
{
75+
filePath = processItem->FileName;
76+
}
77+
else if (moduleItem)
78+
{
79+
filePath = moduleItem->FileName;
80+
}
81+
82+
if (!PhIsNullOrEmptyString(filePath)) // !PH_IS_FAKE_PROCESS_ID(processItem->ProcessId)
6683
{
6784
if (!extension->ResultValid)
6885
{
6986
PPROCESS_DB_OBJECT object;
7087

71-
if (object = FindProcessDbObject(&processItem->FileName->sr))
88+
if (object = FindProcessDbObject(&filePath->sr))
7289
{
7390
extension->Stage1 = TRUE;
7491
extension->ResultValid = TRUE;
@@ -80,9 +97,9 @@ VOID ProcessesUpdatedCallback(
8097

8198
if (!extension->Stage1)
8299
{
83-
if (!VirusTotalGetCachedResult(processItem->FileName))
100+
if (!VirusTotalGetCachedResult(filePath))
84101
{
85-
VirusTotalAddCacheResult(processItem, extension);
102+
VirusTotalAddCacheResult(filePath, extension);
86103
}
87104

88105
extension->Stage1 = TRUE;
@@ -274,8 +291,7 @@ VOID NTAPI ModuleMenuInitializingCallback(
274291
}
275292
}
276293

277-
278-
LONG NTAPI VirusTotalNodeSortFunction(
294+
LONG NTAPI VirusTotalProcessNodeSortFunction(
279295
_In_ PVOID Node1,
280296
_In_ PVOID Node2,
281297
_In_ ULONG SubId,
@@ -290,6 +306,21 @@ LONG NTAPI VirusTotalNodeSortFunction(
290306
return PhCompareStringWithNull(extension1->VirusTotalResult, extension2->VirusTotalResult, TRUE);
291307
}
292308

309+
LONG NTAPI VirusTotalModuleNodeSortFunction(
310+
_In_ PVOID Node1,
311+
_In_ PVOID Node2,
312+
_In_ ULONG SubId,
313+
_In_ PVOID Context
314+
)
315+
{
316+
PPH_MODULE_NODE node1 = Node1;
317+
PPH_MODULE_NODE node2 = Node2;
318+
PPROCESS_EXTENSION extension1 = PhPluginGetObjectExtension(PluginInstance, node1->ModuleItem, EmModuleItemType);
319+
PPROCESS_EXTENSION extension2 = PhPluginGetObjectExtension(PluginInstance, node2->ModuleItem, EmModuleItemType);
320+
321+
return PhCompareStringWithNull(extension1->VirusTotalResult, extension2->VirusTotalResult, TRUE);
322+
}
323+
293324
VOID NTAPI ProcessTreeNewInitializingCallback(
294325
_In_opt_ PVOID Parameter,
295326
_In_opt_ PVOID Context
@@ -298,14 +329,32 @@ VOID NTAPI ProcessTreeNewInitializingCallback(
298329
PPH_PLUGIN_TREENEW_INFORMATION info = Parameter;
299330
PH_TREENEW_COLUMN column;
300331

301-
*(HWND*)Context = info->TreeNewHandle;
332+
//*(HWND*)Context = info->TreeNewHandle;
333+
334+
memset(&column, 0, sizeof(PH_TREENEW_COLUMN));
335+
column.Text = L"VirusTotal";
336+
column.Width = 140;
337+
column.Alignment = PH_ALIGN_CENTER;
338+
column.CustomDraw = TRUE;
339+
PhPluginAddTreeNewColumn(PluginInstance, info->CmData, &column, NETWORK_COLUMN_ID_VIUSTOTAL, NULL, VirusTotalProcessNodeSortFunction);
340+
}
341+
342+
VOID NTAPI ModuleTreeNewInitializingCallback(
343+
_In_opt_ PVOID Parameter,
344+
_In_opt_ PVOID Context
345+
)
346+
{
347+
PPH_PLUGIN_TREENEW_INFORMATION info = Parameter;
348+
PH_TREENEW_COLUMN column;
349+
350+
//*(HWND*)Context = info->TreeNewHandle;
302351

303352
memset(&column, 0, sizeof(PH_TREENEW_COLUMN));
304353
column.Text = L"VirusTotal";
305354
column.Width = 140;
306355
column.Alignment = PH_ALIGN_CENTER;
307356
column.CustomDraw = TRUE;
308-
PhPluginAddTreeNewColumn(PluginInstance, info->CmData, &column, NETWORK_COLUMN_ID_VIUSTOTAL, NULL, VirusTotalNodeSortFunction);
357+
PhPluginAddTreeNewColumn(PluginInstance, info->CmData, &column, NETWORK_COLUMN_ID_VIUSTOTAL_MODULE, NULL, VirusTotalModuleNodeSortFunction);
309358
}
310359

311360
VOID NTAPI TreeNewMessageCallback(
@@ -319,87 +368,110 @@ VOID NTAPI TreeNewMessageCallback(
319368
{
320369
case TreeNewGetCellText:
321370
{
322-
if (message->TreeNewHandle == ProcessTreeNewHandle)
371+
PPH_TREENEW_GET_CELL_TEXT getCellText = message->Parameter1;
372+
373+
switch (message->SubId)
323374
{
324-
PPH_TREENEW_GET_CELL_TEXT getCellText = message->Parameter1;
325-
PPH_PROCESS_NODE processNode = (PPH_PROCESS_NODE)getCellText->Node;
326-
PPROCESS_EXTENSION extension = PhPluginGetObjectExtension(PluginInstance, processNode->ProcessItem, EmProcessItemType);
375+
case NETWORK_COLUMN_ID_VIUSTOTAL:
376+
{
377+
PPH_PROCESS_NODE processNode = (PPH_PROCESS_NODE)getCellText->Node;
378+
PPROCESS_EXTENSION extension = PhPluginGetObjectExtension(PluginInstance, processNode->ProcessItem, EmProcessItemType);
327379

328-
switch (message->SubId)
380+
getCellText->Text = PhGetStringRef(extension->VirusTotalResult);
381+
}
382+
break;
383+
case NETWORK_COLUMN_ID_VIUSTOTAL_MODULE:
329384
{
330-
case NETWORK_COLUMN_ID_VIUSTOTAL:
331-
{
332-
getCellText->Text = PhGetStringRef(extension->VirusTotalResult);
333-
}
334-
break;
385+
PPH_MODULE_NODE processNode = (PPH_MODULE_NODE)getCellText->Node;
386+
PPROCESS_EXTENSION extension = PhPluginGetObjectExtension(PluginInstance, processNode->ModuleItem, EmModuleItemType);
387+
388+
getCellText->Text = PhGetStringRef(extension->VirusTotalResult);
335389
}
390+
break;
336391
}
337392
}
338393
break;
339394
case TreeNewCustomDraw:
340395
{
341-
if (message->TreeNewHandle == ProcessTreeNewHandle)
342-
{
343-
PPH_TREENEW_CUSTOM_DRAW customDraw = message->Parameter1;
344-
PPH_PROCESS_NODE processNode = (PPH_PROCESS_NODE)customDraw->Node;
345-
PPROCESS_EXTENSION extension = PhPluginGetObjectExtension(PluginInstance, processNode->ProcessItem, EmProcessItemType);
346-
PH_STRINGREF text;
347-
SIZE textSize;
348-
349-
if (!VirusTotalScanningEnabled)
350-
{
351-
static PH_STRINGREF disabledText = PH_STRINGREF_INIT(L"VirusTotal disabled");
352-
353-
GetTextExtentPoint32(
354-
customDraw->Dc,
355-
disabledText.Buffer,
356-
(ULONG)disabledText.Length / 2,
357-
&textSize
358-
);
396+
PPH_TREENEW_CUSTOM_DRAW customDraw = message->Parameter1;
397+
PPROCESS_EXTENSION extension = NULL;
398+
PH_STRINGREF text;
399+
SIZE textSize;
359400

360-
DrawText(
361-
customDraw->Dc,
362-
disabledText.Buffer,
363-
(ULONG)disabledText.Length / 2,
364-
&customDraw->CellRect,
365-
DT_CENTER | DT_VCENTER | DT_END_ELLIPSIS | DT_SINGLELINE
366-
);
367-
368-
return;
369-
}
370-
371-
if (extension->Positives)
372-
{
373-
//fontHandle = CommonCreateFont(-14, FW_BOLD, NULL);
374-
SetTextColor(customDraw->Dc, RGB(0xff, 0, 0));
375-
}
376-
else
377-
{
378-
//fontHandle = CommonDuplicateFont((HFONT)SendMessage(ProcessTreeNewHandle, WM_GETFONT, 0, 0));
379-
SetTextColor(customDraw->Dc, RGB(0x64, 0x64, 0x64));
380-
}
381-
382-
// originalFont = SelectObject(customDraw->Dc, fontHandle);
383-
text = PhGetStringRef(extension->VirusTotalResult);
401+
if (!VirusTotalScanningEnabled)
402+
{
403+
static PH_STRINGREF disabledText = PH_STRINGREF_INIT(L"VirusTotal disabled");
384404

385405
GetTextExtentPoint32(
386-
customDraw->Dc,
387-
text.Buffer,
388-
(ULONG)text.Length / 2,
406+
customDraw->Dc,
407+
disabledText.Buffer,
408+
(ULONG)disabledText.Length / 2,
389409
&textSize
390410
);
391411

392412
DrawText(
393413
customDraw->Dc,
394-
text.Buffer,
395-
(ULONG)text.Length / 2,
414+
disabledText.Buffer,
415+
(ULONG)disabledText.Length / 2,
396416
&customDraw->CellRect,
397417
DT_CENTER | DT_VCENTER | DT_END_ELLIPSIS | DT_SINGLELINE
398418
);
399419

400-
//SelectObject(customDraw->Dc, originalFont);
401-
//DeleteObject(fontHandle);
420+
return;
402421
}
422+
423+
switch (message->SubId)
424+
{
425+
case NETWORK_COLUMN_ID_VIUSTOTAL:
426+
{
427+
PPH_PROCESS_NODE processNode = (PPH_PROCESS_NODE)customDraw->Node;
428+
429+
extension = PhPluginGetObjectExtension(PluginInstance, processNode->ProcessItem, EmProcessItemType);
430+
}
431+
break;
432+
case NETWORK_COLUMN_ID_VIUSTOTAL_MODULE:
433+
{
434+
PPH_MODULE_NODE processNode = (PPH_MODULE_NODE)customDraw->Node;
435+
436+
extension = PhPluginGetObjectExtension(PluginInstance, processNode->ModuleItem, EmModuleItemType);
437+
}
438+
break;
439+
}
440+
441+
if (!extension)
442+
break;
443+
444+
if (extension->Positives)
445+
{
446+
//fontHandle = CommonCreateFont(-14, FW_BOLD, NULL);
447+
SetTextColor(customDraw->Dc, RGB(0xff, 0, 0));
448+
}
449+
else
450+
{
451+
//fontHandle = CommonDuplicateFont((HFONT)SendMessage(ProcessTreeNewHandle, WM_GETFONT, 0, 0));
452+
SetTextColor(customDraw->Dc, RGB(0x64, 0x64, 0x64));
453+
}
454+
455+
// originalFont = SelectObject(customDraw->Dc, fontHandle);
456+
text = PhGetStringRef(extension->VirusTotalResult);
457+
458+
GetTextExtentPoint32(
459+
customDraw->Dc,
460+
text.Buffer,
461+
(ULONG)text.Length / 2,
462+
&textSize
463+
);
464+
465+
DrawText(
466+
customDraw->Dc,
467+
text.Buffer,
468+
(ULONG)text.Length / 2,
469+
&customDraw->CellRect,
470+
DT_CENTER | DT_VCENTER | DT_END_ELLIPSIS | DT_SINGLELINE
471+
);
472+
473+
//SelectObject(customDraw->Dc, originalFont);
474+
//DeleteObject(fontHandle);
403475
}
404476
break;
405477
}
@@ -440,6 +512,41 @@ VOID NTAPI ProcessItemDeleteCallback(
440512
PhReleaseQueuedLockExclusive(&ProcessesListLock);
441513
}
442514

515+
VOID NTAPI ModuleItemCreateCallback(
516+
_In_ PVOID Object,
517+
_In_ PH_EM_OBJECT_TYPE ObjectType,
518+
_In_ PVOID Extension
519+
)
520+
{
521+
PPH_MODULE_ITEM moduleItem = Object;
522+
PPROCESS_EXTENSION extension = Extension;
523+
524+
memset(extension, 0, sizeof(PROCESS_EXTENSION));
525+
526+
extension->ModuleItem = moduleItem;
527+
extension->VirusTotalResult = PhReferenceEmptyString();
528+
529+
PhAcquireQueuedLockExclusive(&ProcessesListLock);
530+
InsertTailList(&ProcessListHead, &extension->ListEntry);
531+
PhReleaseQueuedLockExclusive(&ProcessesListLock);
532+
}
533+
534+
VOID NTAPI ModuleItemDeleteCallback(
535+
_In_ PVOID Object,
536+
_In_ PH_EM_OBJECT_TYPE ObjectType,
537+
_In_ PVOID Extension
538+
)
539+
{
540+
PPH_MODULE_ITEM processItem = Object;
541+
PPROCESS_EXTENSION extension = Extension;
542+
543+
PhClearReference(&extension->VirusTotalResult);
544+
545+
PhAcquireQueuedLockExclusive(&ProcessesListLock);
546+
RemoveEntryList(&extension->ListEntry);
547+
PhReleaseQueuedLockExclusive(&ProcessesListLock);
548+
}
549+
443550
LOGICAL DllMain(
444551
_In_ HINSTANCE Instance,
445552
_In_ ULONG Reason,
@@ -514,9 +621,16 @@ LOGICAL DllMain(
514621
PhRegisterCallback(
515622
PhGetGeneralCallback(GeneralCallbackProcessTreeNewInitializing),
516623
ProcessTreeNewInitializingCallback,
517-
&ProcessTreeNewHandle,
624+
NULL,
518625
&ProcessTreeNewInitializingCallbackRegistration
519626
);
627+
PhRegisterCallback(
628+
PhGetGeneralCallback(GeneralCallbackModuleTreeNewInitializing),
629+
ModuleTreeNewInitializingCallback,
630+
NULL,
631+
&ModulesTreeNewInitializingCallbackRegistration
632+
);
633+
520634
PhRegisterCallback(
521635
PhGetPluginCallback(PluginInstance, PluginCallbackTreeNewMessage),
522636
TreeNewMessageCallback,
@@ -532,6 +646,14 @@ LOGICAL DllMain(
532646
ProcessItemDeleteCallback
533647
);
534648

649+
PhPluginSetObjectExtension(
650+
PluginInstance,
651+
EmModuleItemType,
652+
sizeof(PROCESS_EXTENSION),
653+
ModuleItemCreateCallback,
654+
ModuleItemDeleteCallback
655+
);
656+
535657
PhAddSettings(settings, ARRAYSIZE(settings));
536658
}
537659
break;

0 commit comments

Comments
 (0)