Semgrep harden pyyaml (#710)

* semgrep from core codemod to allow list of rules

* new harden pyyaml semgrep codemod

src/codemodder/scripts/generate_docs.py

@@ -325,6 +325,7 @@ class DocMetadata:
     "jwt-decode-verify",
     "use-defusedxml",
     "subprocess-shell-false",
+    "harden-pyyaml",
 ]
 SEMGREP_CODEMODS = {
     name: DocMetadata(

src/core_codemods/__init__.py

@@ -56,6 +56,7 @@
 from .secure_random import SecureRandom
 from .semgrep.semgrep_django_secure_set_cookie import SemgrepDjangoSecureSetCookie
 from .semgrep.semgrep_enable_jinja2_autoescape import SemgrepEnableJinja2Autoescape
+from .semgrep.semgrep_harden_pyyaml import SemgrepHardenPyyaml
 from .semgrep.semgrep_jwt_decode_verify import SemgrepJwtDecodeVerify
 from .semgrep.semgrep_subprocess_shell_false import SemgrepSubprocessShellFalse
 from .semgrep.semgrep_use_defused_xml import SemgrepUseDefusedXml
@@ -206,5 +207,6 @@
         SemgrepUseDefusedXml,
         SemgrepSubprocessShellFalse,
         SemgrepDjangoSecureSetCookie,
+        SemgrepHardenPyyaml,
     ],
 )

src/core_codemods/semgrep/api.py

@@ -20,8 +20,7 @@ def from_core_codemod(
         cls,
         name: str,
         other: CoreCodemod,
-        rule_id: str,
-        rule_name: str,
+        rules: list[ToolRule],
         transformer: BaseTransformerPipeline | None = None,
     ):
         return SemgrepCodemod(
@@ -33,25 +32,20 @@ def from_core_codemod(
                     other.references
                     + [
                         Reference(
-                            url=semgrep_url_from_id(rule_id), description=rule_name
+                            url=semgrep_url_from_id(rule.id), description=rule.name
                         )
+                        for rule in rules
                     ]
                 ),
                 description=other.description,
                 tool=ToolMetadata(
                     name="Semgrep",
-                    rules=[
-                        ToolRule(
-                            id=rule_id,
-                            name=rule_name,
-                            url=semgrep_url_from_id(rule_id),
-                        )
-                    ],
+                    rules=rules,
                 ),
             ),
             transformer=transformer if transformer else other.transformer,
             detector=SemgrepSarifFileDetector(),
-            requested_rules=[rule_id],
+            requested_rules=[rule.id for rule in rules],
         )
 
     @classmethod

src/core_codemods/semgrep/semgrep_django_secure_set_cookie.py

@@ -1,11 +1,19 @@
+from codemodder.codemods.base_codemod import ToolRule
 from core_codemods.defectdojo.semgrep.django_secure_set_cookie import (
     DjangoSecureSetCookie,
 )
-from core_codemods.semgrep.api import SemgrepCodemod
+from core_codemods.semgrep.api import SemgrepCodemod, semgrep_url_from_id
 
 SemgrepDjangoSecureSetCookie = SemgrepCodemod.from_core_codemod(
     name="django-secure-set-cookie",
     other=DjangoSecureSetCookie,
-    rule_id="python.django.security.audit.secure-cookies.django-secure-set-cookie",
-    rule_name="django-secure-set-cookie",
+    rules=[
+        ToolRule(
+            id=(
+                rule_id := "python.django.security.audit.secure-cookies.django-secure-set-cookie"
+            ),
+            name="django-secure-set-cookie",
+            url=semgrep_url_from_id(rule_id),
+        )
+    ],
 )

src/core_codemods/semgrep/semgrep_enable_jinja2_autoescape.py

@@ -1,9 +1,17 @@
+from codemodder.codemods.base_codemod import ToolRule
 from core_codemods.enable_jinja2_autoescape import EnableJinja2Autoescape
