Merge branch 'pattern_tree2'

hasherezade · Feb 8, 2024 · dd6ac87 · dd6ac87
2 parents 35fb3f4 + 40b14c7
commit dd6ac87
Show file tree

Hide file tree

Showing 3 changed files with 154 additions and 42 deletions.
diff --git a/pe_sieve_ver_short.h b/pe_sieve_ver_short.h
@@ -3,6 +3,6 @@
 #define PESIEVE_MAJOR_VERSION 0
 #define PESIEVE_MINOR_VERSION 3
 #define PESIEVE_MICRO_VERSION 8
-#define PESIEVE_PATCH_VERSION 2
+#define PESIEVE_PATCH_VERSION 3
 
-#define PESIEVE_VERSION_STR	"0.3.8.2"
+#define PESIEVE_VERSION_STR	"0.3.8.3"
diff --git a/utils/artefacts_util.cpp b/utils/artefacts_util.cpp
@@ -96,40 +96,22 @@ bool init_64_patterns(Node* rootN64)
 	return true;
 }
 
-size_t search_till_pattern(Node* rootN, const BYTE* loadedData, size_t loadedSize)
+size_t search_till_pattern(Node& rootN, const BYTE* loadedData, size_t loadedSize)
 {
-	if (rootN && loadedData) {
-		for (size_t i = 0; i < loadedSize; i++) {
-			Match m = rootN->getMatching(loadedData + i, loadedSize - i);
-			if (m.sign) return (m.offset + i);
-		}
+	Match m = pattern_tree::find_first_match(rootN, loadedData, loadedSize);
+	if (!m.sign) {
+		return CODE_PATTERN_NOT_FOUND;
 	}
-	return CODE_PATTERN_NOT_FOUND;
+	return m.offset;
 }
 
-size_t search_all_matches(Node* rootN, const BYTE* loadedData, size_t loadedSize, std::vector<Match>& matches)
-{
-	size_t found = 0;
-	if (rootN && loadedData) {
-		for (size_t i = 0; i < loadedSize; i++) {
-			Match m = rootN->getMatching(loadedData + i, loadedSize - i);
-			if (m.sign) {
-				matches.push_back(m);
-				found++;
-			}
-		}
-	}
-	return found;
-}
-
-
 size_t pesieve::util::is_32bit_code(BYTE *loadedData, size_t loadedSize)
 {
 	static Node rootN;
 	if(rootN.isEnd()) {
 		init_32_patterns(&rootN);
 	}
-	return search_till_pattern(&rootN, loadedData, loadedSize);
+	return search_till_pattern(rootN, loadedData, loadedSize);
 }
 
 size_t pesieve::util::is_64bit_code(BYTE* loadedData, size_t loadedSize)
@@ -138,7 +120,7 @@ size_t pesieve::util::is_64bit_code(BYTE* loadedData, size_t loadedSize)
 	if (rootN.isEnd()) {
 		init_64_patterns(&rootN);
 	}
-	return search_till_pattern(&rootN, loadedData, loadedSize);
+	return search_till_pattern(rootN, loadedData, loadedSize);
 }
 
 bool pesieve::util::is_code(BYTE* loadedData, size_t loadedSize)
@@ -153,7 +135,7 @@ bool pesieve::util::is_code(BYTE* loadedData, size_t loadedSize)
 		init_64_patterns(&rootN);
 	}
 
