Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 794a42b

Browse files
KevinHockKevinHock
authored
Merge pull request #155 from bcaller/yield
Yield, YieldFrom, AugAssign propagate taint
2 parents dec57bd + 04b29c6 commit 794a42b

10 files changed

Lines changed: 160 additions & 40 deletions

File tree

‎examples/example_inputs/generator_with_multiple_yields.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
def foo():
2+
a = 1
3+
if a == 1:
4+
yield 0
5+
yield a
6+
7+
foo()

‎examples/vulnerable_code/yield.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import subprocess
2+
from flask import Flask, request
3+
4+
app = Flask(__name__)
5+
6+
7+
def things_to_run():
8+
yield "echo hello"
9+
yield from request.get_json()["commands"]
10+
yield "echo done"
11+
12+
13+
@app.route('/', methods=['POST'])
14+
def home():
15+
script = "; ".join(things_to_run())
16+
subprocess.call(script, shell=True)
17+
return 'Executed'
18+
19+
20+
if __name__ == '__main__':
21+
app.run(debug=True)

‎pyt/cfg/expr_visitor.py‎

Lines changed: 34 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515
IgnoredNode,
1616
Node,
1717
RestoreNode,
18-
ReturnNode
18+
ReturnNode,
19+
YieldNode
1920
)
2021
from .expr_visitor_helper import (
2122
BUILTINS,
@@ -113,22 +114,25 @@ def visit_Yield(self, node):
113114
label = LabelVisitor()
114115
label.visit(node)
115116

116-
try:
117-
rhs_visitor = RHSVisitor()
118-
rhs_visitor.visit(node.value)
119-
except AttributeError:
120-
rhs_visitor.result = 'EmptyYield'
117+
if node.value is None:
118+
rhs_visitor_result = []
119+
else:
120+
rhs_visitor_result = RHSVisitor.result_for_node(node.value)
121121

122+
# Yield is a bit like augmented assignment to a return value
122123
this_function_name = self.function_return_stack[-1]
123-
LHS = 'yield_' + this_function_name
124-
return self.append_node(ReturnNode(
125-
LHS + ' = ' + label.result,
124+
LHS = 'yld_' + this_function_name
125+
return self.append_node(YieldNode(
126+
LHS + ' += ' + label.result,
126127
LHS,
127128
node,
128-
rhs_visitor.result,
129+
rhs_visitor_result + [LHS],
129130
path=self.filenames[-1])
130131
)
131132

133+
def visit_YieldFrom(self, node):
134+
return self.visit_Yield(node)
135+
132136
def visit_Attribute(self, node):
133137
return self.visit_miscelleaneous_node(
134138
node
@@ -449,24 +453,28 @@ def return_handler(
449453
saved_function_call_index(int): Unique number for each call.
450454
first_node(EntryOrExitNode or RestoreNode): Used to connect previous statements to this function.
451455
"""
452-
for node in function_nodes:
456+
if any(isinstance(node, YieldNode) for node in function_nodes):
457+
# Presence of a `YieldNode` means that the function is a generator
458+
rhs_prefix = 'yld_'
459+
elif any(isinstance(node, ConnectToExitNode) for node in function_nodes):
453460
# Only `Return`s and `Raise`s can be of type ConnectToExitNode
454-
if isinstance(node, ConnectToExitNode):
455-
# Create e.g. ~call_1 = ret_func_foo RestoreNode
456-
LHS = CALL_IDENTIFIER + 'call_' + str(saved_function_call_index)
457-
RHS = 'ret_' + get_call_names_as_string(call_node.func)
458-
return_node = RestoreNode(
459-
LHS + ' = ' + RHS,
460-
LHS,
461-
[RHS],
462-
line_number=call_node.lineno,
463-
path=self.filenames[-1]
464-
)
465-
return_node.first_node = first_node
461+
rhs_prefix = 'ret_'
462+
else:
463+
return # No return value
466464

467-
self.nodes[-1].connect(return_node)
468-
self.nodes.append(return_node)
469-
return
465+
# Create e.g. ~call_1 = ret_func_foo RestoreNode
466+
LHS = CALL_IDENTIFIER + 'call_' + str(saved_function_call_index)
467+
RHS = rhs_prefix + get_call_names_as_string(call_node.func)
468+
return_node = RestoreNode(
469+
LHS + ' = ' + RHS,
470+
LHS,
471+
[RHS],
472+
line_number=call_node.lineno,
473+
path=self.filenames[-1]
474+
)
475+
return_node.first_node = first_node
476+
self.nodes[-1].connect(return_node)
477+
self.nodes.append(return_node)
470478

471479
def process_function(self, call_node, definition):
472480
"""Processes a user defined function when it is called.

‎pyt/cfg/stmt_visitor.py‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -244,12 +244,6 @@ def visit_Return(self, node):
244244
label = LabelVisitor()
245245
label.visit(node)
246246

247-
try:
248-
rhs_visitor = RHSVisitor()
249-
rhs_visitor.visit(node.value)
250-
except AttributeError:
251-
rhs_visitor.result = 'EmptyReturn'
252-
253247
this_function_name = self.function_return_stack[-1]
254248
LHS = 'ret_' + this_function_name
255249

@@ -263,14 +257,17 @@ def visit_Return(self, node):
263257
path=self.filenames[-1]
264258
)
265259
return_value_of_call.connect(return_node)
266-
self.nodes.append(return_node)
267-
return return_node
260+
return self.append_node(return_node)
261+
elif node.value is not None:
262+
rhs_visitor_result = RHSVisitor.result_for_node(node.value)
263+
else:
264+
rhs_visitor_result = []
268265

269266
return self.append_node(ReturnNode(
270267
LHS + ' = ' + label.result,
271268
LHS,
272269
node,
273-
rhs_visitor.result,
270+
rhs_visitor_result,
274271
path=self.filenames[-1]
275272
))
276273

@@ -502,11 +499,12 @@ def visit_AugAssign(self, node):
502499
rhs_visitor = RHSVisitor()
503500
rhs_visitor.visit(node.value)
504501

502+
lhs = extract_left_hand_side(node.target)
505503
return self.append_node(AssignmentNode(
506504
label.result,
507-
extract_left_hand_side(node.target),
505+
lhs,
508506
node,
509-
rhs_visitor.result,
507+
rhs_visitor.result + [lhs],
510508
path=self.filenames[-1]
511509
))
512510

‎pyt/core/README.rst‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ This directory contains miscellaneous code that is imported from different parts
99

1010
- `get_call_names`_ used in `vars_visitor.py`_ when visiting a Subscript, and `framework_helper.py`_ on function decorators in `is_flask_route_function`_
1111

12-
- `get_call_names_as_string`_ used in `expr_visitor.py`_ to create ret_function_name as RHS and yield_function_name as LHS, and in stmt_visitor.py when connecting a function to a loop.
12+
- `get_call_names_as_string`_ used in `expr_visitor.py`_ to create ret_function_name as RHS and yld_function_name as LHS, and in stmt_visitor.py when connecting a function to a loop.
1313

1414
- `Arguments`_ used in `expr_visitor.py`_ when processing the arguments of a user defined function and `framework_adaptor.py`_ to taint function definition arguments.
1515

‎pyt/core/node_types.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,3 +285,11 @@ def __init__(
285285
line_number=ast_node.lineno,
286286
path=path
287287
)
288+
289+
290+
class YieldNode(AssignmentNode):
291+
"""CFG Node that represents a yield or yield from.
292+
293+
The presence of a YieldNode means that a function is a generator.
294+
"""
295+
pass

‎pyt/vulnerabilities/vulnerability_helper.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
from enum import Enum
55
from collections import namedtuple
66

7+
from ..core.node_types import YieldNode
8+
79

810
class VulnerabilityType(Enum):
911
FALSE = 0
@@ -56,13 +58,20 @@ def __init__(
5658

5759
self.reassignment_nodes = reassignment_nodes
5860
self._remove_sink_from_secondary_nodes()
61+
self._remove_non_propagating_yields()
5962

6063
def _remove_sink_from_secondary_nodes(self):
6164
try:
6265
self.reassignment_nodes.remove(self.sink)
6366
except ValueError: # pragma: no cover
6467
pass
6568

69+
def _remove_non_propagating_yields(self):
70+
"""Remove yield with no variables e.g. `yield 123` and plain `yield` from vulnerability."""
71+
for node in list(self.reassignment_nodes):
72+
if isinstance(node, YieldNode) and len(node.right_hand_side_variables) == 1:
73+
self.reassignment_nodes.remove(node)
74+
6675
def __str__(self):
6776
"""Pretty printing of a vulnerability."""
6877
reassigned_str = _get_reassignment_str(self.reassignment_nodes)

‎tests/cfg/cfg_test.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -820,6 +820,14 @@ def test_assignment_starred_list(self):
820820
[('a', ['d']), ('b', ['d']), ('c', ['e'])],
821821
)
822822

823+
def test_augmented_assignment(self):
824+
self.cfg_create_from_ast(ast.parse('a+=f(b,c)'))
825+
826+
(node,) = self.cfg.nodes[1:-1]
827+
self.assertEqual(node.label, 'a += f(b, c)')
828+
self.assertEqual(node.left_hand_side, 'a')
829+
self.assertEqual(node.right_hand_side_variables, ['b', 'c', 'a'])
830+
823831

824832
class CFGComprehensionTest(CFGBaseTestCase):
825833
def test_nodes(self):
@@ -995,6 +1003,42 @@ def test_function_multiple_return(self):
9951003
(call_foo, exit_foo),
9961004
(_exit, call_foo)])
9971005

1006+
def test_generator_multiple_yield(self):
1007+
path = 'examples/example_inputs/generator_with_multiple_yields.py'
1008+
self.cfg_create_from_file(path)
1009+
1010+
self.assert_length(self.cfg.nodes, expected_length=9)
1011+
1012+
entry = 0
1013+
entry_foo = 1
1014+
a = 2
1015+
_if = 3
1016+
yld_if = 4
1017+
yld = 5
1018+
exit_foo = 6
1019+
call_foo = 7
1020+
_exit = 8
1021+
1022+
self.assertInCfg([
1023+
(entry_foo, entry),
1024+
(a, entry_foo),
1025+
(_if, a),
1026+
(yld_if, _if),
1027+
(yld, _if),
1028+
(yld, yld_if), # Different from return
1029+
(exit_foo, yld),
1030+
(call_foo, exit_foo),
1031+
(_exit, call_foo)
1032+
])
1033+
1034+
yld_if_node = self.cfg.nodes[yld_if]
1035+
self.assertEqual(yld_if_node.left_hand_side, 'yld_foo')
1036+
self.assertEqual(yld_if_node.right_hand_side_variables, ['yld_foo'])
1037+
yld_node = self.cfg.nodes[yld]
1038+
self.assertEqual(yld_node.left_hand_side, 'yld_foo')
1039+
self.assertEqual(yld_node.right_hand_side_variables, ['a', 'yld_foo'])
1040+
self.assertEqual(self.cfg.nodes[call_foo].right_hand_side_variables, ['yld_foo'])
1041+
9981042
def test_blackbox_call_after_if(self):
9991043
path = 'examples/vulnerable_code/blackbox_call_after_if.py'
10001044
self.cfg_create_from_file(path)

‎tests/main_test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,11 +99,11 @@ def test_targets_with_recursive(self):
9999
excluded_files = ""
100100

101101
included_files = discover_files(targets, excluded_files, True)
102-
self.assertEqual(len(included_files), 30)
102+
self.assertEqual(len(included_files), 31)
103103

104104
def test_targets_with_recursive_and_excluded(self):
105105
targets = ["examples/vulnerable_code/"]
106106
excluded_files = "inter_command_injection.py"
107107

108108
included_files = discover_files(targets, excluded_files, True)
109-
self.assertEqual(len(included_files), 29)
109+
self.assertEqual(len(included_files), 30)

‎tests/vulnerabilities/vulnerabilities_test.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,31 @@ def test_XSS_variable_multiple_assign_result(self):
492492

493493
self.assertAlphaEqual(vulnerability_description, EXPECTED_VULNERABILITY_DESCRIPTION)
494494

495+
def test_yield(self):
496+
vulnerabilities = self.run_analysis('examples/vulnerable_code/yield.py')
497+
self.assert_length(vulnerabilities, expected_length=1)
498+
vuln = vulnerabilities[0]
499+
self.assertEqual(vuln.source.left_hand_side, "yld_things_to_run")
500+
self.assertIn("yld_things_to_run", vuln.source.right_hand_side_variables)
501+
EXPECTED_VULNERABILITY_DESCRIPTION = """
502+
File: examples/vulnerable_code/yield.py
503+
> User input at line 9, source "request.get_json(":
504+
yld_things_to_run += request.get_json()['commands']
505+
Reassigned in:
506+
File: examples/vulnerable_code/yield.py
507+
> Line 15: ~call_2 = yld_things_to_run
508+
File: examples/vulnerable_code/yield.py
509+
> Line 15: ~call_1 = ret_'; '.join(~call_2)
510+
File: examples/vulnerable_code/yield.py
511+
> Line 15: script = ~call_1
512+
File: examples/vulnerable_code/yield.py
513+
> reaches line 16, sink "subprocess.call(":
514+
~call_3 = ret_subprocess.call(script, shell=True)
515+
This vulnerability is unknown due to: Label: ~call_1 = ret_'; '.join(~call_2)
516+
"""
517+
518+
self.assertAlphaEqual(str(vuln), EXPECTED_VULNERABILITY_DESCRIPTION)
519+
495520

496521
class EngineDjangoTest(VulnerabilitiesBaseTestCase):
497522
def run_analysis(self, path):

0 commit comments

Comments
 (0)