mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
bikini
85 lines
3.4 KiB
Python
85 lines
3.4 KiB
Python
from __future__ import annotations
|
|
import ast
|
|
import random
|
|
|
|
class StringEncryptor(ast.NodeTransformer):
|
|
|
|
def __init__(self, rng: random.Random, decrypt_str_name: str, decrypt_bytes_name: str):
|
|
self.rng = rng
|
|
self.dec_s = decrypt_str_name
|
|
self.dec_b = decrypt_bytes_name
|
|
|
|
def _strip_docstring(self, body: list[ast.stmt]) -> list[ast.stmt]:
|
|
if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) and isinstance(body[0].value.value, str):
|
|
return body[1:] or [ast.Pass()]
|
|
return body
|
|
|
|
def visit_Module(self, node: ast.Module) -> ast.AST:
|
|
node.body = self._strip_docstring(node.body)
|
|
self.generic_visit(node)
|
|
return node
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
|
|
node.body = self._strip_docstring(node.body)
|
|
self.generic_visit(node)
|
|
return node
|
|
visit_AsyncFunctionDef = visit_FunctionDef
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
|
|
node.body = self._strip_docstring(node.body)
|
|
self.generic_visit(node)
|
|
return node
|
|
|
|
def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.AST:
|
|
new_values: list[ast.AST] = []
|
|
for v in node.values:
|
|
if isinstance(v, ast.FormattedValue):
|
|
v.value = self.visit(v.value)
|
|
if v.format_spec is not None:
|
|
v.format_spec = self.visit(v.format_spec)
|
|
new_values.append(v)
|
|
else:
|
|
new_values.append(v)
|
|
node.values = new_values
|
|
return node
|
|
|
|
def visit_match_case(self, node: ast.match_case) -> ast.AST:
|
|
if node.guard is not None:
|
|
node.guard = self.visit(node.guard)
|
|
node.body = [self.visit(stmt) for stmt in node.body]
|
|
return node
|
|
|
|
def visit_Constant(self, node: ast.Constant) -> ast.AST:
|
|
v = node.value
|
|
if isinstance(v, bool):
|
|
return node
|
|
if isinstance(v, str):
|
|
return self._encrypt_str(v) if v else node
|
|
if isinstance(v, bytes):
|
|
return self._encrypt_bytes(v) if v else node
|
|
return node
|
|
|
|
def _make_key(self, n_min: int=4, n_max: int=24) -> bytes:
|
|
return bytes((self.rng.randrange(256) for _ in range(self.rng.randint(n_min, n_max))))
|
|
|
|
def _xor(self, data: bytes, key: bytes) -> bytes:
|
|
return bytes((b ^ key[i % len(key)] for i, b in enumerate(data)))
|
|
|
|
def _encrypt_str(self, s: str) -> ast.AST:
|
|
try:
|
|
data = s.encode('utf-8')
|
|
except UnicodeEncodeError:
|
|
return ast.Constant(value=s)
|
|
key = self._make_key()
|
|
enc = self._xor(data, key)
|
|
return ast.Call(func=ast.Name(id=self.dec_s, ctx=ast.Load()), args=[ast.Constant(value=enc), ast.Constant(value=key)], keywords=[])
|
|
|
|
def _encrypt_bytes(self, b: bytes) -> ast.AST:
|
|
key = self._make_key()
|
|
enc = self._xor(b, key)
|
|
return ast.Call(func=ast.Name(id=self.dec_b, ctx=ast.Load()), args=[ast.Constant(value=enc), ast.Constant(value=key)], keywords=[])
|
|
|
|
def build_decrypt_helper(decrypt_str_name: str, decrypt_bytes_name: str) -> list[ast.stmt]:
|
|
src = f'def {decrypt_str_name}(d, k):\n return bytes(b ^ k[i % len(k)] for i, b in enumerate(d)).decode()\ndef {decrypt_bytes_name}(d, k):\n return bytes(b ^ k[i % len(k)] for i, b in enumerate(d))\n'
|
|
return ast.parse(src).body
|