Skip to content

Commit

Permalink
Update __init__.py
Browse files Browse the repository at this point in the history
  • Loading branch information
Aayush-Goel-04 committed Oct 17, 2023
1 parent 7cd5aa1 commit 6dbd376
Showing 1 changed file with 2 additions and 9 deletions.
11 changes: 2 additions & 9 deletions capa/rules/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,24 +207,17 @@ class ComType(Enum):
ComType.INTERFACE: {"db_path": "assets/interfaces.json.gz", "prefix": "IID_"},
}

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


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

Expand Down

0 comments on commit 6dbd376

Please sign in to comment.