Repository: mandiant/capa
File: capa/rules/__init__.py (Python)

Code before:
        return f"invalid rule set: {self.msg}"

    def __repr__(self):
        return str(self)


class ComType(Enum):
    CLASS = "class"
    INTERFACE = "interface"


# COM data source https://github.com/stevemk14ebr/COM-Code-Helper/tree/master
VALID_COM_TYPES = {
    ComType.CLASS: {"db_path": "assets/classes.json.gz", "prefix": "CLSID_"},
    ComType.INTERFACE: {"db_path": "assets/interfaces.json.gz", "prefix": "IID_"},
}

com_db_cache: Dict[ComType, Dict[str, List[str]]] = {}


def load_com_database(com_type: ComType) -> Dict[str, List[str]]:
    com_db_path: Path = capa.main.get_default_root() / VALID_COM_TYPES[com_type]["db_path"]

    if com_type in com_db_cache:
        # If the com database is already in the cache, return it
        return com_db_cache[com_type]

    if not com_db_path.exists():
        raise IOError(f"COM database path '{com_db_path}' does not exist or cannot be accessed")

    try:
        with gzip.open(com_db_path, "rb") as gzfile:
            com_db: Dict[str, List[str]] = json.loads(gzfile.read().decode("utf-8"))
            com_db_cache[com_type] = com_db  # Cache the loaded database
            return com_db
    except Exception as e:
        raise IOError(f"Error loading COM database from '{com_db_path}'") from e


def translate_com_feature(com_name: str, com_type: ComType) -> ceng.Or:
    if com_type not in VALID_COM_TYPES:
        raise InvalidRule(f"Invalid COM type present {com_type}")

    com_db = load_com_database(com_type)
    guid_strings: Optional[List[str]] = com_db.get(com_name)
    if guid_strings is None or len(guid_strings) == 0:
        logger.error(" %s doesn't exist in COM %s database", com_name, com_type)
        raise InvalidRule(f"'{com_name}' doesn't exist in COM {com_type} database")

    com_features: List = []
    for guid_string in guid_strings:

Review comment (performance):
could use `@functools.cache` decorator on this function and not have to implement the cache yourself.

Code after:
        return f"invalid rule set: {self.msg}"

    def __repr__(self):
        return str(self)


class ComType(Enum):
    CLASS = "class"
    INTERFACE = "interface"


# COM data source https://github.com/stevemk14ebr/COM-Code-Helper/tree/master
VALID_COM_TYPES = {
    ComType.CLASS: {"db_path": "assets/classes.json.gz", "prefix": "CLSID_"},
    ComType.INTERFACE: {"db_path": "assets/interfaces.json.gz", "prefix": "IID_"},
}


@lru_cache(maxsize=None)
def load_com_database(com_type: ComType) -> Dict[str, List[str]]:
    com_db_path: Path = capa.main.get_default_root() / VALID_COM_TYPES[com_type]["db_path"]

    if not com_db_path.exists():
        raise IOError(f"COM database path '{com_db_path}' does not exist or cannot be accessed")

    try:
        with gzip.open(com_db_path, "rb") as gzfile:
            return json.loads(gzfile.read().decode("utf-8"))
    except Exception as e:
        raise IOError(f"Error loading COM database from '{com_db_path}'") from e


def translate_com_feature(com_name: str, com_type: ComType) -> ceng.Or:
    com_db = load_com_database(com_type)
    guid_strings: Optional[List[str]] = com_db.get(com_name)
    if guid_strings is None or len(guid_strings) == 0:
        logger.error(" %s doesn't exist in COM %s database", com_name, com_type)
        raise InvalidRule(f"'{com_name}' doesn't exist in COM {com_type} database")

    com_features: List = []
    for guid_string in guid_strings:
        hex_chars = guid_string.replace("-", "")
        h = [hex_chars[i : i + 2] for i in range(0, len(hex_chars), 2)]
        reordered_hex_pairs = [
            h[3],
            h[2],
            h[1],
            h[0],
            h[5],
            h[4],
            h[7],