FilesExpand file tree

main

/

test_transforms.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

100 lines (82 loc) · 3.73 KB

main

/

test_transforms.py

Copy path

Top

File metadata and controls

• Code
• Blame

100 lines (82 loc) · 3.73 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

from __future__ import annotations

import ast

import contextlib

import io

import random

import sys

import unittest

from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

sys.path.insert(0, str(ROOT))

from patchwork.transforms.identifiers import IdentifierRenamer

from patchwork.transforms.normalize import MatchLowerer

from patchwork.transforms.numbers import NumberObfuscator

from patchwork.transforms.strings import StringEncryptor, build_decrypt_helper

def eval_ast_expr(node: ast.AST) -> object:

expr = ast.Expression(body=node)

ast.fix_missing_locations(expr)

return eval(compile(expr, "<test>", "eval"), {})

def run_module(source: str) -> str:

stdout = io.StringIO()

with contextlib.redirect_stdout(stdout):

exec(compile(source, "<test>", "exec"), {})

return stdout.getvalue()

class TransformTests(unittest.TestCase):

def test_new_number_transforms_preserve_values(self) -> None:

values = [-65537, -129, -1, 0, 1, 42, 4096, 2**40 + 123]

methods = ["_t_invert", "_t_bitmask", "_t_divmod", "_t_table"]

for method_name in methods:

for value in values:

with self.subTest(method=method_name, value=value):

obfuscator = NumberObfuscator(random.Random(2026), prob=1.0)

node = getattr(obfuscator, method_name)(value)

self.assertEqual(value, eval_ast_expr(node))

def test_string_encryptor_builds_shared_literal_pool(self) -> None:

source = (

"VALUE = 'PATCHWORK_SECRET_MARKER_2026'\n"

"RAW = b'PATCHWORK_BYTES_MARKER_2026'\n"

"print(VALUE, RAW.decode())\n"

)

tree = ast.parse(source)

encryptor = StringEncryptor(random.Random(17), "_pw_dec_str", "_pw_dec_bytes", "_pw_literal_pool")

tree = encryptor.visit(tree)

pool_stmt = encryptor.build_pool_stmt()

self.assertIsNotNone(pool_stmt)

tree.body = [pool_stmt, *build_decrypt_helper("_pw_dec_str", "_pw_dec_bytes"), *tree.body]

ast.fix_missing_locations(tree)

transformed_source = ast.unparse(tree)

self.assertIn("_pw_literal_pool", transformed_source)

self.assertNotIn("PATCHWORK_SECRET_MARKER_2026", transformed_source)

self.assertNotIn("PATCHWORK_BYTES_MARKER_2026", transformed_source)

stdout = io.StringIO()

with contextlib.redirect_stdout(stdout):

exec(compile(tree, "<test>", "exec"), {})

self.assertEqual("PATCHWORK_SECRET_MARKER_2026 PATCHWORK_BYTES_MARKER_2026\n", stdout.getvalue())

def test_reserved_import_keeps_original_binding(self) -> None:

tree = ast.parse("from collections import deque\nprint(deque([1, 2, 3]))\n")

renamed = IdentifierRenamer(random.Random(5), keep={"deque"}).visit(tree)

ast.fix_missing_locations(renamed)

source = ast.unparse(renamed)

self.assertIn("from collections import deque", source)

stdout = io.StringIO()

with contextlib.redirect_stdout(stdout):

exec(compile(renamed, "<test>", "exec"), {})

self.assertIn("deque([1, 2, 3])", stdout.getvalue())

def test_match_lowerer_preserves_general_sequence_patterns(self) -> None:

source = """

from collections import deque

def classify(value):

match value:

case (first, *middle, last):

return f"{type(value).__name__}:{first}:{middle}:{last}"

case _:

return f"{type(value).__name__}:no"

print(classify(range(4)))

print(classify("abcd"))

print(classify(deque([1, 2, 3])))

"""

tree = MatchLowerer().visit(ast.parse(source))

ast.fix_missing_locations(tree)

self.assertEqual(run_module(source), run_module(ast.unparse(tree)))

if __name__ == "__main__":

unittest.main()