mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
ashton
203 lines
7.5 KiB
Python
203 lines
7.5 KiB
Python
from __future__ import annotations
|
|
import ast
|
|
import random
|
|
from collections.abc import Sequence
|
|
|
|
class StringEncryptor(ast.NodeTransformer):
|
|
|
|
def __init__(self, rng: random.Random, decrypt_str_name: str, decrypt_bytes_name: str, pool_name: str):
|
|
self.rng = rng
|
|
self.dec_s = decrypt_str_name
|
|
self.dec_b = decrypt_bytes_name
|
|
self.pool_name = pool_name
|
|
self.entries: list[tuple[list[bytes], list[bytes], list[int], list[int], int, int]] = []
|
|
|
|
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 _rotate_left(self, data: bytes, amount: int) -> bytes:
|
|
if not data:
|
|
return data
|
|
amount %= len(data)
|
|
return data[amount:] + data[:amount]
|
|
|
|
def _split_chunks(self, data: bytes) -> list[bytes]:
|
|
if len(data) <= 2:
|
|
return [data]
|
|
max_chunks = min(8, len(data))
|
|
min_chunks = 2 if len(data) > 4 else 1
|
|
count = self.rng.randint(min_chunks, max_chunks)
|
|
if count <= 1:
|
|
return [data]
|
|
cuts = sorted(self.rng.sample(range(1, len(data)), count - 1))
|
|
positions = [0, *cuts, len(data)]
|
|
return [data[positions[i]:positions[i + 1]] for i in range(len(positions) - 1)]
|
|
|
|
def _shuffle_chunks(self, chunks: Sequence[bytes]) -> tuple[list[bytes], list[int]]:
|
|
if len(chunks) <= 1:
|
|
return list(chunks), [0]
|
|
shuffled_indexes = list(range(len(chunks)))
|
|
self.rng.shuffle(shuffled_indexes)
|
|
shuffled = [chunks[i] for i in shuffled_indexes]
|
|
order = [shuffled_indexes.index(i) for i in range(len(chunks))]
|
|
return shuffled, order
|
|
|
|
def _bytes_tuple(self, values: Sequence[bytes]) -> ast.Tuple:
|
|
return ast.Tuple(elts=[ast.Constant(value=value) for value in values], ctx=ast.Load())
|
|
|
|
def _int_tuple(self, values: Sequence[int]) -> ast.Tuple:
|
|
return ast.Tuple(elts=[ast.Constant(value=value) for value in values], ctx=ast.Load())
|
|
|
|
def _encoded_entry(self, data: bytes, key: bytes) -> tuple[list[bytes], list[bytes], list[int], list[int], int, int]:
|
|
enc = self._xor(data, key)
|
|
rotation = self.rng.randrange(len(enc)) if enc else 0
|
|
rotated = self._rotate_left(enc, rotation)
|
|
data_chunks, data_order = self._shuffle_chunks(self._split_chunks(rotated))
|
|
key_chunks, key_order = self._shuffle_chunks(self._split_chunks(key))
|
|
mode = self.rng.randrange(3)
|
|
return data_chunks, key_chunks, data_order, key_order, rotation, mode
|
|
|
|
def _store_entry(self, data: bytes, key: bytes) -> int:
|
|
idx = len(self.entries)
|
|
self.entries.append(self._encoded_entry(data, key))
|
|
return idx
|
|
|
|
def _pool_ref(self, idx: int, decoder_name: str) -> ast.Call:
|
|
return ast.Call(
|
|
func=ast.Name(id=decoder_name, ctx=ast.Load()),
|
|
args=[
|
|
ast.Name(id=self.pool_name, ctx=ast.Load()),
|
|
ast.Constant(value=idx),
|
|
],
|
|
keywords=[],
|
|
)
|
|
|
|
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()
|
|
return self._pool_ref(self._store_entry(data, key), self.dec_s)
|
|
|
|
def _encrypt_bytes(self, b: bytes) -> ast.AST:
|
|
key = self._make_key()
|
|
return self._pool_ref(self._store_entry(b, key), self.dec_b)
|
|
|
|
def build_pool_stmt(self) -> ast.stmt | None:
|
|
if not self.entries:
|
|
return None
|
|
entry_nodes = []
|
|
for data_chunks, key_chunks, data_order, key_order, rotation, mode in self.entries:
|
|
entry_nodes.append(
|
|
ast.Tuple(
|
|
elts=[
|
|
self._bytes_tuple(data_chunks),
|
|
self._bytes_tuple(key_chunks),
|
|
self._int_tuple(data_order),
|
|
self._int_tuple(key_order),
|
|
ast.Constant(value=rotation),
|
|
ast.Constant(value=mode),
|
|
],
|
|
ctx=ast.Load(),
|
|
)
|
|
)
|
|
return ast.Assign(
|
|
targets=[ast.Name(id=self.pool_name, ctx=ast.Store())],
|
|
value=ast.Tuple(elts=entry_nodes, ctx=ast.Load()),
|
|
)
|
|
|
|
def build_decrypt_helper(decrypt_str_name: str, decrypt_bytes_name: str) -> list[ast.stmt]:
|
|
src = f'''
|
|
def _pw_join_chunks(_chunks, _order):
|
|
return b''.join(_chunks[_i] for _i in _order)
|
|
|
|
def _pw_unrotate(_data, _amount):
|
|
if not _data:
|
|
return _data
|
|
_amount %= len(_data)
|
|
return _data[-_amount:] + _data[:-_amount] if _amount else _data
|
|
|
|
def _pw_xor_a(_data, _key):
|
|
return bytes(_b ^ _key[_i % len(_key)] for _i, _b in enumerate(_data))
|
|
|
|
def _pw_xor_b(_data, _key):
|
|
_out = bytearray(len(_data))
|
|
for _i, _b in enumerate(_data):
|
|
_out[_i] = _b ^ _key[_i % len(_key)]
|
|
return bytes(_out)
|
|
|
|
def _pw_xor_c(_data, _key):
|
|
return bytes(map(lambda _p: _p[1] ^ _key[_p[0] % len(_key)], enumerate(_data)))
|
|
|
|
def _pw_decode_entry(_pool, _idx):
|
|
_d, _k, _do, _ko, _rot, _mode = _pool[_idx]
|
|
_data = _pw_unrotate(_pw_join_chunks(_d, _do), _rot)
|
|
_key = _pw_join_chunks(_k, _ko)
|
|
if _mode == 1:
|
|
return _pw_xor_b(_data, _key)
|
|
if _mode == 2:
|
|
return _pw_xor_c(_data, _key)
|
|
return _pw_xor_a(_data, _key)
|
|
|
|
def {decrypt_bytes_name}(_pool, _idx):
|
|
return _pw_decode_entry(_pool, _idx)
|
|
|
|
def {decrypt_str_name}(_pool, _idx):
|
|
return _pw_decode_entry(_pool, _idx).decode()
|
|
'''
|
|
return ast.parse(src).body
|