Improvvments to substitution as well as comparing in tests

kam193 · Jan 3, 2025 · 30cc7f1 · 30cc7f1
1 parent 0a70ffa
commit 30cc7f1
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 3 deletions.
diff --git a/ASTGrep/rules/extended/python/renamed-builtin.yml b/ASTGrep/rules/extended/python/renamed-builtin.yml
@@ -1036,6 +1036,8 @@ utils:
           - pattern: choice
           - pattern: system
           - pattern: winreg
+          - pattern: base64
+          - pattern: gc.get_referents
           # TODO: better patterns
           # - pattern: "''.join"
           - pattern: binascii.unhexlify

diff --git a/ASTGrep/rules/helpers/python/substitute-static.yml b/ASTGrep/rules/helpers/python/substitute-static.yml
@@ -1490,6 +1490,46 @@ utils:
       - kind: float
       - kind: string
       - kind: "none"
+      - kind: identifier
+        not:
+          # Handled by renamed-builtin ->  TODO: merge both!
+          any:
+            - pattern: eval
+            - pattern: exec
+            - pattern: str
+            - pattern: map
+            - pattern: ord
+            - pattern: tuple
+            - pattern: type
+            - pattern: globals
+            - pattern: locals
+            - pattern: __import__
+            - pattern: compile
+            - pattern: getattr
+            - pattern: dir
+            - pattern: vars
+            - pattern: bool
+            - pattern: float
+            - pattern: int
+            - pattern: bytes
+            - pattern: list
+            - pattern: pow
+            - pattern: round
+            - pattern: range
+            - pattern: open
+            - pattern: chr
+            - pattern: request
+            - pattern: getenv
+            - pattern: listdir
+            - pattern: exit
+            - pattern: choice
+            - pattern: system
+            - pattern: winreg
+            - pattern: base64
+            - pattern: gc.get_referents
+            # TODO: better patterns
+            # - pattern: "''.join"
+            - pattern: binascii.unhexlify
 
 metadata:
   extended-obfuscation: no

diff --git a/ASTGrep/rules/templates/python-replace-static-var.yml.j2 b/ASTGrep/rules/templates/python-replace-static-var.yml.j2
@@ -94,10 +94,24 @@ rule:
                     - kind: "false"
                     - kind: float
                     - kind: "none"
+                    - pattern: {{ VALUE | replace("$", "{{{DOLLARPLACEHOLDER}}}") | tojson }}
                   # TODO: good regex sanitization to exclude only specific string
                   #has:
                   #  kind: string_content
                   #  {# regex: {{ VALUE | replace('?', '\\?') | replace('(', '\\(') | replace(')', '\\)') | replace('[', '\\[') | replace(']', '\\]') | replace('{', '\\{') | replace('}', '\\}') | replace('*', '\\*') | replace("$", "{{{DOLLARPLACEHOLDER}}}") | tojson }} #}
+      - inside:
+          kind: expression_statement
+          stopBy: end
+          follows:
+            stopBy: end
+            any:
+              - kind: for_statement
+            has:
+              stopBy: end
+              kind: assignment
+              has:
+                field: left
+                pattern: {{ VAR }}
 # TODO: better way to detect if the variable is used in a loop
 # or if it is used in a condition that modifies the variable
 #      - inside:

diff --git a/ASTGrep/tests/other_examples/python/helpers/python-substitute-static-vars.in b/ASTGrep/tests/other_examples/python/helpers/python-substitute-static-vars.in
@@ -22,7 +22,7 @@ def f2(var1):
 var20 = 2
 self.var20 = 2
 
-# TODO: this should be ignored
+# this should be ignored
 x = 7
 for i in range(3):
     x = i
@@ -71,3 +71,6 @@ func(var10)
 
 encrypted_payload = b"abcdefghijklmnopqrstuvwxyz"
 encrypted_payload = some_func(key=b"\x95",nonce=b"\xff").decrypt(encrypted_payload[12:])
+
+lIIIlIIIIlllIlIIII=gc
+lIllIIIlIlIIIIlIlI=lIIIlIIIIlllIlIIII.get_referents
diff --git a/ASTGrep/tests/other_examples/python/helpers/python-substitute-static-vars.out b/ASTGrep/tests/other_examples/python/helpers/python-substitute-static-vars.out
@@ -26,7 +26,7 @@ self.var20 = 2
 x = 7
 for i in range(3):
     x = i
-print(7)
+print(x)
 
 global var2
 var1 += "a"
@@ -71,3 +71,6 @@ func(var10)
 
 encrypted_payload = b"abcdefghijklmnopqrstuvwxyz"
 encrypted_payload = some_func(key=b"\x95",nonce=b"\xff").decrypt(b'mnopqrstuvwxyz')
+
+lIIIlIIIIlllIlIIII=gc
+lIllIIIlIlIIIIlIlI=gc.get_referents
diff --git a/ASTGrep/tests/test_deobfuscation.py b/ASTGrep/tests/test_deobfuscation.py
@@ -61,6 +61,12 @@ def detect_sample(
     ), f"Sample was {'not' if expected else ''} detected as {metadata_field} by {rule_id}"
 
 
+def rstrip_lines(data: str | list[str]) -> str:
+    if not isinstance(data, list):
+        data = data.splitlines()
+    return "\n".join(s.rstrip() for s in data)
+
+
 @pytest.fixture
 def deobfuscate_example(deobfuscator):
     def _check_example(
@@ -76,7 +82,9 @@ def _check_example(
             language = f"code/{lang_name}"
         results = list(deobfuscator.deobfuscate_file(f"{path}.in", language))
         if check_output:
-            assert results[-1][0].strip() == open(f"{path}.out", "r").read().strip()
+            assert rstrip_lines(results[-1][0].strip()) == rstrip_lines(
+                open(f"{path}.out", "r").read().strip()
+            )
         if deobfuscator.work_time > warning_time:
             warnings.warn(f"Deobfuscation took {deobfuscator.work_time:.3f} seconds")
         assert (