11import datetime
22import logging
33import string
4+ from itertools import chain
45from typing import Dict , Iterable , List
56
67from volatility3 .framework import constants , exceptions
@@ -147,30 +148,33 @@ def _check_csrss_handles(
147148 ) -> Dict [int , extensions .EPROCESS ]:
148149 ret : List [extensions .EPROCESS ] = []
149150
151+ handles_plugin = handles .Handles (
152+ context = self .context , config_path = self .config_path
153+ )
154+
155+ type_map = handles_plugin .get_type_map (self .context , layer_name , symbol_table )
156+
157+ cookie = handles_plugin .find_cookie (
158+ context = self .context ,
159+ layer_name = layer_name ,
160+ symbol_table = symbol_table ,
161+ )
162+
150163 for p in tasks :
151164 name = self ._proc_name_to_string (p )
152- if name == "csrss.exe" :
153- try :
154- if p .has_member ("ObjectTable" ):
155- handles_plugin = handles .Handles (
156- context = self .context , config_path = self .config_path
157- )
158- hndls = list (handles_plugin .handles (p .ObjectTable ))
159- for h in hndls :
160- if (
161- h .get_object_type (
162- handles_plugin .get_type_map (
163- self .context , layer_name , symbol_table
164- )
165- )
166- == "Process"
167- ):
168- ret .append (h .Body .cast ("_EPROCESS" ))
169-
170- except exceptions .InvalidAddressException :
171- vollog .log (
172- constants .LOGLEVEL_VVV , "Cannot access eprocess object table"
173- )
165+ if name != "csrss.exe" :
166+ continue
167+
168+ try :
169+ ret += [
170+ handle .Body .cast ("_EPROCESS" )
171+ for handle in handles_plugin .handles (p .ObjectTable )
172+ if handle .get_object_type (type_map , cookie ) == "Process"
173+ ]
174+ except exceptions .InvalidAddressException :
175+ vollog .log (
176+ constants .LOGLEVEL_VVV , "Cannot access eprocess object table"
177+ )
174178
175179 return self ._proc_list_to_dict (ret )
176180
@@ -187,7 +191,7 @@ def _generator(self):
187191 )
188192
189193 # get processes from each source
190- processes = {}
194+ processes : Dict [ str , Dict [ int , extensions . EPROCESS ]] = {}
191195
192196 processes ["pslist" ] = self ._check_pslist (kdbg_list_processes )
193197 processes ["psscan" ] = self ._check_psscan (layer_name , symbol_table )
@@ -196,27 +200,20 @@ def _generator(self):
196200 kdbg_list_processes , layer_name , symbol_table
197201 )
198202
199- # print results
200-
201- # list of lists of offsets
202- offsets = [list (processes [source ].keys ()) for source in processes ]
203-
204- # flatten to one list
205- offsets = sum (offsets , [])
206-
207- # remove duplicates
208- offsets = set (offsets )
203+ # Unique set of all offsets from all sources
204+ offsets = set (chain (* (mapping .keys () for mapping in processes .values ())))
209205
210206 for offset in offsets :
211- proc = None
207+ # We know there will be at least one process mapped to each offset
208+ proc : extensions .EPROCESS = next (
209+ mapping [offset ] for mapping in processes .values () if offset in mapping
210+ )
212211
213212 in_sources = {src : False for src in processes }
214213
215- for source in processes :
216- if offset in processes [ source ] :
214+ for source , process_mapping in processes . items () :
215+ if offset in process_mapping :
217216 in_sources [source ] = True
218- if not proc :
219- proc = processes [source ][offset ]
220217
221218 pid = proc .UniqueProcessId
222219 name = self ._proc_name_to_string (proc )
0 commit comments