add: tests and fix obfuscators

Author: deoktr

.gitignore

@@ -121,7 +121,7 @@ celerybeat.pid
 
 # Environments
 .env
-.venv
+.venv/
 env/
 venv/
 ENV/

pof/obfuscator/builtins.py

@@ -18,7 +18,18 @@
 # TODO (deoktr): add `__name__.__class__.__class__.__base__.__subclasses__()` variant
 
 import random
-from tokenize import LPAR, LSQB, NAME, NUMBER, OP, RPAR, RSQB, STRING
+from tokenize import (
+    FSTRING_END,
+    FSTRING_START,
+    LPAR,
+    LSQB,
+    NAME,
+    NUMBER,
+    OP,
+    RPAR,
+    RSQB,
+    STRING,
+)
 from typing import ClassVar
 
 from pof.logger import logger
@@ -470,6 +481,7 @@ def obfuscate_builtins(self, tokval):
     def obfuscate_tokens(self, tokens):
         result = []
         parenthesis_depth = 0  # parenthesis depth
+        fstring_depth = 0
         prev_tokval = None
         for index, (toknum, tokval, *_) in enumerate(tokens):
             new_tokens = [(toknum, tokval)]
@@ -482,14 +494,20 @@ def obfuscate_tokens(self, tokens):
             elif toknum == OP and tokval == ")":
                 parenthesis_depth -= 1
 
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+            elif toknum == FSTRING_END:
+                fstring_depth -= 1
+
             if (
                 toknum == NAME
                 and tokval in self.BUILTINS
-                and prev_tokval != "."  # avoid changing class/imports functions
+                and prev_tokval not in (".", "def", "class")
                 and (
                     parenthesis_depth == 0
                     or (parenthesis_depth > 0 and next_tokval != "=")
                 )
+                and fstring_depth == 0
             ):
                 new_tokens = self.obfuscate_builtins(tokval)
 

pof/obfuscator/constants.py

@@ -38,7 +38,18 @@
 #
 
 import random
-from tokenize import DEDENT, ENCODING, INDENT, NAME, NEWLINE, NUMBER, OP, STRING
+from tokenize import (
+    DEDENT,
+    ENCODING,
+    FSTRING_END,
+    FSTRING_START,
+    INDENT,
+    NAME,
+    NEWLINE,
+    NUMBER,
+    OP,
+    STRING,
+)
 
 from pof.utils.generator import BasicGenerator
 from pof.utils.tokens import merge_implicit_strings
@@ -238,6 +249,7 @@ def obfuscate_tokens(self, tokens):
         variables = {}
         result = []
         parenthesis_depth = 0  # parenthesis depth
+        fstring_depth = 0
         prev_tokval = None
         prev_toknum = None
         for index, (toknum, tokval, *_) in enumerate(tokens):
@@ -252,12 +264,17 @@ def obfuscate_tokens(self, tokens):
             elif toknum == OP and tokval == ")":
                 parenthesis_depth -= 1
 
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+            elif toknum == FSTRING_END:
+                fstring_depth -= 1
+
             # obfuscation
