Merge branch 'watch_file'

Author: hasherezade

driver_comm.cpp

@@ -3,7 +3,12 @@
 #include <iostream>
 
 struct ProcessData {
-	unsigned long Id;
+	DWORD Id;
+};
+
+struct ProcessDataEx {
+	DWORD Id;
+	LONGLONG fileId;
 };
 
 #define DRIVER_PATH  L"\\\\.\\MalUnpackCompanion"
@@ -124,9 +129,26 @@ bool driver::fetch_watched_files(DWORD startPID, LONGLONG out_buffer[], size_t o
 }
 
 
-bool driver::watch_pid(DWORD pid)
+bool driver::watch_pid(DWORD pid, ULONGLONG fileId)
 {
-	return driver::request_action_on_pid(IOCTL_MUNPACK_COMPANION_ADD_TO_WATCHED, pid);
+	if (!pid) {
+		return false;
+	}
+	HANDLE hDevice = CreateFileW(DRIVER_PATH, GENERIC_WRITE, FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, 0, nullptr);
+	if (hDevice == INVALID_HANDLE_VALUE) {
+		std::cerr << "Failed to open device" << std::endl;
+		return false;
+	}
+
+	ProcessDataEx data = { 0 };
+	data.Id = pid;
+	data.fileId = fileId;
+
+	BOOL success = FALSE;
+	DWORD returned = 0;
+	success = DeviceIoControl(hDevice, IOCTL_MUNPACK_COMPANION_ADD_TO_WATCHED, &data, sizeof(data), nullptr, 0, &returned, nullptr);
+	CloseHandle(hDevice);
+	return success == TRUE ? true : false;
 }
 
 bool driver::kill_watched_pid(DWORD pid)

driver_comm.h

@@ -6,7 +6,7 @@ namespace driver {
 
 	bool is_ready();
 
-	bool watch_pid(DWORD pid);
+	bool watch_pid(DWORD pid, ULONGLONG fileId);
 
 	bool kill_watched_pid(DWORD pid);
 

file_util.cpp

@@ -53,9 +53,79 @@ namespace file_util {
 		return (got_len != 0) ? true: false;
 	}
 
+	NTSTATUS FetchFileId(HANDLE hFile, LONGLONG& FileId)
+	{
+		FileId = FILE_INVALID_FILE_ID;
+
+		if (!hFile) {
+			return STATUS_INVALID_PARAMETER;
+		}
+
+		NTSTATUS status = STATUS_UNSUCCESSFUL;
+		__try
+		{
+			IO_STATUS_BLOCK ioStatusBlock;
+			FILE_INTERNAL_INFORMATION fileIdInfo;
+			status = ZwQueryInformationFile(
+				hFile,
+				&ioStatusBlock,
+				&fileIdInfo,
+				sizeof(fileIdInfo),
+				FileInternalInformation
+			);
+			if (NT_SUCCESS(status)) {
+				FileId = fileIdInfo.IndexNumber.QuadPart;
+			}
+		}
+		__except (EXCEPTION_EXECUTE_HANDLER)
+		{
+			status = STATUS_UNSUCCESSFUL;
+		}
+		return status;
+	}
 };
 
 
