File size: 3,103 Bytes
fc32096 | 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 | 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()
|