From 6db02513462ff0b3b9317e0bf012fd341238d017 Mon Sep 17 00:00:00 2001 From: Filip Olszak Date: Fri, 9 Apr 2021 01:51:31 +0200 Subject: [PATCH] simplify code to improve readability --- agent/AgentWorker.cpp | 153 ++++++++++++++++----------------------- agent/DetectionLogic.cpp | 21 +++--- agent/DetectionLogic.h | 19 ++++- agent/Helpers.h | 1 + 4 files changed, 92 insertions(+), 102 deletions(-) diff --git a/agent/AgentWorker.cpp b/agent/AgentWorker.cpp index 27ebbc9f..10265d93 100644 --- a/agent/AgentWorker.cpp +++ b/agent/AgentWorker.cpp @@ -5,44 +5,50 @@ #include "DetectionLogic.h" #include "YaraInstance.h" -void report_detection(int detId, map evt_body) { +using namespace krabs; + +void report_detection(int detId, GenericEvent evt) { using std::to_string; + using std::string; - std::string sDump; - std::string sOutBody; + string sDump; + string sOutBody; - std::string procId = to_string(evt_body[L"CallingProcessId"]); - std::string procImage = get_pname(evt_body[L"CallingProcessId"]); - std::string targetProcId = to_string(evt_body[L"TargetProcessId"]); - std::string targetProcImage = get_pname(evt_body[L"TargetProcessId"]); - std::string protMask = itohs(evt_body[L"ProtectionMask"]); - std::string allocType = itohs(evt_body[L"AllocationType"]); - std::string size = to_string(evt_body[L"RegionSize"]); - std::string baseAddr = itohs(evt_body[L"BaseAddress"]); + string procId = to_string(evt.fields[L"CallingProcessId"]); + string procImage = get_pname(evt.fields[L"CallingProcessId"]); + string targetProcId = to_string(evt.fields[L"TargetProcessId"]); + string targetProcImage = get_pname(evt.fields[L"TargetProcessId"]); + string protMask = itohs(evt.fields[L"ProtectionMask"]); + string allocType = itohs(evt.fields[L"AllocationType"]); + string size = to_string(evt.fields[L"RegionSize"]); + string baseAddr = itohs(evt.fields[L"BaseAddress"]); switch (detId) { - case ALLOCVM_REMOTE_META_GENERIC: - sDump = dump_memory_ascii(evt_body[L"TargetProcessId"], evt_body[L"BaseAddress"], MEM_STR_SIZE); - sOutBody = "\n\n\n\nANOMALOUS MEMORY ALLOCATION DETECTED \n\n"; - sOutBody += "[+] Source: " + procImage + " (PID: " + procId + ")\n"; - sOutBody += "[+] Target: " + targetProcImage + " (PID: " + targetProcId + ")\n"; - sOutBody += "[+] Protection: " + protMask + "\n"; - sOutBody += "[+] Allocation: " + allocType + "\n"; - sOutBody += "[+] Region size: " + size + "\n"; - sOutBody += "[+] Base address: " + baseAddr + "\n"; - sOutBody += "[+] MZ-header: "; - if (sDump.rfind("MZ", 0) == 0) { - sOutBody += "YES\n\n"; - } - else { - sOutBody += "NO\n\n"; - } - sOutBody += "[+] Memory at location: \n\n"; - sOutBody += sDump; - break; - case ALLOCVM_REMOTE_SIGNATURES: - default: - return; + case ALLOCVM_REMOTE_META_GENERIC: + sDump = dump_memory_ascii(evt.fields[L"TargetProcessId"], evt.fields[L"BaseAddress"], MEM_STR_SIZE); + + sOutBody = "\n\n\n\nANOMALOUS MEMORY ALLOCATION DETECTED \n\n"; + sOutBody += "[+] Source: " + procImage + " (PID: " + procId + ")\n"; + sOutBody += "[+] Target: " + targetProcImage + " (PID: " + targetProcId + ")\n"; + sOutBody += "[+] Protection: " + protMask + "\n"; + sOutBody += "[+] Allocation: " + allocType + "\n"; + sOutBody += "[+] Region size: " + size + "\n"; + sOutBody += "[+] Base address: " + baseAddr + "\n"; + sOutBody += "[+] MZ-header: "; + + if (sDump.rfind("MZ", 0) == 0) { + sOutBody += "YES\n\n"; + } + else { + sOutBody += "NO\n\n"; + } + + sOutBody += "[+] Memory at location: \n\n"; + sOutBody += sDump; + break; + case ALLOCVM_REMOTE_SIGNATURES: + default: + return; } if (sOutBody.empty()) { @@ -56,73 +62,40 @@ void report_detection(int detId, map evt_body) { return; } -// Parse KERNEL_THREATINT_TASK_ALLOCVM_REMOTE -map parse_allocvm_remote(krabs::schema schema, krabs::parser parser) { - map zero_map; - map allocation_fields = { {(wstring)L"CallingProcessId",0}, - {(wstring)L"TargetProcessId",0}, - {(wstring)L"AllocationType",0}, - {(wstring)L"ProtectionMask",0}, - {(wstring)L"RegionSize",0}, - {(wstring)L"BaseAddress",0} - }; +VOID parse_generic_event(const EVENT_RECORD& record, const trace_context& trace_context) { + schema schema(record, trace_context.schema_locator); + parser parser(schema); + + GenericEvent new_event; + + new_event.type = schema.event_id(); try { - for (krabs::property property : parser.properties()) { - std::wstring wsPropertyName = property.name(); - if (allocation_fields.find(wsPropertyName) != allocation_fields.end()) { + for (property property : parser.properties()) { + wstring wsPropertyName = property.name(); + if (new_event.fields.find(wsPropertyName) != new_event.fields.end()) { switch (property.type()) { - case TDH_INTYPE_UINT32: - allocation_fields[wsPropertyName] = parser.parse(wsPropertyName); - break; - case TDH_INTYPE_POINTER: - allocation_fields[wsPropertyName] = parser.parse(wsPropertyName).address; - break; + case TDH_INTYPE_UINT32: + new_event.fields[wsPropertyName] = parser.parse(wsPropertyName); + break; + case TDH_INTYPE_POINTER: + new_event.fields[wsPropertyName] = parser.parse(wsPropertyName).address; + break; } } } - return allocation_fields; } + catch (...) { log_debug(L"Error parsing the event\n"); - return zero_map; + return; } -} - -VOID parse_single_event(const EVENT_RECORD& record, const krabs::trace_context& trace_context) { - krabs::schema schema(record, trace_context.schema_locator); - krabs::parser parser(schema); - - map parsed_event; - int eid = schema.event_id(); - - switch (eid) { - case KERNEL_THREATINT_TASK_ALLOCVM_REMOTE: - parsed_event = parse_allocvm_remote(schema, parser); - break; - // Currently unsupported event types - case KERNEL_THREATINT_TASK_PROTECTVM_REMOTE: - case KERNEL_THREATINT_TASK_MAPVIEW_REMOTE: - case KERNEL_THREATINT_TASK_QUEUEUSERAPC_REMOTE: - case KERNEL_THREATINT_TASK_SETTHREADCONTEXT_REMOTE: - case KERNEL_THREATINT_TASK_ALLOCVM_LOCAL: - case KERNEL_THREATINT_TASK_PROTECTVM_LOCAL: - case KERNEL_THREATINT_TASK_MAPVIEW_LOCAL: - case KERNEL_THREATINT_TASK_QUEUEUSERAPC_LOCAL: - case KERNEL_THREATINT_TASK_SETTHREADCONTEXT_LOCAL: - case KERNEL_THREATINT_TASK_READVM_LOCAL: - case KERNEL_THREATINT_TASK_WRITEVM_LOCAL: - case KERNEL_THREATINT_TASK_READVM_REMOTE: - case KERNEL_THREATINT_TASK_WRITEVM_REMOTE: - default: - return; - } - if (parsed_event.empty()) { + if (new_event.fields.empty()) { log_debug(L"TiEtwAgent: Failed to parse an event\n"); } else { - detect_event(parsed_event, eid); + detect_event(new_event); } return; } @@ -132,13 +105,13 @@ DWORD agent_worker() DWORD ret{ 0 }; log_debug(L"TiEtwAgent: Started the agent worker\n"); - krabs::user_trace trace(ETS_NAME); - krabs::provider<> provider(L"Microsoft-Windows-Threat-Intelligence"); - krabs::event_filter filter(krabs::predicates::id_is((int)KERNEL_THREATINT_TASK_ALLOCVM_REMOTE)); + user_trace trace(ETS_NAME); + provider<> provider(L"Microsoft-Windows-Threat-Intelligence"); + event_filter filter(predicates::id_is((int)KERNEL_THREATINT_TASK_ALLOCVM_REMOTE)); try { log_debug(L"TiEtwAgent: Setting up the trace session\n"); - provider.add_on_event_callback(parse_single_event); + provider.add_on_event_callback(parse_generic_event); provider.add_filter(filter); trace.enable(provider); diff --git a/agent/DetectionLogic.cpp b/agent/DetectionLogic.cpp index 262e22d7..13469d1c 100644 --- a/agent/DetectionLogic.cpp +++ b/agent/DetectionLogic.cpp @@ -6,28 +6,29 @@ const int ALLOC_PROTECTION{ PAGE_EXECUTE_READWRITE }; const int ALLOC_TYPE{ MEM_RESERVE | MEM_COMMIT }; const int MIN_REGION_SIZE{ 10240 }; -VOID allocvm_remote_meta_generic(std::map alloc_event) { - if (alloc_event[L"RegionSize"] >= MIN_REGION_SIZE) { - if (alloc_event[L"AllocationType"] == ALLOC_TYPE) { - if (alloc_event[L"ProtectionMask"] == ALLOC_PROTECTION) { +DWORD allocvm_remote_meta_generic(GenericEvent alloc_event) { + if (alloc_event.fields[L"RegionSize"] >= MIN_REGION_SIZE) { + if (alloc_event.fields[L"AllocationType"] == ALLOC_TYPE) { + if (alloc_event.fields[L"ProtectionMask"] == ALLOC_PROTECTION) { report_detection(ALLOCVM_REMOTE_META_GENERIC, alloc_event); + return TRUE; } } } - return; + return FALSE; } // Trigger Yara scan of the remotely allocated memory page -VOID allocvm_remote_signatures(std::map alloc_event) { +VOID allocvm_remote_signatures(GenericEvent alloc_event) { return; } -VOID detect_event(std::map parsed_event, int eid) { +VOID detect_event(GenericEvent evt) { // Run detection functions depending on source event type - switch (eid) { + switch (evt.type) { case KERNEL_THREATINT_TASK_ALLOCVM_REMOTE: - allocvm_remote_meta_generic(parsed_event); - allocvm_remote_signatures(parsed_event); + allocvm_remote_meta_generic(evt); + allocvm_remote_signatures(evt); break; case KERNEL_THREATINT_TASK_PROTECTVM_REMOTE: case KERNEL_THREATINT_TASK_MAPVIEW_REMOTE: diff --git a/agent/DetectionLogic.h b/agent/DetectionLogic.h index 7419273a..5e7457db 100644 --- a/agent/DetectionLogic.h +++ b/agent/DetectionLogic.h @@ -35,5 +35,20 @@ enum DETECTIONS { extern map allocation_fields; -VOID report_detection(int evtId, map evt_body); -VOID detect_event(std::map parsed_event, int eid); +class GenericEvent { +public: + uint8_t type; + + map fields = { + {(wstring)L"CallingProcessId",0}, + {(wstring)L"TargetProcessId",0}, + {(wstring)L"AllocationType",0}, + {(wstring)L"ProtectionMask",0}, + {(wstring)L"RegionSize",0}, + {(wstring)L"BaseAddress",0} + }; +}; + + +VOID report_detection(int detId, GenericEvent evt); +VOID detect_event(GenericEvent evt); diff --git a/agent/Helpers.h b/agent/Helpers.h index e62cb26b..333cc9f6 100644 --- a/agent/Helpers.h +++ b/agent/Helpers.h @@ -5,6 +5,7 @@ #define MAX_BUF_SIZE 2048 #define MEM_STR_SIZE 512 +#define GET_VARIABLE_NAME(Variable) (#Variable) std::string itohs(uint64_t i); std::string get_pname(uint64_t pid);