-from core_codemods.semgrep.api import SemgrepCodemod
+from core_codemods.semgrep.api import SemgrepCodemod, semgrep_url_from_id
 
 SemgrepEnableJinja2Autoescape = SemgrepCodemod.from_core_codemod(
     name="enable-jinja2-autoescape",
     other=EnableJinja2Autoescape,
-    rule_id="python.flask.security.xss.audit.direct-use-of-jinja2.direct-use-of-jinja2",
-    rule_name="direct-use-of-jinja2",
+    rules=[
+        ToolRule(
+            id=(
+                rule_id := "python.flask.security.xss.audit.direct-use-of-jinja2.direct-use-of-jinja2"
+            ),
+            name="direct-use-of-jinja2",
+            url=semgrep_url_from_id(rule_id),
+        )
+    ],
 )

src/core_codemods/semgrep/semgrep_harden_pyyaml.py

@@ -0,0 +1,24 @@
+from codemodder.codemods.base_codemod import ToolRule
+from core_codemods.harden_pyyaml import HardenPyyaml
+from core_codemods.semgrep.api import SemgrepCodemod, semgrep_url_from_id
+
+SemgrepHardenPyyaml = SemgrepCodemod.from_core_codemod(
+    name="harden-pyyaml",
+    other=HardenPyyaml,
+    rules=[
+        ToolRule(
+            id=(
+                rule_id := "python.lang.security.deserialization.avoid-pyyaml-load.avoid-pyyaml-load"
+            ),
+            name=" avoid-pyyaml-load",
+            url=semgrep_url_from_id(rule_id),
+        ),
+        ToolRule(
+            id=(
+                rule_id := "python.django.security.audit.avoid-insecure-deserialization.avoid-insecure-deserialization"
+            ),
+            name="avoid-insecure-deserialization",
+            url=semgrep_url_from_id(rule_id),
+        ),
+    ],
+)

src/core_codemods/semgrep/semgrep_jwt_decode_verify.py

@@ -1,14 +1,22 @@
+from codemodder.codemods.base_codemod import ToolRule
 from codemodder.codemods.libcst_transformer import LibcstTransformerPipeline
 from core_codemods.jwt_decode_verify import (
     JwtDecodeVerify,
     JwtDecodeVerifySASTTransformer,
 )
-from core_codemods.semgrep.api import SemgrepCodemod
+from core_codemods.semgrep.api import SemgrepCodemod, semgrep_url_from_id
 
 SemgrepJwtDecodeVerify = SemgrepCodemod.from_core_codemod(
     name="jwt-decode-verify",
     other=JwtDecodeVerify,
-    rule_id="python.jwt.security.unverified-jwt-decode.unverified-jwt-decode",
-    rule_name="unverified-jwt-decode",
+    rules=[
+        ToolRule(
+            id=(
+                rule_id := "python.jwt.security.unverified-jwt-decode.unverified-jwt-decode"
+            ),
+            name="unverified-jwt-decode",
+            url=semgrep_url_from_id(rule_id),
+        )
+    ],
     transformer=LibcstTransformerPipeline(JwtDecodeVerifySASTTransformer),
 )

src/core_codemods/semgrep/semgrep_subprocess_shell_false.py

@@ -1,9 +1,17 @@
-from core_codemods.semgrep.api import SemgrepCodemod
+from codemodder.codemods.base_codemod import ToolRule
+from core_codemods.semgrep.api import SemgrepCodemod, semgrep_url_from_id
 from core_codemods.subprocess_shell_false import SubprocessShellFalse
 
 SemgrepSubprocessShellFalse = SemgrepCodemod.from_core_codemod(
     name="subprocess-shell-false",
     other=SubprocessShellFalse,
-    rule_id="python.lang.security.audit.subprocess-shell-true.subprocess-shell-true",
-    rule_name="subprocess-shell-true",
+    rules=[
+        ToolRule(
+            id=(
+                rule_id := "python.lang.security.audit.subprocess-shell-true.subprocess-shell-true"
+            ),
+            name="subprocess-shell-true",
+            url=semgrep_url_from_id(rule_id),
+        )
+    ],
 )

tests/codemods/semgrep/test_semgrep_harden_pyyaml.py