-	if ((search_till_pattern(&rootN, loadedData, loadedSize)) != CODE_PATTERN_NOT_FOUND) {
+	if ((search_till_pattern(rootN, loadedData, loadedSize)) != CODE_PATTERN_NOT_FOUND) {
 		return true;
 	}
 	return false;

diff --git a/utils/pattern_tree.h b/utils/pattern_tree.h
@@ -4,6 +4,7 @@
 #include <iostream>
 #include <string>
 #include <map>
+#include <vector>
 
 namespace pattern_tree {
 
@@ -20,16 +21,23 @@ namespace pattern_tree {
 			this->pattern_size = _pattern_size;
 		}
 
+		size_t size()
+		{
+			return pattern_size;
+		}
+
+		std::string name;
+
+	protected:
 		size_t pattern_size;
 		BYTE* pattern;
-		std::string name;
 	};
 
 	class Match
 	{
 	public:
 		Match()
-			: offset(-1), sign(nullptr)
+			: offset(0), sign(nullptr)
 		{
 		}
 
@@ -48,6 +56,60 @@ namespace pattern_tree {
 		Signature* sign;
 	};
 
+	template<class Element> class ShortList
+	{
+	public:
+		ShortList()
+			: elCount(0)
+		{
+		}
+
+		bool push_back(Element n)
+		{
+			if (elCount >= _countof(list)) {
+				return false;
+			}
+			if (find(n)) {
+				return true;
+			}
+			list[elCount] = n;
+			elCount++;
+			return true;
+		}
+
+		Element at(size_t i)
+		{
+			if (i < _countof(list)) {
+				return list[i];
+			}
+			return nullptr;
+		}
+
+		Element find(Element& searched)
+		{
+			for (size_t i = 0; i < elCount; i++) {
+				if (list[i] == searched) {
+					return list[i];
+				}
+			}
+			return nullptr;
+		}
+
+		void clear()
+		{
+			elCount = 0;
+		}
+
+		size_t size()
+		{
+			return elCount;
+		}
+
+	protected:
+		size_t elCount;
+		Element list[100];
+	};
+
 	class Node
 	{
 	public:
@@ -122,30 +184,69 @@ namespace pattern_tree {
 			}
 		}
 
-		Match getMatching(const BYTE* data, size_t data_size)
+#define SEARCH_BACK
+		size_t getMatching(const BYTE* data, size_t data_size, std::vector<Match> &matches, bool stopOnFirst)
 		{
-			Match empty;
-			Node* curr = this;
+			size_t processed = 0;
+			//
+			ShortList<Node*> level;
+			level.push_back(this);
+			ShortList<Node*> level2;
+
+			auto level1_ptr = &level;
+			auto level2_ptr = &level2;
+
 			for (size_t i = 0; i < data_size; i++)
 			{
-				if (curr->iSign()) {
-					const size_t match_start = i - curr->sign->pattern_size;
-					return Match(match_start, curr->sign);
+				processed = i; // processed bytes
+				level2_ptr->clear();
+				for (size_t k = 0; k < level1_ptr->size(); k++) {
+					Node * curr = level1_ptr->at(k);
+					if (curr->isSign()) {
+						size_t match_start = i - curr->sign->size();
+						Match m(match_start, curr->sign);
+						matches.push_back(m);
+						if (stopOnFirst) {
+							return match_start;
+						}
+					}
+					Node* prev = curr;
+					curr = prev->getNode(data[i]);
+					if (curr) {
+						level2_ptr->push_back(curr);
+					}
+#ifdef SEARCH_BACK
+					if (prev != this) {
+						// the current value may also be a beginning of a new pattern:
+						Node* start = this->getNode(data[i]);
+						if (start) {
+							level2_ptr->push_back(start);
+						}
+					}
+#endif
+				}
+				if (!level2_ptr->size()) {
+#ifdef SEARCH_BACK
+					// restart search from the beginning
+					level2_ptr->push_back(this);
+#else
+					return results;
+#endif //SEARCH_BACK
 				}
-				if (curr->isEnd()) return empty;
-				BYTE val = data[i];
-				curr = curr->getNode(val);
-				if (!curr) return empty;
+				//swap:
+				auto tmp = level1_ptr;
+				level1_ptr = level2_ptr;
+				level2_ptr = tmp;
 			}
-			return empty;
+			return processed;
 		}
 
 		bool isEnd()
 		{
 			return this->immediates.size() ? false : true;
 		}
 
-		bool iSign()
+		bool isSign()
 		{
 			return sign ? true : false;
 		}
@@ -157,4 +258,33 @@ namespace pattern_tree {
 		std::map<BYTE, Node*> immediates;
 	};
 
+
+	inline size_t find_all_matches(Node& rootN, const BYTE* loadedData, size_t loadedSize, std::vector<Match> &allMatches)
+	{
+		if (!loadedData || !loadedSize) {
+			return 0;
+		}
+		size_t counter = 0;
+		rootN.getMatching(loadedData, loadedSize, allMatches, false);
+		if (allMatches.size()) {
+			counter += allMatches.size();
+		}
+		return counter;
+	}
+
+	inline Match find_first_match(Node& rootN, const BYTE* loadedData, size_t loadedSize)
+	{
+		Match empty;
+		if (!loadedData || !loadedSize) {
+			return empty;
+		}
+		std::vector<Match> allMatches;
+		rootN.getMatching(loadedData, loadedSize, allMatches, true);
+		if (allMatches.size()) {
+			auto itr = allMatches.begin();
+			return *itr;
+		}
+		return empty;
+	}
+
 }; //namespace pattern_tree