-            if (
+            if fstring_depth == 0 and (
                 (
                     toknum == NAME
                     and tokval in self.BUILTINS
-                    and prev_tokval != "."  # avoid changing class/imports functions
+                    and prev_tokval not in (".", "def", "class")
                     and (
                         parenthesis_depth == 0
                         or (parenthesis_depth > 0 and next_tokval != "=")

pof/obfuscator/controlflow/control_flow_flatten.py

@@ -40,6 +40,8 @@ def _should_skip_function(body: list[ast.stmt]) -> bool:
                     ast.AsyncFor,
                     ast.AsyncWith,
                     ast.Await,
+                    ast.Nonlocal,
+                    ast.Global,
                 ),
             ):
                 return True

pof/obfuscator/esoteric/call.py

@@ -15,7 +15,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import keyword
-from tokenize import NAME, OP
+from tokenize import FSTRING_END, FSTRING_START, NAME, OP
 
 
 # TODO (deoktr): add frequency
@@ -39,18 +39,25 @@ class CallObfuscator:
     def obfuscate_tokens(cls, tokens):
         result = []  # obfuscated tokens
         prev_tokval = None
+        fstring_depth = 0
         for index, (toknum, tokval, *_) in enumerate(tokens):
             new_tokens = [(toknum, tokval)]
             next_tokval = None
             if len(tokens) > index + 1:
                 _, next_tokval, *__ = tokens[index + 1]
 
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+            elif toknum == FSTRING_END:
+                fstring_depth -= 1
+
             if (
                 # ensure it's not a definition
                 (prev_tokval is None or prev_tokval not in ["def", "class"])
                 and toknum == NAME
                 and tokval not in cls.RESERVED
                 and next_tokval == "("
+                and fstring_depth == 0
             ):
                 new_tokens.extend(
                     [

pof/obfuscator/extract_variables.py

@@ -15,7 +15,19 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import keyword
-from tokenize import DEDENT, ENCODING, INDENT, NAME, NEWLINE, NL, NUMBER, OP, STRING
+from tokenize import (
+    DEDENT,
+    ENCODING,
+    FSTRING_END,
+    FSTRING_START,
+    INDENT,
+    NAME,
+    NEWLINE,
+    NL,
+    NUMBER,
+    OP,
+    STRING,
+)
 
 from pof.utils.generator import BasicGenerator
 from pof.utils.tokens import merge_implicit_strings
@@ -207,12 +219,13 @@ def __init__(self, generator=None) -> None:
     def generate_new_name(self):
         return next(self.generator)
 
-    def obfuscate_tokens(self, tokens):  # noqa: C901
+    def obfuscate_tokens(self, tokens):  # noqa: C901, PLR0912
         tokens = merge_implicit_strings(tokens)
         result = []
         new_line_buffer = []
         line_buffer = []
         parenthesis_depth = 0
+        fstring_depth = 0
         prev_toknum = None
         in_decorator = False
 
@@ -224,6 +237,11 @@ def obfuscate_tokens(self, tokens):  # noqa: C901
             elif toknum == OP and tokval == ")":
                 parenthesis_depth -= 1
 
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+            elif toknum == FSTRING_END:
+                fstring_depth -= 1
+
             # track decorator context, suppress flushing between @ and def/class
             if toknum == OP and tokval == "@":
                 in_decorator = True
@@ -244,8 +262,10 @@ def obfuscate_tokens(self, tokens):  # noqa: C901
             on_continuation_line = first_name_in_line in self.CONTINUATION_KEYWORDS
 
             if (
-                (toknum == STRING and not is_docstring) or toknum == NUMBER
-            ) and not on_continuation_line:
+                ((toknum == STRING and not is_docstring) or toknum == NUMBER)
+                and not on_continuation_line
+                and fstring_depth == 0
+            ):
                 random_name = self.generate_new_name()
                 new_line_buffer.extend(
                     [

pof/obfuscator/junk/add_newlines.py

@@ -14,7 +14,6 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-# TODO (deoktr): add parameter to disable spacing imports
 import random
 from tokenize import DEDENT, INDENT, NEWLINE
 
@@ -35,7 +34,6 @@ def obfuscate_tokens(self, tokens):
                     new_tokens.extend([(NEWLINE, "\n")])
 
                 else:
-                    # FIXME (deoktr): add newlines BEFORE class/function decorators
                     next_non_indent_tokval = None
                     for i in range(index, len(tokens)):
                         tn, tv, *__ = tokens[i]

pof/obfuscator/numbers.py

@@ -16,7 +16,19 @@
 
 import random
 from enum import Enum
-from tokenize import COMMA, LPAR, NAME, NUMBER, OP, PLUS, RPAR, STRING, untokenize
+from tokenize import (
+    COMMA,
+    FSTRING_END,
+    FSTRING_START,
+    LPAR,
+    NAME,
+    NUMBER,
+    OP,
+    PLUS,
+    RPAR,
+    STRING,
+    untokenize,
+)
 
 from pof.errors import PofError
 from pof.logger import logger
@@ -274,10 +286,16 @@ def obfuscate_number(self, toknum, tokval):  # noqa: C901 PLR0912
 
     def obfuscate_tokens(self, tokens):
         result = []
+        fstring_depth = 0
         for toknum, tokval, *_ in tokens:
             new_tokens = [(toknum, tokval)]
 
-            if toknum == NUMBER:
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+            elif toknum == FSTRING_END:
+                fstring_depth -= 1
+
+            if toknum == NUMBER and fstring_depth == 0:
                 new_tokens = self.obfuscate_number(toknum, tokval)
 
             if new_tokens:

pof/obfuscator/other/tokens.py

@@ -14,7 +14,17 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-from tokenize import LPAR, NAME, NEWLINE, NUMBER, OP, RPAR, STRING
+from tokenize import (
+    FSTRING_END,
+    FSTRING_START,
+    LPAR,
+    NAME,
+    NEWLINE,
+    NUMBER,
+    OP,
+    RPAR,
+    STRING,
+)
 
 
 class TokensObfuscator:
@@ -33,7 +43,33 @@ def import_tokens():
     def generate_tokens_list(tokens):
         tokens_list = []
         tokens_list.append((OP, "["))
+        fstring_depth = 0
+        fstring_parts: list[str] = []
         for toknum, tokval, *_ in tokens:
+            if toknum == FSTRING_START:
+                fstring_depth += 1
+                fstring_parts = [tokval]
+                continue
+            if fstring_depth > 0:
+                if toknum == FSTRING_END:
+                    fstring_depth -= 1
+                    if fstring_depth == 0:
+                        fstring_parts.append(tokval)
+                        merged = "".join(fstring_parts)
+                        tokens_list.extend(
+                            [
+                                (LPAR, "("),
+                                (NUMBER, str(STRING)),
+                                (OP, ","),
+                                (STRING, repr(merged)),
+                                (RPAR, ")"),
+                                (OP, ","),
+                            ],
+                        )
+                        fstring_parts = []
+                        continue
+                fstring_parts.append(tokval)
+                continue
             tokens_list.extend(
                 [
                     (LPAR, "("),

pof/obfuscator/remove/print.py

@@ -46,7 +46,12 @@ def obfuscate_tokens(tokens):
         for toknum, tokval, *_ in tokens:
             new_tokens = [(toknum, tokval)]
 
-            if not inside_print and toknum == NAME and tokval == "print":
+            if (
+                not inside_print
+                and toknum == NAME
+                and tokval == "print"
+                and prev_tokval not in ("def", ".")
+            ):
                 new_tokens = None
                 inside_print = True
                 print_par_depth = parenthesis_depth