Skip to content

Commit 993c151

Browse files
clavedelunaclavedeluna
committed
harden-pyyaml can detect inheriting unsafe loaders
1 parent 48a7831 commit 993c151

3 files changed

Lines changed: 193 additions & 15 deletions

File tree

‎src/codemodder/codemods/api/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,8 @@ def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
135135

136136
def leave_Assign(self, original_node, updated_node):
137137
return self._new_or_updated_node(original_node, updated_node)
138+
139+
def leave_ClassDef(
140+
self, original_node: cst.ClassDef, updated_node: cst.ClassDef
141+
) -> cst.ClassDef:
142+
return self._new_or_updated_node(original_node, updated_node)

‎src/core_codemods/harden_pyyaml.py‎

Lines changed: 67 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Union
2+
import libcst as cst
13
from codemodder.codemods.base_codemod import ReviewGuidance
24
from codemodder.codemods.api import SemgrepCodemod
35
from codemodder.codemods.utils_mixin import NameResolutionMixin
@@ -6,8 +8,8 @@
68
class HardenPyyaml(SemgrepCodemod, NameResolutionMixin):
79
NAME = "harden-pyyaml"
810
REVIEW_GUIDANCE = ReviewGuidance.MERGE_WITHOUT_REVIEW
9-
SUMMARY = "Use SafeLoader in `yaml.load()` Calls"
10-
DESCRIPTION = "Ensures all calls to yaml.load use `SafeLoader`."
11+
SUMMARY = "Replace unsafe `pyyaml` loader with `SafeLoader`"
12+
DESCRIPTION = "Replace unsafe `pyyaml` loader with `SafeLoader` in calls to `yaml.load` or custom loader classes."
1113
REFERENCES = [
1214
{
1315
"url": "https://owasp.org/www-community/vulnerabilities/Deserialization_of_untrusted_data",
@@ -50,18 +52,70 @@ def rule(cls):
5052
- pattern: yaml.BaseLoader
5153
- pattern: yaml.FullLoader
5254
- pattern: yaml.UnsafeLoader
55+
- patterns:
56+
- pattern: |
57+
class $X(...,$LOADER, ...):
58+
...
59+
- metavariable-pattern:
60+
metavariable: $LOADER
61+
patterns:
62+
- pattern-either:
63+
- pattern: yaml.Loader
64+
- pattern: yaml.BaseLoader
65+
- pattern: yaml.FullLoader
66+
- pattern: yaml.UnsafeLoader
5367
5468
"""
5569

56-
def on_result_found(self, original_node, updated_node):
57-
maybe_name = self.get_aliased_prefix_name(original_node, self._module_name)
58-
maybe_name = maybe_name or self._module_name
59-
if maybe_name == self._module_name:
60-
self.add_needed_import(self._module_name)
61-
new_args = [
62-
*updated_node.args[:1],
63-
updated_node.args[1].with_changes(
64-
value=self.parse_expression(f"{maybe_name}.SafeLoader")
65-
),
70+
def on_result_found(
71+
self,
72+
original_node: Union[cst.Call, cst.ClassDef],
73+
updated_node: Union[cst.Call, cst.ClassDef],
74+
):
75+
# TODO: provide different change description for each case.
76+
match original_node:
77+
case cst.Call():
78+
maybe_name = self.get_aliased_prefix_name(
79+
original_node, self._module_name
80+
)
81+
maybe_name = maybe_name or self._module_name
82+
if maybe_name == self._module_name:
83+
self.add_needed_import(self._module_name)
84+
new_args = [
85+
*updated_node.args[:1],
86+
updated_node.args[1].with_changes(
87+
value=self.parse_expression(f"{maybe_name}.SafeLoader")
88+
),
89+
]
90+
return self.update_arg_target(updated_node, new_args)
91+
case cst.ClassDef():
92+
return updated_node.with_changes(
93+
bases=self._update_bases(original_node)
94+
)
95+
return updated_node
96+
97+
def _update_bases(self, original_node: cst.ClassDef) -> list[cst.Arg]:
98+
new = []
99+
base_names = [
100+
f"yaml.{klas}"
101+
for klas in ("UnsafeLoader", "Loader", "BaseLoader", "FullLoader")
66102
]
67-
return self.update_arg_target(updated_node, new_args)
103+
for base_arg in original_node.bases:
104+
base_name = self.find_base_name(base_arg.value)
105+
if base_name not in base_names:
106+
new.append(base_arg)
107+
continue
108+
109+
match base_arg.value:
110+
case cst.Name():
111+
self.add_needed_import(self._module_name, "SafeLoader")
112+
self.remove_unused_import(base_arg.value)
113+
base_arg = base_arg.with_changes(
114+
value=base_arg.value.with_changes(value="SafeLoader")
115+
)
116+
case cst.Attribute():
117+
base_arg = base_arg.with_changes(
118+
value=base_arg.value.with_changes(attr=cst.Name("SafeLoader"))
119+
)
120+
new.append(base_arg)
121+
return new

‎tests/codemods/test_harden_pyyaml.py‎

Lines changed: 121 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
UNSAFE_LOADERS = yaml.loader.__all__.copy() # type: ignore
77
UNSAFE_LOADERS.remove("SafeLoader")
8+
loaders = pytest.mark.parametrize("loader", UNSAFE_LOADERS)
89

910

1011
class TestHardenPyyaml(BaseSemgrepCodemodTest):
@@ -19,8 +20,9 @@ def test_safe_loader(self, tmpdir):
1920
deserialized_data = yaml.load(data, Loader=yaml.SafeLoader)
2021
"""
2122
self.run_and_assert(tmpdir, input_code, input_code)
23+
assert len(self.file_context.codemod_changes) == 0
2224

23-
@pytest.mark.parametrize("loader", UNSAFE_LOADERS)
25+
@loaders
2426
def test_all_unsafe_loaders_arg(self, tmpdir, loader):
2527
input_code = f"""import yaml
2628
data = b'!!python/object/apply:subprocess.Popen \\n- ls'
@@ -32,8 +34,9 @@ def test_all_unsafe_loaders_arg(self, tmpdir, loader):
3234
deserialized_data = yaml.load(data, yaml.SafeLoader)
3335
"""
3436
self.run_and_assert(tmpdir, input_code, expected)
37+
assert len(self.file_context.codemod_changes) == 1
3538

36-
@pytest.mark.parametrize("loader", UNSAFE_LOADERS)
39+
@loaders
3740
def test_all_unsafe_loaders_kwarg(self, tmpdir, loader):
3841
input_code = f"""import yaml
3942
data = b'!!python/object/apply:subprocess.Popen \\n- ls'
@@ -45,6 +48,7 @@ def test_all_unsafe_loaders_kwarg(self, tmpdir, loader):
4548
deserialized_data = yaml.load(data, Loader=yaml.SafeLoader)
4649
"""
4750
self.run_and_assert(tmpdir, input_code, expected)
51+
assert len(self.file_context.codemod_changes) == 1
4852

4953
def test_import_alias(self, tmpdir):
5054
input_code = """import yaml as yam
@@ -60,6 +64,7 @@ def test_import_alias(self, tmpdir):
6064
deserialized_data = yam.load(data, Loader=yam.SafeLoader)
6165
"""
6266
self.run_and_assert(tmpdir, input_code, expected)
67+
assert len(self.file_context.codemod_changes) == 1
6368

6469
def test_preserve_custom_loader(self, tmpdir):
6570
expected = input_code = """
@@ -70,6 +75,7 @@ def test_preserve_custom_loader(self, tmpdir):
7075
"""
7176

7277
self.run_and_assert(tmpdir, input_code, expected)
78+
assert len(self.file_context.codemod_changes) == 0
7379

7480
def test_preserve_custom_loader_kwarg(self, tmpdir):
7581
expected = input_code = """
@@ -80,3 +86,116 @@ def test_preserve_custom_loader_kwarg(self, tmpdir):
8086
"""
8187

8288
self.run_and_assert(tmpdir, input_code, expected)
89+
assert len(self.file_context.codemod_changes) == 0
90+
91+
92+
class TestHardenPyyamlClassInherit(BaseSemgrepCodemodTest):
93+
codemod = HardenPyyaml
94+
95+
def test_safe_loader(self, tmpdir):
96+
input_code = """\
97+
import yaml
98+
99+
class MyCustomLoader(yaml.SafeLoader):
100+
def __init__(self, *args, **kwargs):
101+
super(MyCustomLoader, self).__init__(*args, **kwargs)
102+
103+
"""
104+
105+
self.run_and_assert(tmpdir, input_code, input_code)
106+
assert len(self.file_context.codemod_changes) == 0
107+
108+
@loaders
109+
def test_unsafe_loaders(self, tmpdir, loader):
110+
input_code = f"""\
111+
import yaml
112+
113+
class MyCustomLoader(yaml.{loader}):
114+
def __init__(self, *args, **kwargs):
115+
super(MyCustomLoader, self).__init__(*args, **kwargs)
116+
"""
117+
expected = """\
118+
import yaml
119+
120+
class MyCustomLoader(yaml.SafeLoader):
121+
def __init__(self, *args, **kwargs):
122+
super(MyCustomLoader, self).__init__(*args, **kwargs)
123+
"""
124+
self.run_and_assert(tmpdir, input_code, expected)
125+
assert len(self.file_context.codemod_changes) == 1
126+
127+
def test_from_import(self, tmpdir):
128+
input_code = """\
129+
from yaml import UnsafeLoader
130+
131+
class MyCustomLoader(UnsafeLoader):
132+
def __init__(self, *args, **kwargs):
133+
super(MyCustomLoader, self).__init__(*args, **kwargs)
134+
"""
135+
expected = """\
136+
from yaml import SafeLoader
137+
138+
class MyCustomLoader(SafeLoader):
139+
def __init__(self, *args, **kwargs):
140+
super(MyCustomLoader, self).__init__(*args, **kwargs)
141+
"""
142+
self.run_and_assert(tmpdir, input_code, expected)
143+
assert len(self.file_context.codemod_changes) == 1
144+
145+
def test_import_alias(self, tmpdir):
146+
input_code = """\
147+
import yaml as yam
148+
149+
class MyCustomLoader(yam.UnsafeLoader):
150+
def __init__(self, *args, **kwargs):
151+
super(MyCustomLoader, self).__init__(*args, **kwargs)
152+
"""
153+
expected = """\
154+
import yaml as yam
155+
156+
class MyCustomLoader(yam.SafeLoader):
157+
def __init__(self, *args, **kwargs):
158+
super(MyCustomLoader, self).__init__(*args, **kwargs)
159+
"""
160+
self.run_and_assert(tmpdir, input_code, expected)
161+
assert len(self.file_context.codemod_changes) == 1
162+
163+
def test_multiple_bases(self, tmpdir):
164+
input_code = """\
165+
from abc import ABC
166+
import yaml as yam
167+
from whatever import Loader
168+
169+
class MyCustomLoader(ABC, yam.UnsafeLoader, Loader):
170+
def __init__(self, *args, **kwargs):
171+
super(MyCustomLoader, self).__init__(*args, **kwargs)
172+
"""
173+
expected = """\
174+
from abc import ABC
175+
import yaml as yam
176+
from whatever import Loader
177+
178+
class MyCustomLoader(ABC, yam.SafeLoader, Loader):
179+
def __init__(self, *args, **kwargs):
180+
super(MyCustomLoader, self).__init__(*args, **kwargs)
181+
"""
182+
self.run_and_assert(tmpdir, input_code, expected)
183+
assert len(self.file_context.codemod_changes) == 1
184+
185+
def test_different_yaml(self, tmpdir):
186+
input_code = """\
187+
from yaml import UnsafeLoader
188+
import whatever as yaml
189+
190+
class MyLoader(UnsafeLoader, yaml.Loader):
191+
...
192+
"""
193+
expected = """\
194+
from yaml import SafeLoader
195+
import whatever as yaml
196+
197+
class MyLoader(SafeLoader, yaml.Loader):
198+
...
199+
"""
200+
self.run_and_assert(tmpdir, input_code, expected)
201+
assert len(self.file_context.codemod_changes) == 1

0 commit comments

Comments
 (0)