From 3d812edc4da6e008c40072d43e6fa194ab9a248d Mon Sep 17 00:00:00 2001
From: Yacine Elhamer <elhamer.yacine@gmail.com>
Date: Thu, 27 Jul 2023 15:52:39 +0100
Subject: [PATCH] use weakrefs for `Scopes` instantiation; fix test_rules()

---
 capa/helpers.py        | 21 +++++++++++++++++++++
 capa/rules/__init__.py |  4 ++--
 tests/test_rules.py    |  2 +-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/capa/helpers.py b/capa/helpers.py
index f7978b54..b47d9852 100644
--- a/capa/helpers.py
+++ b/capa/helpers.py
@@ -8,6 +8,8 @@
 import json
 import inspect
 import logging
+import weakref
+import functools
 import contextlib
 import importlib.util
 from typing import NoReturn
@@ -130,6 +132,25 @@ def redirecting_print_to_tqdm(disable_progress):
         inspect.builtins.print = old_print  # type: ignore
 
 
+def weak_lru(maxsize=128, typed=False):
+    """
+    LRU Cache decorator that keeps a weak reference to 'self'
+    """
+
+    def wrapper(func):
+        @functools.lru_cache(maxsize, typed)
+        def _func(_self, *args, **kwargs):
+            return func(_self(), *args, **kwargs)
+
+        @functools.wraps(func)
+        def inner(self, *args, **kwargs):
+            return _func(weakref.ref(self), *args, **kwargs)
+
+        return inner
+
+    return wrapper
+
+
 def log_unsupported_format_error():
     logger.error("-" * 80)
     logger.error(" Input file does not appear to be a PE or ELF file.")
diff --git a/capa/rules/__init__.py b/capa/rules/__init__.py
index c327c18c..47c29764 100644
--- a/capa/rules/__init__.py
+++ b/capa/rules/__init__.py
@@ -16,7 +16,7 @@ import collections
 from enum import Enum
 from pathlib import Path
 
-from capa.helpers import assert_never
+from capa.helpers import weak_lru, assert_never
 
 try:
     from functools import lru_cache
@@ -115,7 +115,7 @@ class Scopes:
     static: Optional[str] = None
     dynamic: Optional[str] = None
 
-    @lru_cache()  # type: ignore
+    @weak_lru()
     def __new__(cls, *args, **kwargs):
         return super().__new__(cls)
 
diff --git a/tests/test_rules.py b/tests/test_rules.py
index 262be9a5..34c05204 100644
--- a/tests/test_rules.py
+++ b/tests/test_rules.py
@@ -357,7 +357,7 @@ def test_rules_flavor_filtering():
         ),
     ]
 
-    static_rules = capa.rules.RuleSet(rules, rules_filter_func=lambda rule: rule.scopes.static)
+    static_rules = capa.rules.RuleSet(rules.copy(), rules_filter_func=lambda rule: rule.scopes.static)
     dynamic_rules = capa.rules.RuleSet(rules, rules_filter_func=lambda rule: rule.scopes.dynamic)
 
     # only static rule