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

Code before:
        self.__cause__ = None

    def __str__(self):
        return f"invalid rule: {self.path}: {self.msg}"


class InvalidRuleSet(ValueError):
    def __init__(self, msg):
        super().__init__()
        self.msg = msg

    def __str__(self):
        return f"invalid rule set: {self.msg}"

    def __repr__(self):
        return str(self)


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


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

    CD = Path(__file__).resolve().parent.parent.parent
    com_db_path = CD / VALID_COM_TYPES[com_type]["db_path"]
    if not com_db_path.exists():
        logger.error("Using COM %s database '%s', but it doesn't exist", com_type, com_db_path)
        raise IOError(f"COM database path '{com_db_path}' does not exist or cannot be accessed")

    with gzip.open(com_db_path, "rb") as gzfile:
        com_db: Dict[str, List[str]] = json.loads(gzfile.read().decode("utf-8"))
        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],

Review comment (suggestion):
```suggestion
def translate_com_feature(com_name: str, com_type: Literal["class", "interface"]) -> ceng.Or:
```

alternatively introduce an enum for this.

Code after:
        self.__cause__ = None

    def __str__(self):
        return f"invalid rule: {self.path}: {self.msg}"


class InvalidRuleSet(ValueError):
    def __init__(self, msg):
        super().__init__()
        self.msg = msg

    def __str__(self):
        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")