Docs and tests for flask-send-file-path-parameterization

integration_tests/base_test.py

@@ -204,7 +204,7 @@ def _lines_from_codepath(code_path):
 def _replace_lines_with(lines, replacements):
     total_lines = len(lines)
     for lineno, replacement in replacements:
-        if lineno > total_lines:
+        if lineno >= total_lines:
             lines.extend(replacement)
             continue
         lines[lineno] = replacement

integration_tests/test_flask_send_file_path_parameterization.py

@@ -0,0 +1,53 @@
+from core_codemods.flask_send_file_path_parameterization import (
+    FlaskSendFilePathParameterization,
+)
+from integration_tests.base_test import (
+    BaseIntegrationTest,
+    original_and_expected_from_code_path,
+)
+
+
+class TestFlaskSendFilePathParameterization(BaseIntegrationTest):
+    codemod = FlaskSendFilePathParameterization
+    code_path = "tests/samples/flask_send_file_path_parameterization.py"
+    original_code, expected_new_code = original_and_expected_from_code_path(
+        code_path,
+        [
+            (0, """from flask import Flask\n"""),
+            (1, """import flask\n"""),
+            (2, """from pathlib import Path\n"""),
+            (3, """\n"""),
+            (4, """app = Flask(__name__)\n"""),
+            (5, """\n"""),
+            (6, """@app.route("/uploads/<path:name>")\n"""),
+            (7, """def download_file(name):\n"""),
+            (
+                8,
+                """    return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name)\n""",
+            ),
+        ],
+    )
+
+    # fmt: off
+    expected_diff =(
+    """--- \n"""
+    """+++ \n"""
+    """@@ -1,7 +1,9 @@\n"""
+    """-from flask import Flask, send_file\n"""
+    """+from flask import Flask\n"""
+    """+import flask\n"""
+    """+from pathlib import Path\n"""
+    """ \n"""
+    """ app = Flask(__name__)\n"""
+    """ \n"""
+    """ @app.route("/uploads/<path:name>")\n"""
+    """ def download_file(name):\n"""
+    """-    return send_file(f'path/to/{name}.txt')\n"""
+    """+    return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name)\n"""
+
+    )
+    # fmt: on
+
+    expected_line_change = "7"
+    change_description = FlaskSendFilePathParameterization.CHANGE_DESCRIPTION
+    num_changed_files = 1

src/codemodder/scripts/generate_docs.py

@@ -202,6 +202,10 @@ class DocMetadata:
         importance="Low",
         guidance_explained="This change fixes deprecated uses and is safe.",
     ),
+    "flask-send-file-path-parameterization": DocMetadata(
+        importance="Medium",
+        guidance_explained="We believe this change is safe and will not cause any issues.",
+    ),
 }
 
 

src/core_codemods/docs/pixee_python_flask-send-file-path-parameterization.md

@@ -0,0 +1,19 @@
+The `send_file` function from Flask is susceptible to a path injection attack if its input is not proper validated. 
+In a path injection attack, the malicious agent can craft a path containing special paths like `./` or `../` to resolve an unintended file. This allows the agent to overwrite, delete or read arbitrary files.
+Flask offers a native solution with the `flask.send_from_directory` function that validates the given path.
+
+Our changes look something like this:
+
+```diff
+-from flask import Flask, send_file
++from flask import Flask
++import flask
++from pathlib import Path
+
+app = Flask(__name__)
+
+@app.route("/uploads/<path:name>")
+def download_file(name):
+-    return send_file(f'path/to/{name}.txt')
++    return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name)
+```

src/core_codemods/flask_send_file_path_parameterization.py

@@ -49,12 +49,15 @@ def leave_Call(
                         ),
                     ]
                     self.report_change(original_node)
+                    self.add_needed_import("flask")
+                    self.remove_unused_import(original_node)
                     new_func = cst.parse_expression("flask.send_from_directory")
                     return updated_node.with_changes(func=new_func, args=new_args)
 
         return updated_node
 
     def _wrap_in_path(self, expr) -> cst.Call:
+        self.add_needed_import("pathlib", "Path")
         return cst.Call(func=cst.Name(value="Path"), args=[cst.Arg(expr)])
 
     def _attribute_reference(self, expr, attribute: str) -> cst.Attribute:

tests/codemods/test_flask_send_file_path_parameterization.py

@@ -0,0 +1,178 @@
+from core_codemods.flask_send_file_path_parameterization import (
+    FlaskSendFilePathParameterization,
+)
+from tests.codemods.base_codemod_test import BaseCodemodTest
+from textwrap import dedent
+
+
+class TestFlaskSendFilePathParameterization(BaseCodemodTest):
+    codemod = FlaskSendFilePathParameterization
+
+    def test_name(self):
+        assert self.codemod.name() == "flask-send-file-path-parameterization"
+
+    def test_direct_string(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return send_file(f'path/to/{name}.txt')
+        """
+        expected = """\
+        from flask import Flask
+        import flask
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_direct_string_convert_arguments(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return send_file(f'path/to/{name}.txt', None, False, download_name = True)
+        """
+        expected = """\
+        from flask import Flask
+        import flask
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name, mimetype = None, as_attachment = False, download_name = True)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_direct_path(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return send_file(Path(f'path/to/{name}.txt'))
+        """
+        expected = """\
+        from flask import Flask
+        from pathlib import Path
+        import flask
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return flask.send_from_directory((p := Path(f'path/to/{name}.txt')).parent, p.name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_indirect_path(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = Path(f'path/to/{name}.txt')
+            return send_file(path)
+        """
+        expected = """\
+        from flask import Flask
+        from pathlib import Path
+        import flask
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = Path(f'path/to/{name}.txt')
+            return flask.send_from_directory(path.parent, path.name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_indirect_path_alias(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file as send
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = Path(f'path/to/{name}.txt')
+            return send(path)
+        """
+        expected = """\
+        from flask import Flask
+        from pathlib import Path
+        import flask
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = Path(f'path/to/{name}.txt')
+            return flask.send_from_directory(path.parent, path.name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_indirect_string(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = f'path/to/{name}.txt'
+            return send_file(path)
+        """
+        expected = """\
+        from flask import Flask
+        import flask
+        from pathlib import Path
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            path = f'path/to/{name}.txt'
+            return flask.send_from_directory((p := Path(path)).parent, p.name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(expected))
+        assert len(self.file_context.codemod_changes) == 1
+
+    def test_unknown_type(self, tmpdir):
+        input_code = """\
+        from flask import Flask, send_file
+
+        app = Flask(__name__)
+
+        @app.route("/uploads/<path:name>")
+        def download_file(name):
+            return send_file(name)
+        """
+        self.run_and_assert(tmpdir, dedent(input_code), dedent(input_code))
+        assert len(self.file_context.codemod_changes) == 0

tests/samples/flask_send_file_path_parameterization.py

@@ -1,9 +1,7 @@
 from flask import Flask, send_file
-from pathlib import Path
 
 app = Flask(__name__)
 
 @app.route("/uploads/<path:name>")
 def download_file(name):
-    return send_file(f'path/to/{name}.txt', None, True, max_age=10)
-
+    return send_file(f'path/to/{name}.txt')