Make sure code is safe by walking through the AST. Code considered unsafe if: * it has restricted AST nodes (only nodes defined in ALLOWED_AST_NODES are allowed) * it is trying to assign to attributes * it is trying to access resricted attributes Adopted from h
| 1338 | |
| 1339 | |
| 1340 | class SafeVisitor(ast.NodeVisitor): |
| 1341 | """ |
| 1342 | Make sure code is safe by walking through the AST. |
| 1343 | |
| 1344 | Code considered unsafe if: |
| 1345 | * it has restricted AST nodes (only nodes defined in ALLOWED_AST_NODES are allowed) |
| 1346 | * it is trying to assign to attributes |
| 1347 | * it is trying to access resricted attributes |
| 1348 | |
| 1349 | Adopted from http://www.zafar.se/bkz/uploads/safe.txt (public domain, Babar K. Zafar) |
| 1350 | * Using ast rather than compiler tree, for jython and Py3 support since Py2.6 |
| 1351 | * Simplified with ast.NodeVisitor class |
| 1352 | """ |
| 1353 | |
| 1354 | def __init__(self, *args, **kwargs): |
| 1355 | "Initialize visitor by generating callbacks for all AST node types." |
| 1356 | super().__init__(*args, **kwargs) |
| 1357 | self.errors = [] |
| 1358 | |
| 1359 | def walk(self, tree, filename): |
| 1360 | "Validate each node in AST and raise SecurityError if the code is not safe." |
| 1361 | self.filename = filename |
| 1362 | self.visit(tree) |
| 1363 | if self.errors: |
| 1364 | raise SecurityError("\n".join([str(err) for err in self.errors])) |
| 1365 | |
| 1366 | def generic_visit(self, node): |
| 1367 | nodename = type(node).__name__ |
| 1368 | if nodename not in ALLOWED_AST_NODES: |
| 1369 | self.fail_node(node, nodename) |
| 1370 | super().generic_visit(node) |
| 1371 | |
| 1372 | def visit_Name(self, node): |
| 1373 | if node.id.startswith("__"): |
| 1374 | self.fail_name(node) |
| 1375 | |
| 1376 | def visit_Attribute(self, node): |
| 1377 | attrname = self.get_node_attr(node) |
| 1378 | if self.is_unallowed_attr(attrname): |
| 1379 | self.fail_attribute(node, attrname) |
| 1380 | super().generic_visit(node) |
| 1381 | |
| 1382 | def visit_Assign(self, node): |
| 1383 | self.check_assign_targets(node) |
| 1384 | |
| 1385 | def visit_AugAssign(self, node): |
| 1386 | self.check_assign_target(node) |
| 1387 | |
| 1388 | def check_assign_targets(self, node): |
| 1389 | for target in node.targets: |
| 1390 | self.check_assign_target(target) |
| 1391 | super().generic_visit(node) |
| 1392 | |
| 1393 | def check_assign_target(self, targetnode): |
| 1394 | targetname = type(targetnode).__name__ |
| 1395 | if targetname == "Attribute": |
| 1396 | attrname = self.get_node_attr(targetnode) |
| 1397 | self.fail_attribute(targetnode, attrname) |