Skip to content

Commit

Permalink
simplify code to improve readability
Browse files Browse the repository at this point in the history
  • Loading branch information
Filip Olszak committed Apr 8, 2021
1 parent de2952c commit 6db0251
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 102 deletions.
153 changes: 63 additions & 90 deletions agent/AgentWorker.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,44 +5,50 @@
#include "DetectionLogic.h"
#include "YaraInstance.h"

void report_detection(int detId, map<wstring, uint64_t> 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()) {
Expand All @@ -56,73 +62,40 @@ void report_detection(int detId, map<wstring, uint64_t> evt_body) {
return;
}

// Parse KERNEL_THREATINT_TASK_ALLOCVM_REMOTE
map<wstring, uint64_t> parse_allocvm_remote(krabs::schema schema, krabs::parser parser) {
map<wstring, uint64_t> zero_map;
map<wstring, uint64_t> 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<std::uint32_t>(wsPropertyName);
break;
case TDH_INTYPE_POINTER:
allocation_fields[wsPropertyName] = parser.parse<krabs::pointer>(wsPropertyName).address;
break;
case TDH_INTYPE_UINT32:
new_event.fields[wsPropertyName] = parser.parse<uint32_t>(wsPropertyName);
break;
case TDH_INTYPE_POINTER:
new_event.fields[wsPropertyName] = parser.parse<pointer>(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<wstring, uint64_t> 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;
}
Expand All @@ -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);

Expand Down
21 changes: 11 additions & 10 deletions agent/DetectionLogic.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::wstring, uint64_t> 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<std::wstring, uint64_t> alloc_event) {
VOID allocvm_remote_signatures(GenericEvent alloc_event) {
return;
}

VOID detect_event(std::map<std::wstring, uint64_t> 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:
Expand Down
19 changes: 17 additions & 2 deletions agent/DetectionLogic.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,5 +35,20 @@ enum DETECTIONS {

extern map<wstring, uint64_t> allocation_fields;

VOID report_detection(int evtId, map<wstring, uint64_t> evt_body);
VOID detect_event(std::map<std::wstring, uint64_t> parsed_event, int eid);
class GenericEvent {
public:
uint8_t type;

map<wstring, uint64_t> 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);
1 change: 1 addition & 0 deletions agent/Helpers.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 6db0251

Please sign in to comment.