mirror of
https://github.com/bikini/patchwork
synced 2026-06-27 08:08:41 +00:00
bikini
65 lines
2.8 KiB
Python
65 lines
2.8 KiB
Python
from __future__ import annotations
|
|
import ast
|
|
import random
|
|
_JUNK_NAMES = ('_v', '_t', '_acc', '_state', '_buf', '_h', '_z', '_w', '_q', '_n')
|
|
_JUNK_STMT_TEMPLATES = ('{v} = ({a} ^ {b}) + ({c} & {d})', '{v} = ((({a} * 0x01000193) ^ {b}) & 0xFFFFFFFF)', '{v} = sum(((({a} >> _k) & 0xFF) for _k in range(0, 32, 8)))', '{v} = ({a} | {b}) - (({c} & {d}) >> 1)', '{v} = pow({a} | 1, 1, ({b} | 0x10001))', '{v} = (({a} + {b}) ^ ({c} - {d})) & 0x7FFFFFFF', '{v} = bytes(((_i * {a}) ^ {b}) & 0xFF for _i in range(8))', '{v} = (~({a}) | ({b} & {c})) & 0xFFFFFF')
|
|
|
|
def _rint(rng: random.Random) -> int:
|
|
return rng.randrange(1, 1 << 30)
|
|
|
|
class JunkBranchInjector(ast.NodeTransformer):
|
|
|
|
def __init__(self, rng: random.Random, seed_name: str, prob: float=0.4, max_depth: int=0):
|
|
self.rng = rng
|
|
self.seed_name = seed_name
|
|
self.prob = prob
|
|
self.injected = False
|
|
self._depth = 0
|
|
self._max_depth = max_depth
|
|
|
|
def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
|
|
return self._inject_into_body(node)
|
|
visit_AsyncFunctionDef = visit_FunctionDef
|
|
|
|
def visit_Module(self, node: ast.Module) -> ast.AST:
|
|
return self._inject_into_body(node)
|
|
|
|
def _inject_into_body(self, node):
|
|
for child in ast.iter_child_nodes(node):
|
|
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
|
|
self.visit(child)
|
|
if not getattr(node, 'body', None):
|
|
return node
|
|
if self.rng.random() > self.prob:
|
|
return node
|
|
block = self._build_block()
|
|
idx = self.rng.randint(0, len(node.body))
|
|
node.body.insert(idx, block)
|
|
self.injected = True
|
|
return node
|
|
|
|
def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
|
|
for child in node.body:
|
|
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
self.visit(child)
|
|
return node
|
|
|
|
def _build_block(self) -> ast.stmt:
|
|
s = self.seed_name
|
|
kind = self.rng.choice(('parity', 'flt3', 'flt5'))
|
|
if kind == 'parity':
|
|
cond_src = f'({s} * ({s} + 1) + 1) % 2 == 0'
|
|
elif kind == 'flt3':
|
|
cond_src = f'({s} * {s} * {s} - {s} + 1) % 3 == 0'
|
|
else:
|
|
cond_src = f'({s} * {s} * {s} * {s} * {s} - {s} + 1) % 5 == 0'
|
|
cond = ast.parse(cond_src, mode='eval').body
|
|
body = [self._random_stmt() for _ in range(self.rng.randint(2, 4))]
|
|
return ast.If(test=cond, body=body, orelse=[])
|
|
|
|
def _random_stmt(self) -> ast.stmt:
|
|
v = self.rng.choice(_JUNK_NAMES)
|
|
tmpl = self.rng.choice(_JUNK_STMT_TEMPLATES)
|
|
src = tmpl.format(v=v, a=_rint(self.rng), b=_rint(self.rng), c=_rint(self.rng), d=_rint(self.rng))
|
|
return ast.parse(src).body[0]
|