Skip to content

Instantly share code, notes, and snippets.

@EBNull EBNull/comfind.py

Created Dec 5, 2012
Embed
What would you like to do?
Given a directory, find files containing COM objects
import os
import sys
import glob
import ctypes
import logging
log = logging.getLogger(__name__)
from collections import namedtuple
def path_iter(path, root=None, dirsortkey=None):
"""Given a directory, yield all files recusively. Directories are yielded as their filename, followed by their contents."""
path = os.path.realpath(path)
if root is None:
root = os.path.realpath(path)
dir_todo = []
for fn in os.listdir(path):
full_fn = os.path.join(path, fn)
rel_fn = os.path.relpath(full_fn, root)
if os.path.isdir(full_fn):
dir_todo.append(fn)
else:
yield rel_fn
for fn in sorted(dir_todo, key=dirsortkey):
full_fn = os.path.join(path, fn)
rel_fn = os.path.relpath(full_fn, root)
yield rel_fn
for fn in path_iter(full_fn, root, dirsortkey=dirsortkey):
yield fn
def path_iter_sorted(path, key=None):
"""Given a directory yield all files recursively.
First lists all files (sorted), then for each directory lists
the directory followed by it's files (sorted)."""
if key is None:
key = lambda x: x.lower()
files = []
for fn in path_iter(path, dirsortkey=lambda x: x.lower()):
full_fn = os.path.join(path, fn)
if os.path.isdir(full_fn):
for f in sorted(files, key=key): #Sort all files in this directory and yield those
yield f
files = []
yield fn #Yield the new directory's name
else:
files.append(fn)
for f in sorted(files, key=key):
yield f #Yield the last directory
DONT_RESOLVE_DLL_REFERENCES = 0x1
#XXX: Do not use LOAD_LIBRARY_AS_IMAGE_RESOURCE or GetProcAddress will fail.
# See http://blogs.msdn.com/b/oldnewthing/archive/2005/02/14/372266.aspx
LOAD_LIBRARY_AS_IMAGE_RESOURCE = 0x20
Module = namedtuple("Module", "relname filename")
LoadedModule = namedtuple("LoadedModule", "relname filename hmod")
LoadedDLL = namedtuple("LoadedDLL", "relname filename hmod addr_DllGetClassObject")
class ComFind(object):
def __init__(self):
self.log = logging.getLogger(__name__ + '.ComFind')
self.log_search = logging.getLogger(__name__ + '.ComFind.search')
self.modules = []
self.dll_com_servers = []
def load_modules(self, dir, check_func=lambda x: True):
modules = []
dir = unicode(dir)
for fn in path_iter_sorted(dir):
full_fn = os.path.join(dir, fn)
if not os.path.isdir(full_fn):
if not check_func(full_fn):
continue
#self.log_search.debug("Trying to load %s", fn)
hmod = ctypes.windll.kernel32.LoadLibraryExW(unicode(full_fn), None, DONT_RESOLVE_DLL_REFERENCES)
if hmod:
mod = LoadedModule(fn, full_fn, hmod)
modules.append(mod)
self.log_search.debug("0x%08x: Loaded %s", mod.hmod, mod.relname)
self._find_objects(mod)
self.modules = modules
def _find_objects(self, mod):
addr = ctypes.windll.kernel32.GetProcAddress(mod.hmod, 'DllGetClassObject')
if addr:
self.dll_com_servers.append(LoadedDLL(*(mod + (addr, ))))
self.log.info("0x%08x: %s: has DllGetClassObject at 0x%08x", mod.hmod, mod.relname, addr)
def dump_objects(self, stream):
stream.write("DLLs containing COM objects:\n\n")
for dll in sorted(self.dll_com_servers, key=lambda x: x.relname.lower()):
stream.write(dll.relname + '\n')
any_with_std_ext = lambda x: x.lower().endswith(('dll', 'exe', 'ocx'))
any_with_ext = lambda x: x.rindex('.') > x.rindex(os.sep)
any = lambda x: True
def any_with_magic(filename):
with open(filename, "rb") as file:
hdr = file.read(2)
if hdr == 'MZ':
return True
return False
def main(argv):
logging.basicConfig(level=logging.NOTSET)
fn_check = any_with_magic
if '--fast' in argv:
fn_check = any_with_std_ext
if '--ext' in argv:
fn_check = any_with_ext
if '--all' in argv:
fn_check = any
cf = ComFind()
cf.load_modules(argv[1], fn_check)
cf.dump_objects(sys.stdout)
if __name__ == '__main__':
sys.exit(main(sys.argv))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.