| import ast |
| import builtins |
| import json |
| import resource |
| import random |
| import sys |
| from pathlib import Path |
|
|
| ALLOWED_IMPORTS = {'math','cmath','re','string','collections','itertools','functools','heapq','bisect','typing', |
| 'statistics','decimal','fractions','operator','random','copy','datetime','calendar','hashlib', |
| 'json','base64','unicodedata','dataclasses','enum','array'} |
| BLOCKED_CALLS = {'open','exec','eval','compile','globals','locals','vars','getattr','setattr','delattr', |
| '__import__','breakpoint','help','input','exit','quit'} |
|
|
|
|
| def validate(source): |
| tree = ast.parse(source) |
| for node in ast.walk(tree): |
| if isinstance(node,(ast.Import,ast.ImportFrom)): |
| modules = [alias.name.split('.')[0] for alias in node.names] if isinstance(node,ast.Import) else [(node.module or '').split('.')[0]] |
| if isinstance(node,ast.ImportFrom) and node.level: |
| raise ValueError('unsupported_relative_import') |
| if any(module not in ALLOWED_IMPORTS for module in modules): |
| raise ValueError('unsupported_import:'+','.join(modules)) |
| if isinstance(node,ast.Name) and node.id.startswith('__') and node.id != '__name__': |
| raise ValueError('unsupported_introspection') |
| if isinstance(node,ast.Attribute) and (node.attr.startswith('_') or isinstance(node.ctx,ast.Store)): |
| raise ValueError('unsupported_attribute_access') |
| if isinstance(node,ast.Call): |
| if isinstance(node.func,ast.Name) and node.func.id in BLOCKED_CALLS: |
| raise ValueError('unsupported_call:'+node.func.id) |
| if isinstance(node.func,ast.Attribute) and node.func.attr in {'attrgetter','methodcaller'}: |
| raise ValueError('unsupported_introspection') |
| return tree |
|
|
|
|
| def main(): |
| resource.setrlimit(resource.RLIMIT_CPU,(2,2)) |
| resource.setrlimit(resource.RLIMIT_AS,(384*1024**2,384*1024**2)) |
| resource.setrlimit(resource.RLIMIT_FSIZE,(65536,65536)) |
| resource.setrlimit(resource.RLIMIT_NOFILE,(64,64)) |
| resource.setrlimit(resource.RLIMIT_CORE,(0,0)) |
| payload = json.loads(Path('/work/payload.json').read_text()) |
| nonce = sys.argv[1] |
| random.seed(0) |
| result = {'passed':False} |
| try: |
| source = payload['code'] |
| tree = validate(source) |
| namespace = {'__builtins__':dict(vars(builtins)),'__name__':'candidate'} |
| exec(compile(tree,'<candidate>','exec'),namespace) |
| candidate = namespace.get(payload['entry_point']) |
| if not callable(candidate): |
| raise ValueError('missing_entry_point') |
| test_namespace = {'__builtins__':dict(vars(builtins)),'__name__':'tests'} |
| exec(compile(payload['test_code'],'<tests>','exec'),test_namespace) |
| test_namespace['check'](candidate) |
| result = {'passed':True} |
| except BaseException as error: |
| result = {'passed':False,'error_type':type(error).__name__,'error':str(error)[:300]} |
| print('__WSP_RESULT__'+nonce+json.dumps(result),flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