+std::wstring file_util::get_file_path(const char* file_name)
+{
+	wchar_t full_path[MAX_PATH] = { 0 };
+
+	HANDLE hFile = CreateFileA(file_name, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
+	DWORD got_len = GetFinalPathNameByHandleW(hFile, full_path, MAX_PATH, VOLUME_NAME_DOS);
+	CloseHandle(hFile);
+
+	if (!got_len) {
+		return L"";
+	}
+	
+	wchar_t* prefix = L"\\\\?\\";
+	size_t prefix_len = wcslen(prefix);
+
+	if (got_len < prefix_len) {
+		return full_path;
+	}
+	if (wcsncmp(prefix, full_path, prefix_len) == 0) {
+		// skip the prefix
+		return full_path + prefix_len;
+	}
+	return full_path;
+}
+
+ULONGLONG file_util::get_file_id(const char* img_path)
+{
+	HANDLE file = CreateFileA(img_path, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
+
+	LONGLONG FileId = FILE_INVALID_FILE_ID;
+	NTSTATUS status = FetchFileId(file, FileId);
+
+	CloseHandle(file);
+
+	if (NT_SUCCESS(status)) {
+		return FileId;
+	}
+	return FILE_INVALID_FILE_ID;
+}
+
 size_t file_util::file_ids_to_names(std::set<LONGLONG>& filesIds, std::set<std::wstring>& names)
 {
 	FILE_ID_DESCRIPTOR FileDesc = { 0 };

file_util.h

@@ -7,6 +7,10 @@
 
 namespace file_util {
 
+	std::wstring get_file_path(const char* filename);
+
+	ULONGLONG get_file_id(const char* img_path);
+
 	size_t file_ids_to_names(std::set<LONGLONG>& filesIds, std::set<std::wstring> &names);
 
 	size_t delete_dropped_files(std::set<std::wstring>& names);

main.cpp

@@ -17,6 +17,7 @@
 #include "process_util.h"
 #include "util.h"
 #include "version.h"
+#include "file_util.h"
 
 #define WAIT_FOR_PROCESS_TIMEOUT 5000
 
@@ -46,6 +47,30 @@ void init_defaults(t_params_struct &params)
 #endif
 }
 
+bool get_watched_file_id(const t_params_struct& params, ULONGLONG &file_id)
+{
+    bool is_diffrent = false;
+    file_id = FILE_INVALID_FILE_ID;
+
+    if (strnlen(params.img_path, MAX_PATH) > 0
+        && strncmp(params.img_path, params.exe_path, MAX_PATH) != 0)
+    {
+        file_id = file_util::get_file_id(params.img_path);
+        is_diffrent = true;
+    }
+    if (file_id == FILE_INVALID_FILE_ID) {
+        file_id = file_util::get_file_id(params.exe_path);
+        is_diffrent = false;
+    }
+    if (is_diffrent) {
+        std::cout << "[*] Watch respawns from the IMG: " << params.img_path << "\n";
+    }
+    else {
+        std::cout << "[*] Watch respawns from main EXE file: " << params.exe_path << "\n";
+    }
+    return is_diffrent;
+}
+
 int main(int argc, char* argv[])
 {
     UnpackParams uParams(VERSION);
@@ -66,10 +91,10 @@ int main(int argc, char* argv[])
     }
     uParams.fillStruct(params);
     if (params.hh_args.pesieve_args.use_cache) {
-        std::cerr << "[*] Cache is Enabled!" << std::endl;
+        std::cout << "[*] Cache is Enabled!" << std::endl;
     }
     else {
-        std::cerr << "[*] Cache is Disabled!" << std::endl;
+        std::cout << "[*] Cache is Disabled!" << std::endl;
     }
     params.hh_args.kill_suspicious = true;
     // if the timeout was chosen as the trigger, don't interfere in the process:
@@ -86,7 +111,7 @@ int main(int argc, char* argv[])
     std::cout << "Starting the process: " << params.exe_path << std::endl;
     std::cout << "With commandline: \"" << params.exe_cmd << "\"" << std::endl;
 
-    char* file_name = get_file_name(params.exe_path);
+    const char* file_name = get_file_name(params.exe_path);
     std::cout << "Exe name: " << file_name << std::endl;
 
     DWORD timeout = params.timeout;
@@ -98,13 +123,17 @@ int main(int argc, char* argv[])
 
     t_pesieve_res ret_code = PESIEVE_ERROR;
     const DWORD flags = DETACHED_PROCESS | CREATE_NO_WINDOW;
-    HANDLE proc = make_new_process(params.exe_path, params.exe_cmd, flags);
+
+    ULONGLONG file_id;
+    bool is_img_diff = get_watched_file_id(params, file_id);
+    HANDLE proc = make_new_process(params.exe_path, params.exe_cmd, flags, file_id);
     if (!proc) {
         std::cerr << "Could not start the process!" << std::endl;
         return ret_code;
     }
-
-    params.hh_args.pname = file_name;
+    params.hh_args.is_main_module = is_img_diff ? false : true;
+    params.hh_args.module_path = (is_img_diff) ? file_util::get_file_path(params.img_path) : file_util::get_file_path(params.exe_path);
+    std::wcout << "Module Path retrieved: " << params.hh_args.module_path << "\n";
     params.hh_args.start_pid = GetProcessId(proc);
 
     std::string out_dir = make_dir_name(root_dir, time(NULL));
@@ -150,7 +179,7 @@ int main(int argc, char* argv[])
     finalStats.scanTime = GetTickCount64() - start_tick;
 
     //this works only with the companion driver:
-    if (scanner.collectDroppedFiles()) {
+    if (scanner.collectDroppedFiles(file_id)) {
         std::cout << "The process dropped some files!\n";
     }
 

params.h

@@ -9,6 +9,7 @@ using namespace paramkit;
 #define DEFAULT_TIMEOUT 1000
 
 #define PARAM_EXE "exe"
+#define PARAM_IMG "img"
 #define PARAM_CMD "cmd"
 #define PARAM_TIMEOUT "timeout"
 #define PARAM_OUT_DIR "dir"
@@ -32,6 +33,7 @@ typedef struct {
     char exe_path[MAX_PATH];
     char exe_cmd[MAX_PATH];
     char out_dir[MAX_PATH];
+    char img_path[MAX_PATH];
     DWORD timeout;
     t_term_trigger trigger;
     UnpackScanner::t_unp_params hh_args;
@@ -94,6 +96,10 @@ class UnpackParams : public Params
         this->addParam(new StringParam(PARAM_EXE, true));
         this->setInfo(PARAM_EXE, "Input exe (to be run)");
 
+        this->addParam(new StringParam(PARAM_IMG, false));
+        this->setInfo(PARAM_IMG, "Path to the image from which the malware was run (may be a DLL or EXE).\n\t   DEFAULT: same as /exe", 
+            "\t   This allows to follow processes respawned from the given image.");
+
         this->addParam(new StringParam(PARAM_CMD, false));
         this->setInfo(PARAM_CMD, "Commandline arguments for the input exe");
 
@@ -197,6 +203,7 @@ class UnpackParams : public Params
     void fillStruct(t_params_struct &ps)
     {
         copyCStr<StringParam>(PARAM_EXE, ps.exe_path, sizeof(ps.exe_path));
+        copyCStr<StringParam>(PARAM_IMG, ps.img_path, sizeof(ps.img_path));
         copyCStr<StringParam>(PARAM_CMD, ps.exe_cmd, sizeof(ps.exe_cmd));
         copyCStr<StringParam>(PARAM_OUT_DIR, ps.out_dir, sizeof(ps.out_dir));