text: #include "code_scanner.h"
  
  #include <peconv.h>
  
  #include "patch_analyzer.h"
  #include "../utils/artefacts_util.h"
  //---
  #include <iostream>
  
  using namespace pesieve;
  using namespace pesieve::util;
  
  size_t pesieve::CodeScanReport::generateTags(std::string reportPath)
  {
  	if (patchesList.size() == 0) {
  		return 0;
  	}
  	std::ofstream patch_report;
  	patch_report.open(reportPath);
  	if (patch_report.is_open() == false) {
  		return 0;
  	}
  	size_t patches = patchesList.toTAGs(patch_report, ';');
  	if (patch_report.is_open()) {
  		patch_report.close();
  	}
  	return patches;
  }
  //---
  
  bool pesieve::CodeScanner::clearIAT(PeSection &originalSec, PeSection &remoteSec)
  {
  	// collect IAT fields:
  	std::set<DWORD> impThunkRVAs;
  	moduleData.loadImportThunks(impThunkRVAs);
  	if (impThunkRVAs.size() == 0) {
  		return false;
  	}
  
  	const size_t thunk_size = moduleData.is64bit() ? sizeof(ULONGLONG) : sizeof(DWORD);
  	std::set<DWORD>::iterator itr;
  	for (itr = impThunkRVAs.begin(); itr != impThunkRVAs.end(); ++itr) {
  		const DWORD iat_field = *itr;
  		// clear fields one by one:
  		if (originalSec.isContained(iat_field, thunk_size)) {
  			const DWORD offset = iat_field - originalSec.rva;
  			memset(originalSec.loadedSection + offset, 0, thunk_size);
  			memset(remoteSec.loadedSection + offset, 0, thunk_size);
  		}
  	}
  	return true;
  }
  
  bool pesieve::CodeScanner::clearLoadConfig(PeSection &originalSec, PeSection &remoteSec)
  {
  	// check if the Guard flag is enabled:
  	WORD charact = peconv::get_dll_characteristics(moduleData.original_module);
  	if ((charact & 0x4000) == 0) {
  		return false; //no guard flag
  	}
  	BYTE *ldconf_ptr = peconv::get_load_config_ptr(moduleData.original_module, moduleData.original_size);
  	if (!ldconf_ptr) return false;
  
  	peconv::t_load_config_ver ver = peconv::get_load_config_version(moduleData.original_module, moduleData.original_size, ldconf_ptr);
  	if (ver != peconv::LOAD_CONFIG_W8_VER && ver != peconv::LOAD_CONFIG_W10_VER) {
  		return false; // nothing to cleanup
  	}
  	ULONGLONG cflag_va = 0;
  	size_t field_size = 0;
  	if (this->moduleData.is64bit()) {
  		peconv::IMAGE_LOAD_CONFIG_DIR64_W8* ldc = (peconv::IMAGE_LOAD_CONFIG_DIR64_W8*) ldconf_ptr;
  		cflag_va = ldc->GuardCFCheckFunctionPointer;
  		field_size = sizeof(ULONGLONG);
  	}
  	else {
  		peconv::IMAGE_LOAD_CONFIG_DIR32_W8* ldc = (peconv::IMAGE_LOAD_CONFIG_DIR32_W8*) ldconf_ptr;
  		cflag_va = ldc->GuardCFCheckFunctionPointer;
  		field_size = sizeof(DWORD);
  	}
  	if (cflag_va == 0) return false;
  
  	const ULONGLONG module_base = (ULONG_PTR)moduleData.moduleHandle;
  	const ULONGLONG cflag_rva = cflag_va - module_base;
  	if (!originalSec.isContained(cflag_rva, field_size)) {
  		return false;
  	}
  	//clear the field:
  	size_t sec_offset = size_t(cflag_rva - originalSec.rva);
  	memset(originalSec.loadedSection + sec_offset, 0, field_size);
  	memset(remoteSec.loadedSection + sec_offset, 0, field_size);
  	return true;
  }
  
  bool pesieve::CodeScanner::clearExports(PeSection &originalSec, PeSection &remoteSec)
  {
  	IMAGE_DATA_DIRECTORY* dir = peconv::get_directory_entry(moduleData.original_module, IMAGE_DIRECTORY_ENTRY_EXPORT);
  	if (!dir) {
  		return false;
  	}
  	DWORD iat_rva = dir->VirtualAddress;
  	DWORD iat_size = dir->Size;
  
  	if (originalSec.isContained(iat_rva, iat_size))
  	{
  #ifdef _DEBUG
  		std::cout << "Exports are in the Code section!" << std::endl;
  #endif
  		DWORD offset = iat_rva - originalSec.rva;
  		IMAGE_EXPORT_DIRECTORY *exports = (IMAGE_EXPORT_DIRECTORY*) ((ULONGLONG)originalSec.loadedSection + offset);
  		if (!peconv::validate_ptr(originalSec.loadedSection, originalSec.loadedSize, exports, sizeof(IMAGE_EXPORT_DIRECTORY))) {
  			return false;
  		}
  		DWORD functions_offset = exports->AddressOfFunctions - originalSec.rva;
  		DWORD functions_count = exports->NumberOfFunctions;
  
  		const size_t func_area_size = functions_count * sizeof(DWORD);
  		if (!peconv::validate_ptr(originalSec.loadedSection, originalSec.loadedSize, 
  			originalSec.loadedSection + functions_offset, 
  			func_area_size))
  		{
  			return false;
  		}
  		memset(originalSec.loadedSection + functions_offset, 0, func_area_size);
  		memset(remoteSec.loadedSection + functions_offset, 0, func_area_size);
  	}
  	return true;
  }
  
  size_t pesieve::CodeScanner::collectPatches(DWORD section_rva, PBYTE orig_code, PBYTE patched_code, size_t code_size, OUT PatchList &patchesList)
  {
  	PatchAnalyzer analyzer(moduleData, section_rva, patched_code, code_size);
  	PatchList::Patch *currPatch = nullptr;
  
  	for (DWORD i = 0; i < (DWORD) code_size; i++) {
  		if (orig_code[i] == patched_code[i]) {
  			if (currPatch != nullptr) {
  				// close the patch
  				currPatch->setEnd(section_rva + i);
  				currPatch = nullptr;
  			}
  			continue;
  		}
  		if (currPatch == nullptr) {
  			//open a new patch
  			currPatch = new PatchList::Patch(moduleData.moduleHandle, patchesList.size(), (DWORD) section_rva + i);
  			patchesList.insert(currPatch);
  			DWORD parsed_size = (DWORD) analyzer.analyze(*currPatch);
  			if (parsed_size > 0) {
  				currPatch->setEnd(section_rva + i + parsed_size);
  				currPatch = nullptr; // close this patch
  				i += (parsed_size - 1); //substract 1 because of i++ executed after continue
  				continue;
  			}
  		}
  	}
  	// if there is still unclosed patch, close it now:
  	if (currPatch != nullptr) {
  		//this happens if the patch lasts till the end of code, so, its end is the end of code
  		currPatch->setEnd(section_rva + (DWORD) code_size);
  		currPatch = nullptr;
  	}
  	return patchesList.size();
  }
  
  namespace pesieve {
  	inline BYTE* first_different(const BYTE *buf_ptr, size_t bif_size, const BYTE padding)
  	{
  		for (size_t i = 0; i < bif_size; i++) {
  			if (buf_ptr[i] != padding) {
  				return (BYTE*)(buf_ptr + i);
  			}
  		}
  		return nullptr;
  	}
  };
  
  CodeScanReport::t_section_status pesieve::CodeScanner::scanSection(PeSection &originalSec, PeSection &remoteSec, OUT PatchList &patchesList)
  {
  	if (!originalSec.isInitialized() || !remoteSec.isInitialized()) {
  		return CodeScanReport::SECTION_SCAN_ERR;
  	}
  	clearIAT(originalSec, remoteSec);
  	clearExports(originalSec, remoteSec);
  	clearLoadConfig(originalSec, remoteSec);
  	//TODO: handle sections that have inside Delayed Imports (they give false positives)
  
  	const size_t smaller_size = originalSec.loadedSize > remoteSec.loadedSize ? remoteSec.loadedSize : originalSec.loadedSize;
  #ifdef _DEBUG
  	std::cout << "Code RVA: " 
  		<< std::hex << originalSec.rva
  		<< " to "
  		<< std::hex << originalSec.loadedSize
  		<< std::endl;
  #endif
  	//check if the code of the loaded module is same as the code of the module on the disk:
  	int res = memcmp(remoteSec.loadedSection, originalSec.loadedSection, smaller_size);
  
  	if ((originalSec.rawSize == 0 || peconv::is_padding(originalSec.loadedSection, smaller_size, 0))
  		&& !peconv::is_padding(remoteSec.loadedSection, smaller_size, 0))
  	{
  		return pesieve::CodeScanReport::SECTION_UNPACKED; // modified
  	}
  
  	if (res != 0) {
  		collectPatches(originalSec.rva, originalSec.loadedSection, remoteSec.loadedSection, smaller_size, patchesList);
  	}
  
  	if (remoteSec.loadedSize > originalSec.loadedSize) {
  		
  		const size_t diff = remoteSec.loadedSize - originalSec.loadedSize;
  		const BYTE *diff_bgn = remoteSec.loadedSection + originalSec.loadedSize;
  		
  		BYTE *not_padding = first_different(diff_bgn, diff, 0);
  		if (not_padding) {
  			const size_t found_offset = not_padding - remoteSec.loadedSection;
  			const DWORD found_rva = remoteSec.rva + found_offset;
  			PatchList::Patch* currPatch = new PatchList::Patch(moduleData.moduleHandle, patchesList.size(), found_rva);
  			currPatch->setEnd((DWORD)remoteSec.rva + remoteSec.loadedSize);
  			patchesList.insert(currPatch);
  		}
  	}
  	if (patchesList.size()) {
  		return pesieve::CodeScanReport::SECTION_PATCHED; // modified
  	}
  	if (res == 0) {
  		return pesieve::CodeScanReport::SECTION_NOT_MODIFIED; //not modified
  	}
  	return pesieve::CodeScanReport::SECTION_UNPACKED; // modified
  }
  
  size_t pesieve::CodeScanner::collectExecutableSections(RemoteModuleData &_remoteModData, std::map<size_t, PeSection*> &sections, CodeScanReport &my_report)
  {
  	size_t initial_size = sections.size();
  	const size_t sec_count = peconv::get_sections_count(_remoteModData.headerBuffer, _remoteModData.getHeaderSize());
  	for (size_t i = 0; i < sec_count; i++) {
  		PIMAGE_SECTION_HEADER section_hdr = peconv::get_section_hdr(_remoteModData.headerBuffer, _remoteModData.getHeaderSize(), i);
  		if (section_hdr == nullptr) {
  			continue;
  		}
  
  		const bool is_entry = _remoteModData.isSectionEntry(i);
  
  		if (!is_entry // entry section may be set as non executable, but it will still be executed
  			&& !(section_hdr->Characteristics & IMAGE_SCN_MEM_EXECUTE)
  			&& !_remoteModData.isSectionExecutable(i, this->isScanData, this->isScanInaccessible))
  		{
  			//not executable, skip it
  			continue;
  		}
  
  		//get the code section from the remote module:
  		PeSection *remoteSec = new PeSection(_remoteModData, i);
  		if (!remoteSec->isInitialized()) {
  			my_report.sectionToResult[i] = CodeScanReport::SECTION_SCAN_ERR;
  			continue;
  		}
  		if (is_entry // always scan section containing Entry Point
  			|| is_code(remoteSec->loadedSection, remoteSec->loadedSize))
  		{
  			sections[i] = remoteSec;
  		}
  	}
  	//corner case: PEs without sections
  	if (sec_count == 0) {
  		PeSection *remoteSec = new PeSection(_remoteModData, 0);
  		if (remoteSec->isInitialized()) {
  			sections[0] = remoteSec;
  		}
  	}
  	return sections.size() - initial_size;
  }
  
  void pesieve::CodeScanner::freeExecutableSections(std::map<size_t, PeSection*> &sections)
  {
  	std::map<size_t, PeSection*>::iterator itr;
  	for (itr = sections.begin(); itr != sections.end(); ++itr) {
  		PeSection *sec = itr->second;
  		delete sec;
  	}
  	sections.clear();
  }
  
  t_scan_status pesieve::CodeScanner::scanUsingBase(
  	IN ULONGLONG load_base,
  	IN std::map<size_t, PeSection*> &remote_code,
  	OUT std::map<DWORD, CodeScanReport::t_section_status> &sectionToResult,
  	OUT PatchList &patchesList)
  {
  	t_scan_status last_res = SCAN_NOT_SUSPICIOUS;
  
  	// before scanning, ensure that the original module is relocated to the base where it was loaded
  	if (!moduleData.relocateToBase(load_base)) {
  		return SCAN_ERROR;
  	}
  
  	size_t errors = 0;
  	size_t modified = 0;
  	std::map<size_t, PeSection*>::iterator itr;
  
  	for (itr = remote_code.begin(); itr != remote_code.end(); ++itr) {
  		size_t sec_indx = itr->first;
  		PeSection *remoteSec = itr->second;
  
  		PeSection originalSec(moduleData, sec_indx);
  
  		CodeScanReport::t_section_status sec_status = scanSection(originalSec, *remoteSec, patchesList);
  		sectionToResult[originalSec.rva] = sec_status; //save the status for the section
  
  		if (sec_status == pesieve::CodeScanReport::SECTION_SCAN_ERR) errors++;
  		else if (sec_status != pesieve::CodeScanReport::SECTION_NOT_MODIFIED) {
  			modified++;
  		}
  	}
  
  	if (modified > 0) {
  		last_res = SCAN_SUSPICIOUS; //the highest priority for modified
  	}
  	else if (errors > 0) {
  		last_res = SCAN_ERROR;
  	}
  	return last_res;
  }
  
  pesieve::CodeScanReport* pesieve::CodeScanner::scanRemote()
  {
  	if (!moduleData.isInitialized()) {
  		std::cerr << "[-] Module not initialized" << std::endl;
  		return nullptr;
  	}
  	if (!remoteModData.isInitialized()) {
  		std::cerr << "[-] Failed to read the module header" << std::endl;
  		return nullptr;
  	}
  	CodeScanReport *my_report = new CodeScanReport(this->processHandle, moduleData.moduleHandle, remoteModData.getModuleSize());
  	if (!my_report) return nullptr; //this should not happen...
  
  	my_report->isDotNetModule = moduleData.isDotNet();
  	
  	t_scan_status last_res = SCAN_NOT_SUSPICIOUS;
  	std::map<size_t, PeSection*> remote_code;
  
  	if (!collectExecutableSections(remoteModData, remote_code, *my_report)) {
  		my_report->status = last_res;
  		if (my_report->countInaccessibleSections() > 0) {
  			my_report->status = SCAN_ERROR;
  		}
  		return my_report;
  	}
  	ULONGLONG load_base = (ULONGLONG)moduleData.moduleHandle;
  	ULONGLONG hdr_base = remoteModData.getHdrImageBase();
  
  	my_report->relocBase = load_base;
  	last_res = scanUsingBase(load_base, remote_code, my_report->sectionToResult, my_report->patchesList);
  	
  	if (load_base != hdr_base && my_report->patchesList.size() > 0) {
  #ifdef _DEBUG
  		std::cout << "[WARNING] Load Base: " << std::hex << load_base << " is different than the Hdr Base: " << hdr_base << "\n";
  #endif
  		PatchList list2;
  		std::map<DWORD, CodeScanReport::t_section_status> section_to_result;
  		t_scan_status scan_res2 = scanUsingBase(hdr_base, remote_code, section_to_result, list2);
  		if (list2.size() < my_report->patchesList.size()) {
  			my_report->relocBase = hdr_base;
  			my_report->patchesList = list2;
  			my_report->sectionToResult = section_to_result;
  			last_res = scan_res2;
  		}
  #ifdef _DEBUG
  		std::cout << "Using patches list for the base: " << my_report->relocBase << " list size: " << my_report->patchesList.size() << "\n";
  #endif
  	}
  
  	this->freeExecutableSections(remote_code);
  	//post-process collected patches:
  	postProcessScan(*my_report);
  
  	my_report->status = last_res;
  	return my_report; // last result
  }
  
  bool pesieve::CodeScanner::postProcessScan(IN OUT CodeScanReport &report)
  {
  	// we need only exports from the current module, not the global mapping
  	if (report.patchesList.size() == 0) {
  		return false;
  	}
  	peconv::ExportsMapper local_mapper;
  	local_mapper.add_to_lookup(moduleData.szModName, (HMODULE) moduleData.original_module, (ULONGLONG) moduleData.moduleHandle);
  	report.patchesList.checkForHookedExports(local_mapper);
  	return true;
  }
  