From e0b7f0b5f0be2e6705c4ecee5109959d7b6f2348 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Sun, 22 May 2016 15:28:12 +0100
Subject: [PATCH] Add in helper functions for symbolspaces/symboltables to
 locate specific symbols.

---
 volatility/framework/interfaces/symbols.py | 29 +++++++++++++++++++---
 volatility/framework/symbols/__init__.py   | 12 +++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/volatility/framework/interfaces/symbols.py b/volatility/framework/interfaces/symbols.py
index 9c0d562d0..e8ab029cf 100644
--- a/volatility/framework/interfaces/symbols.py
+++ b/volatility/framework/interfaces/symbols.py
@@ -3,6 +3,7 @@ Created on 4 May 2013
 
 @author: mike
 """
+import bisect
 
 from volatility.framework import validity, exceptions, constants
 from volatility.framework.interfaces import configuration
@@ -55,10 +56,6 @@ class SymbolTableInterface(validity.ValidityRoutines):
         """
         raise NotImplementedError("Abstract property get_symbol not implemented by subclass.")
 
-    def get_symbol_type(self, name):
-        """Resolves a symbol name into a symbol and then resolves the symbol's type"""
-        return self.get_type(self.get_symbol(name).type_name)
-
     @property
     def symbols(self):
         """Returns an iterator of the Symbols"""
@@ -111,6 +108,30 @@ class SymbolTableInterface(validity.ValidityRoutines):
         """Removes the associated class override for a specific Symbol type"""
         raise NotImplementedError("Abstract method del_type_class not implemented yet.")
 
+    # ## Convenience functions for location symbols
+
+    def get_symbol_type(self, name):
+        """Resolves a symbol name into a symbol and then resolves the symbol's type"""
+        return self.get_type(self.get_symbol(name).type_name)
+
+    def get_symbols_by_type(self, type_name):
+        """Returns the name of all symbols in this table that have type matching type_name"""
+        for symbol in self.symbols:
+            # This allows for searching with and without the table name (in case multiple tables contain
+            # the same symbol name and we've not specifically been told which one)
+            if symbol.type_name == type_name or (symbol.type_name.endswith(constants.BANG + type_name)):
+                yield symbol.name
+
+    def get_symbols_by_location(self, offset):
+        """Returns the name of all symbols in this table that have type matching type_name"""
+        sort_symbols = [(s.offset, s) for s in sorted(self.symbols, key = lambda x: x.offset)]
+        result = bisect.bisect_left(sort_symbols, offset)
+        if result == len(sort_symbols):
+            raise StopIteration
+        closest_symbol = sort_symbols[result][1]
+        if closest_symbol.offset == offset:
+            yield closest_symbol.name
+
 
 class NativeTableInterface(SymbolTableInterface):
     """Class to distinguish NativeSymbolLists from other symbol lists"""
diff --git a/volatility/framework/symbols/__init__.py b/volatility/framework/symbols/__init__.py
index 3e21bc61c..2b4c6f0f3 100644
--- a/volatility/framework/symbols/__init__.py
+++ b/volatility/framework/symbols/__init__.py
@@ -33,6 +33,18 @@ class SymbolSpace(collections.abc.Mapping):
         # Permanently cache all resolved symbols
         self._resolved = {}
 
+    def get_symbols_by_type(self, type_name):
+        """Returns all symbols based """
+        for table in self._dict.keys():
+            for symbol_name in self._dict[table].get_symbols_by_type(type_name):
+                yield table + constants.BANG + symbol_name
+
+    def get_symbols_by_location(self, offset):
+        """Returns all symbols that exist at a specific relative offset"""
+        for table in self._dict.values():
+            for symbol_name in self._dict[table].get_symbols_by_location(offset = offset):
+                yield table + constants.BANG + symbol_name
+
     @property
     def natives(self):
         """Returns the native_types for this symbol space"""