@@ -0,0 +1,124 @@
+import json
+
+from codemodder.codemods.test import BaseSASTCodemodTest
+from core_codemods.semgrep.semgrep_harden_pyyaml import SemgrepHardenPyyaml
+
+
+class TestSemgrepHardenPyyaml(BaseSASTCodemodTest):
+    codemod = SemgrepHardenPyyaml
+    tool = "semgrep"
+
+    def test_name(self):
+        assert self.codemod.name == "harden-pyyaml"
+
+    def test_pyyaml(self, tmpdir):
+        input_code = """\
+        import yaml
+        data = b'!!python/object/apply:subprocess.Popen \\n- ls'
+        deserialized_data = yaml.load(data, Loader=yaml.Loader)
+        """
+        expected_output = """\
+        import yaml
+        data = b'!!python/object/apply:subprocess.Popen \\n- ls'
+        deserialized_data = yaml.load(data, Loader=yaml.SafeLoader)
+        """
+
+        results = {
+            "runs": [
+                {
+                    "results": [
+                        {
+                            "fingerprints": {"matchBasedId/v1": "123"},
+                            "locations": [
+                                {
+                                    "physicalLocation": {
+                                        "artifactLocation": {
+                                            "uri": "code.py",
+                                            "uriBaseId": "%SRCROOT%",
+                                        },
+                                        "region": {
+                                            "endColumn": 56,
+                                            "endLine": 3,
+                                            "snippet": {
+                                                "text": "deserialized_data = yaml.load(data, Loader=yaml.Loader)"
+                                            },
+                                            "startColumn": 21,
+                                            "startLine": 3,
+                                        },
+                                    }
+                                }
+                            ],
+                            "message": {
+                                "text": "Detected a possible YAML deserialization vulnerability. `yaml.unsafe_load`, `yaml.Loader`, `yaml.CLoader`, and `yaml.UnsafeLoader` are all known to be unsafe methods of deserializing YAML. An attacker with control over the YAML input could create special YAML input that allows the attacker to run arbitrary Python code. This would allow the attacker to steal files, download and install malware, or otherwise take over the machine. Use `yaml.safe_load` or `yaml.SafeLoader` instead."
+                            },
+                            "properties": {},
+                            "ruleId": "python.lang.security.deserialization.avoid-pyyaml-load.avoid-pyyaml-load",
+                        }
+                    ]
+                }
+            ]
+        }
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected_output,
+            results=json.dumps(results),
+        )
+
+    def test_pyyaml_django(self, tmpdir):
+        input_code = """\
+        import yaml
+        
+        def index(request):
+            cookie = request.cookies.get('cookie')
+            return "Hey there! {}!".format(yaml.load(cookie))
+        """
+        expected_output = """\
+        import yaml
+        
+        def index(request):
+            cookie = request.cookies.get('cookie')
+            return "Hey there! {}!".format(yaml.load(cookie, Loader=yaml.SafeLoader))
+        """
+
+        results = {
+            "runs": [
+                {
+                    "results": [
+                        {
+                            "fingerprints": {"matchBasedId/v1": "123"},
+                            "locations": [
+                                {
+                                    "physicalLocation": {
+                                        "artifactLocation": {
+                                            "uri": "code.py",
+                                            "uriBaseId": "%SRCROOT%",
+                                        },
+                                        "region": {
+                                            "endColumn": 53,
+                                            "endLine": 5,
+                                            "snippet": {
+                                                "text": '    return "Hey there! {}!".format(yaml.load(cookie))'
+                                            },
+                                            "startColumn": 36,
+                                            "startLine": 5,
+                                        },
+                                    }
+                                }
+                            ],
+                            "message": {
+                                "text": "Avoid using insecure deserialization library, backed by `pickle`, `_pickle`, `cpickle`, `dill`, `shelve`, or `yaml`, which are known to lead to remote code execution vulnerabilities."
+                            },
+                            "properties": {},
+                            "ruleId": "python.django.security.audit.avoid-insecure-deserialization.avoid-insecure-deserialization",
+                        }
+                    ]
+                }
+            ]
+        }
+        self.run_and_assert(
+            tmpdir,
+            input_code,
+            expected_output,
+            results=json.dumps(results),
+        )