SNOW-1055591: Sanitize terminal output to avoid ASCII escapes (#1238)

RELEASE-NOTES.md

@@ -27,6 +27,7 @@
 * Passing a directory to `snow app deploy` will now deploy any contained file or subfolder specified in the application's artifact rules
 * Fixes markup escaping errors in `snow sql` that may occur when users use unintentionally markup-like escape tags.
 * Fixed case where `snow app teardown` could leave behind orphan applications if they were not created by the Snowflake CLI
+* Improve terminal output sanitization to avoid ASCII escape codes.
 
 # v2.5.0
 ## Backward incompatibility

src/snowflake/cli/api/commands/snow_typer.py

@@ -27,12 +27,14 @@
 from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS
 from snowflake.cli.api.exceptions import CommandReturnTypeError
 from snowflake.cli.api.output.types import CommandResult
+from snowflake.cli.api.sanitizers import sanitize_for_terminal
 
 log = logging.getLogger(__name__)
 
 
 class SnowTyper(typer.Typer):
     def __init__(self, /, **kwargs):
+        self._sanitize_kwargs(kwargs)
         super().__init__(
             **kwargs,
             context_settings=DEFAULT_CONTEXT_SETTINGS,
@@ -41,6 +43,21 @@ def __init__(self, /, **kwargs):
             add_completion=True,
         )
 
+    @staticmethod
+    def _sanitize_kwargs(kwargs: Dict):
+        # Sanitize all string options that are visible in terminal output
+        known_keywords = [
+            "help",
+            "short_help",
+            "options_metavar",
+            "rich_help_panel",
+            "epilog",
+        ]
+        for kw in known_keywords:
+            if kw in kwargs:
+                kwargs[kw] = sanitize_for_terminal(kwargs[kw])
+        return kwargs
+
     @wraps(typer.Typer.command)
     def command(
         self,
@@ -55,11 +72,16 @@ def command(
         logic before and after execution as well as process the result and act on possible
         errors.
         """
+        name = sanitize_for_terminal(name)
+        self._sanitize_kwargs(kwargs)
         if is_enabled is not None and not is_enabled():
             return lambda func: func
 
         def custom_command(command_callable):
             """Custom command wrapper similar to Typer.command."""
+            # Sanitize doc string which is used to create help in terminal
+            command_callable.__doc__ = sanitize_for_terminal(command_callable.__doc__)
+
             if requires_connection:
                 command_callable = global_options_with_connection(command_callable)
             elif requires_global_options:

src/snowflake/cli/api/sanitizers.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed 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.
+
+from __future__ import annotations
+
+import re
+
+# 7-bit C1 ANSI sequences
+_ANSI_ESCAPE = re.compile(
+    r"""
+    \x1B  # ESC
+    (?:   # 7-bit C1 Fe (except CSI)
+        [@-Z\\-_]
+    |     # or [ for CSI, followed by a control sequence
+        \[
+        [0-?]*  # Parameter bytes
+        [ -/]*  # Intermediate bytes
+        [@-~]   # Final byte
+    )
+""",
+    re.VERBOSE,
+)
+
+
+def sanitize_for_terminal(text: str) -> str | None:
+    """
+    Escape ASCII escape codes in string. This should be always used
+    when printing output to terminal.
+    """
+    if text is None:
+        return None
+    return _ANSI_ESCAPE.sub("", text)

src/snowflake/cli/app/printing.py

@@ -36,6 +36,7 @@
     ObjectResult,
     QueryResult,
 )
+from snowflake.cli.api.sanitizers import sanitize_for_terminal
 
 NO_ITEMS_FOUND: str = "No data"
 
@@ -50,6 +51,8 @@ class CustomJSONEncoder(JSONEncoder):
     """Custom JSON encoder handling serialization of non-standard types"""
 
     def default(self, o):
+        if isinstance(o, str):
+            return sanitize_for_terminal(o)
         if isinstance(o, (ObjectResult, MessageResult)):
             return o.result
         if isinstance(o, (CollectionResult, MultipleResults)):
@@ -134,7 +137,7 @@ def print_unstructured(obj: CommandResult | None):
     elif not obj.result:
         rich_print("No data")
     elif isinstance(obj, MessageResult):
-        rich_print(obj.message)
+        rich_print(sanitize_for_terminal(obj.message))
     else:
         if isinstance(obj, ObjectResult):
             _print_single_table(obj)
@@ -149,7 +152,9 @@ def _print_single_table(obj):
     table.add_column("key", overflow="fold")
     table.add_column("value", overflow="fold")
     for key, value in obj.result.items():
-        table.add_row(str(key), str(value))
+        table.add_row(
+            sanitize_for_terminal(str(key)), sanitize_for_terminal(str(value))
+        )
     rich_print(table)
 
 

tests/__snapshots__/test_help_messages.ambr

@@ -7105,6 +7105,79 @@
   ╰──────────────────────────────────────────────────────────────────────────────╯
 
 
+  '''
+# ---
+# name: test_help_messages[streamlit.'\x07']
+  '''
+
+   Usage: default streamlit '' [OPTIONS]                                          
+
+   Deploys a Streamlit app defined in the project definition file                 
+   (snowflake.yml). By default, the command uploads environment.yml and any other 
+   pages or folders, if present. If you don’t specify a stage name, the           
+   `streamlit` stage is used. If the specified stage does not exist, the command  
+   creates it.                                                                    
+
+  ╭─ Options ────────────────────────────────────────────────────────────────────╮
+  │ --replace                Replace the Streamlit app if it already exists.     │
+  │ --open                   Whether to open the Streamlit app in a browser.     │
+  │ --project  -p      TEXT  Path where the Streamlit app project resides.       │
+  │                          Defaults to current working directory.              │
+  │ --help     -h            Show this message and exit.                         │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+  ╭─ Connection configuration ───────────────────────────────────────────────────╮
+  │ --connection,--environment  -c      TEXT  Name of the connection, as defined │
+  │                                           in your `config.toml`. Default:    │
+  │                                           `default`.                         │
+  │ --account,--accountname             TEXT  Name assigned to your Snowflake    │
+  │                                           account. Overrides the value       │
+  │                                           specified for the connection.      │
+  │ --user,--username                   TEXT  Username to connect to Snowflake.  │
+  │                                           Overrides the value specified for  │
+  │                                           the connection.                    │
+  │ --password                          TEXT  Snowflake password. Overrides the  │
+  │                                           value specified for the            │
+  │                                           connection.                        │
+  │ --authenticator                     TEXT  Snowflake authenticator. Overrides │
+  │                                           the value specified for the        │
+  │                                           connection.                        │
+  │ --private-key-path                  TEXT  Snowflake private key path.        │
+  │                                           Overrides the value specified for  │
+  │                                           the connection.                    │
+  │ --database,--dbname                 TEXT  Database to use. Overrides the     │
+  │                                           value specified for the            │
+  │                                           connection.                        │
+  │ --schema,--schemaname               TEXT  Database schema to use. Overrides  │
+  │                                           the value specified for the        │
+  │                                           connection.                        │
+  │ --role,--rolename                   TEXT  Role to use. Overrides the value   │
+  │                                           specified for the connection.      │
+  │ --warehouse                         TEXT  Warehouse to use. Overrides the    │
+  │                                           value specified for the            │
+  │                                           connection.                        │
+  │ --temporary-connection      -x            Uses connection defined with       │
+  │                                           command line parameters, instead   │
+  │                                           of one defined in config           │
+  │ --mfa-passcode                      TEXT  Token to use for multi-factor      │
+  │                                           authentication (MFA)               │
+  │ --enable-diag                             Run python connector diagnostic    │
+  │                                           test                               │
+  │ --diag-log-path                     TEXT  Diagnostic report path             │
+  │ --diag-allowlist-path               TEXT  Diagnostic report path to optional │
+  │                                           allowlist                          │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+  ╭─ Global configuration ───────────────────────────────────────────────────────╮
+  │ --format           [TABLE|JSON]  Specifies the output format.                │
+  │                                  [default: TABLE]                            │
+  │ --verbose  -v                    Displays log entries for log levels `info`  │
+  │                                  and higher.                                 │
+  │ --debug                          Displays log entries for log levels `debug` │
+  │                                  and higher; debug logs contains additional  │
+  │                                  information.                                │
+  │ --silent                         Turns off intermediate output to console.   │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+
+
   '''
 # ---
 # name: test_help_messages[streamlit.deploy]

tests/api/__snapshots__/test_sanitizers.ambr

@@ -0,0 +1,24 @@
+# serializer version: 1
+# name: test_snow_typer_help_sanitization
+  '''
+
+   Usage: root [OPTIONS] COMMAND [ARGS]...                                        
+
+  ╭─ Options ────────────────────────────────────────────────────────────────────╮
+  │ --install-completion            Install completion for the current shell.    │
+  │ --show-completion               Show completion for the current shell, to    │
+  │                                 copy it or customize the installation.       │
+  │ --help                -h        Show this message and exit.                  │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+  ╭─ Commands ───────────────────────────────────────────────────────────────────╮
+  │ ''      doc 4                                                                │
+  │ '🤯?'   doc 3                                                                │
+  │ func1   doc 1 ''                                                             │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+  ╭─ '' ─────────────────────────────────────────────────────────────────────────╮
+  │ func2   ''                                                                   │
+  ╰──────────────────────────────────────────────────────────────────────────────╯
+
+
+  '''
+# ---

tests/api/test_sanitizers.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed 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 pytest
+from snowflake.cli.api.commands.snow_typer import SnowTyper
+from snowflake.cli.api.sanitizers import sanitize_for_terminal
+from typer.testing import CliRunner
+
+
+@pytest.mark.parametrize(
+    "text, expected",
+    [
+        ("'\033[0i\007'", "'\x07'"),
+        ("'\033[5i\007'", "'\x07'"),
+        ("'🤯\033[1000;b\077'", "'🤯?'"),
+        (
+            "'\033[H\007''\033]1337;ClearScrollback\077''\033[2J\007''\033[1;31m'Error, enable MaliciousPlugin in your config'\033[#F\007'",
+            "'\x07''1337;ClearScrollback?''\x07'''Error, enable MaliciousPlugin in your config'\x07'",
+        ),
+    ],
+)
+def test_sanitize_for_terminal(text, expected):
+    result = sanitize_for_terminal(text)
+    assert result == expected
+
+
+def test_snow_typer_help_sanitization(snapshot):
+    app = SnowTyper()
+
+    escape_text = "'\033[0i\007'"
+
+    @app.command(epilog=escape_text, options_metavar=escape_text)
+    def func1():
+        """doc 1 '\033[0i\007'"""
+        return 42
+
+    @app.command(help=escape_text, rich_help_panel=escape_text, short_help=escape_text)
+    def func2():
+        return escape_text
+
+    # Escape as arg not kwarg
+    @app.command("'🤯\033[1000;b\077'")
+    def func3():
+        """doc 3"""
+        return 42
+
+    @app.command(name="'\033[5i\007'")
+    def func4():
+        """doc 4"""
+        return 42
+
+    runner = CliRunner()
+    result = runner.invoke(app, ["--help"])
+    assert result.output == snapshot
+
+    result = runner.invoke(app, ["func2"])
+    assert result.output == ""

tests/notebook/test_notebook_commands.py

@@ -53,6 +53,7 @@ def test_open(mock_launch, mock_url, runner):
 def test_create(mock_create, runner):
     notebook_name = "my_notebook"
     notebook_file = "@stage/notebook.ipynb"
+    mock_create.return_value = "created"
 
     result = runner.invoke(
         ("notebook", "create", notebook_name, "--notebook-file", notebook_file)