fix(#13734): Properly escape special characters in CSV output

superset/charts/api.py

@@ -484,12 +484,7 @@ def get_data_response(
         if result_format == ChartDataResultFormat.CSV:
             # return the first result
             data = result["queries"][0]["data"]
-            return CsvResponse(
-                data,
-                status=200,
-                headers=generate_download_headers("csv"),
-                mimetype="application/csv",
-            )
+            return CsvResponse(data, headers=generate_download_headers("csv"))
 
         if result_format == ChartDataResultFormat.JSON:
             response_data = simplejson.dumps(

superset/common/query_context.py

@@ -35,6 +35,7 @@
 )
 from superset.extensions import cache_manager, security_manager
 from superset.stats_logger import BaseStatsLogger
+from superset.utils import csv
 from superset.utils.cache import generate_cache_key, set_and_log_cache
 from superset.utils.core import (
     ChartDataResultFormat,
@@ -151,7 +152,9 @@ def df_metrics_to_num(df: pd.DataFrame, query_object: QueryObject) -> None:
     def get_data(self, df: pd.DataFrame,) -> Union[str, List[Dict[str, Any]]]:
         if self.result_format == ChartDataResultFormat.CSV:
             include_index = not isinstance(df.index, pd.RangeIndex)
-            result = df.to_csv(index=include_index, **config["CSV_EXPORT"])
+            result = csv.df_to_escaped_csv(
+                df, index=include_index, **config["CSV_EXPORT"]
+            )
             return result or ""
 
         return df.to_dict(orient="records")

superset/utils/csv.py

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import re
+from typing import Any
+
+import pandas as pd
+
+
+def escape_value(value: str) -> str:
+    """
+    Escapes a set of special characters.
+
+    Taken from: https://github.com/raphaelm/defusedcsv/blob/5911a6ae54a82e71fad4db5eb8769738e9620e31/defusedcsv/csv.py#L15-L23 # pylint: disable=line-too-long
+    """
+    if (
+        value
+        and value[0] in ("@", "+", "-", "=", "|", "%")
+        and not re.match("^-?[0-9,\\.]+$", value)
+    ):
+        value = value.replace("|", "\\|")
+        value = "'" + value
+    return value
+
+
+def df_to_escaped_csv(df: pd.DataFrame, **kwargs: Any) -> Any:
+    escape_values = lambda v: escape_value(v) if isinstance(v, str) else v
+
+    # Escape csv headers
+    df = df.rename(columns=escape_values)
+
+    # Escape csv rows
+    df = df.applymap(escape_values)
+
+    return df.to_csv(**kwargs)

superset/views/base.py

@@ -479,6 +479,7 @@ class CsvResponse(Response):  # pylint: disable=too-many-ancestors
     """
 
     charset = conf["CSV_EXPORT"].get("encoding", "utf-8")
+    default_mimetype = "text/csv"
 
 
 def check_ownership(obj: Any, raise_if_false: bool = True) -> bool:

superset/views/core.py

@@ -98,7 +98,7 @@
 from superset.sql_validators import get_validator_by_name
 from superset.tasks.async_queries import load_explore_json_into_cache
 from superset.typing import FlaskResponse
-from superset.utils import core as utils
+from superset.utils import core as utils, csv
 from superset.utils.async_query_manager import AsyncQueryTokenException
 from superset.utils.cache import etag_cache
 from superset.utils.core import ReservedUrlParameters
@@ -437,10 +437,7 @@ def generate_json(
     ) -> FlaskResponse:
         if response_type == utils.ChartDataResultFormat.CSV:
             return CsvResponse(
-                viz_obj.get_csv(),
-                status=200,
-                headers=generate_download_headers("csv"),
-                mimetype="application/csv",
+                viz_obj.get_csv(), headers=generate_download_headers("csv")
             )
 
         if response_type == utils.ChartDataResultType.QUERY:
@@ -2628,18 +2625,14 @@ def csv(self, client_id: str) -> FlaskResponse:  # pylint: disable=no-self-use
             columns = [c["name"] for c in obj["columns"]]
             df = pd.DataFrame.from_records(obj["data"], columns=columns)
             logger.info("Using pandas to convert to CSV")
-            csv = df.to_csv(index=False, **config["CSV_EXPORT"])
         else:
             logger.info("Running a query to turn into CSV")
             sql = query.select_sql or query.executed_sql
             df = query.database.get_df(sql, query.schema)
-            # TODO(bkyryliuk): add compression=gzip for big files.
-            csv = df.to_csv(index=False, **config["CSV_EXPORT"])
-        response = Response(csv, mimetype="text/csv")
+        csv_data = csv.df_to_escaped_csv(df, index=False, **config["CSV_EXPORT"])
         quoted_csv_name = parse.quote(query.name)
-        response.headers["Content-Disposition"] = (
-            f'attachment; filename="{quoted_csv_name}.csv"; '
-            f"filename*=UTF-8''{quoted_csv_name}.csv"
+        response = CsvResponse(
+            csv_data, headers=generate_download_headers("csv", quoted_csv_name)
         )
         event_info = {
             "event_type": "data_export",

superset/viz.py

@@ -66,7 +66,7 @@
 from superset.models.cache import CacheKey
 from superset.models.helpers import QueryResult
 from superset.typing import QueryObjectDict, VizData, VizPayload
-from superset.utils import core as utils
+from superset.utils import core as utils, csv
 from superset.utils.cache import set_and_log_cache
 from superset.utils.core import (
     DTTM_ALIAS,
@@ -615,7 +615,7 @@ def data(self) -> Dict[str, Any]:
     def get_csv(self) -> Optional[str]:
         df = self.get_df_payload()["df"]  # leverage caching logic
         include_index = not isinstance(df.index, pd.RangeIndex)
-        return df.to_csv(index=include_index, **config["CSV_EXPORT"])
+        return csv.df_to_escaped_csv(df, index=include_index, **config["CSV_EXPORT"])
 
     def get_data(self, df: pd.DataFrame) -> VizData:
         return df.to_dict(orient="records")

tests/utils/csv_tests.py

@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=no-self-use
+import io
+
+import pandas as pd
+import pytest
+
+from superset.utils import csv
+
+
+def test_escape_value():
+    result = csv.escape_value("value")
+    assert result == "value"
+
+    result = csv.escape_value("-10")
+    assert result == "-10"
+
+    result = csv.escape_value("@value")
+    assert result == "'@value"
+
+    result = csv.escape_value("+value")
+    assert result == "'+value"
+
+    result = csv.escape_value("-value")
+    assert result == "'-value"
+
+    result = csv.escape_value("=value")
+    assert result == "'=value"
+
+    result = csv.escape_value("|value")
+    assert result == "'\|value"
+
+    result = csv.escape_value("%value")
+    assert result == "'%value"
+
+    result = csv.escape_value("=cmd|' /C calc'!A0")
+    assert result == "'=cmd\|' /C calc'!A0"
+
+
+def test_df_to_escaped_csv():
+    csv_str = "col_a,=col_b\n-10,=cmd|' /C calc'!A0\na,=b"
+    df = pd.read_csv(io.StringIO(csv_str))
+
+    result = csv.df_to_escaped_csv(df, encoding="utf8", index=False)
+    assert result == "col_a,'=col_b\n-10,'=cmd\|' /C calc'!A0\na,'=b\n"