Given a directory, find files containing COM objects

comfind.py

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))