Repository: apache/airflow

Issue:
Mask secrets in stdout for `airflow tasks test ...` CLI command (#17476)
The problem with `airflow tasks test ...` is that any `stdout` that comes from a task inside an operator or other callable does not go through a logger.  It runs in the main process.  This results in secrets that are printed to `stdout` to not be masked according to the `airflow.task` logger.

I wanted a way to capture and filter stdout without having to route everything through a secondary logger.  I decided a context manager that filters `stdout` would be the best way to go about this.  The context manager redacts all secrets according to the filters on the `airflow.task` logger.

Closes #17476 

Sensitive variables don't get masked when rendered with airflow tasks test
**Apache Airflow version**: 2.1.2


**Kubernetes version (if you are using kubernetes)** (use `kubectl version`): No

**Environment**:

- **Cloud provider or hardware configuration**: No
- **OS** (e.g. from /etc/os-release): MacOS Big Sur 11.4
- **Kernel** (e.g. `uname -a`):  -
- **Install tools**: -
- **Others**: -

**What happened**:

With the following code:
```
from airflow import DAG
from airflow.models import Variable
from airflow.operators.python import PythonOperator

from datetime import datetime, timedelta

def _extract():
    partner = Variable.get("my_dag_partner_secret")
    print(partner)

with DAG("my_dag", start_date=datetime(2021, 1 , 1), schedule_interval="@daily") as dag:

    extract = PythonOperator(
        task_id="extract",
        python_callable=_extract
    )
```

By executing the command

`
airflow tasks test my_dag extract 2021-01-01
`

The value of the variable my_dag_partner_secret gets rendered in the logs whereas it shouldn't

```
[2021-08-06 19:05:30,088] {taskinstance.py:1303} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=my_dag
AIRFLOW_CTX_TASK_ID=extract
AIRFLOW_CTX_EXECUTION_DATE=2021-01-01T00:00:00+00:00
partner_a
[2021-08-06 19:05:30,091] {python.py:151} INFO - Done. Returned value was: None
[2021-08-06 19:05:30,096] {taskinstance.py:1212} INFO - Marking task as SUCCESS. dag_id=my_dag, task_id=extract, execution_date=20210101T000000, start_date=20210806T131013, end_date=20210806T190530
```
**What you expected to happen**:

The value should be masked like on the UI or in the logs

**How to reproduce it**:

DAG given above


**Anything else we need to know**:

Nop



Discussion:

I can taking a look into this. Will try to reproduce.
I repeated above instructions and was able reproduce. So it looks like the airflow/util/log/secrets_masker.py:should_hide_value_for_key() is where the check for hiding a sensitive variable happens. Looking into the workflow of running DAGs from cli and if it hits this function/class
So from [1]: "The automatic masking is triggered by Connection or Variable access. This means that if you pass a sensitive value via XCom or any other side-channel it will not be masked when printed in the downstream task." and going into task-command::task_test it looks like logs are being propagated so they are no longer masked. Seems to be WAI
[1] https://airflow.apache.org/docs/apache-airflow/stable/security/secrets/index.html#masking-sensitive-data
Hello! 
Any news? 
It is critical for us to display secrets in the CI log (
@artful88533 -> suggestion - maybe you or @ShakaibKhan can take a look and provide fix for that? If you care about it, this is the most "certain" way to make it land in 2.2.3 - way more certain than pinging here - and becoming one of the > 1800 contributors to Airflow is a great way to pay back for the free software you use.

IMHO it is not crirtical (when you have access to the CLI you already can read all the secretes in whatever way you want) but if you need to get it for your CI process, maybe that's a good incentive for you to fix it ? It seems that this is just a question of adding secret masker as filter when the CLI commands are run.
Actually, I think this might be more critical than it looks at first glance @potiuk.  By running the `airflow tasks test ...` command through a BashOperator in a separate DAG, which tests the task from the DAG that Marc linked in the original issue, it's actually possible to display unmasked secrets through the Airflow Web UI.
> Actually, I think this might be more critical than it looks at first glance @potiuk. By running the `airflow tasks test ...` command through a BashOperator in a separate DAG, which tests the task from the DAG that Marc linked in the original issue, it's actually possible to display unmasked secrets through the Airflow Web UI.

Why would you want to run "airflow tasks test" in DAG? Is this a valid casethat is likely? Maybe I am not understnding something, but I am not sure I see the case when it could be used in "production" in a valid scenario? 

Just to give a bit of context  - there are many ways you could print the unmasked values for connections, variables, even secrets. For example you could easily launch a subprocess calling "python -c print(Connection.get('conn_id'))" or just running "airlfow connection list" a as a command to print unmasked paswords.  

The way how  masking is done currently will not prevent this if the Connection is not used in the task before.  So "Masking" is not "total prevention" of showing the secret values, it just prevents from accidental printing of those in the "regular use cases".  There are many ways how DAG writer could print those and bypass secrets masker deliberately. So my question is - how likely and "normal" it is to run 'airlfow test" inside the "execute" method of a task. I think very unlikely.

And BTW in the future when we implement DB-less mode and maybe even (as a follow up) we will further harden Airlfow to not be able to reach out to read Airflow DB at all, this might be more "hardened" but as of currently we have no mechanism to prevent the DAG writers to print any secret they want to the task log. That's simply impossible.  
That's very true.  Thanks for explaining your reasoning about this.

Resolving patch:
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -53,6 +53,7 @@
 )
 from airflow.utils.dates import timezone
 from airflow.utils.log.logging_mixin import StreamLogWriter
+from airflow.utils.log.secrets_masker import RedactedIO
 from airflow.utils.net import get_hostname
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.state import DagRunState
@@ -546,10 +547,11 @@ def task_test(args, dag=None):
     ti, dr_created = _get_ti(task, args.execution_date_or_run_id, args.map_index, create_if_necessary="db")
 
     try:
-        if args.dry_run:
-            ti.dry_run()
-        else:
-            ti.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)
+        with redirect_stdout(RedactedIO()):
+            if args.dry_run:
+                ti.dry_run()
+            else:
+                ti.run(ignore_task_deps=True, ignore_ti_state=True, test_mode=True)
     except Exception:
         if args.post_mortem:
             debugger = _guess_debugger()
diff --git a/airflow/utils/log/secrets_masker.py b/airflow/utils/log/secrets_masker.py
--- a/airflow/utils/log/secrets_masker.py
+++ b/airflow/utils/log/secrets_masker.py
@@ -18,18 +18,17 @@
 import collections
 import logging
 import re
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+import sys
+from typing import Any, Dict, Iterable, List, Optional, Set, TextIO, Tuple, TypeVar, Union
 
 from airflow import settings
 from airflow.compat.functools import cache, cached_property
 
-if TYPE_CHECKING:
-    RedactableItem = Union[str, Dict[Any, Any], Tuple[Any, ...], List[Any]]
-
+Redactable = TypeVar("Redactable", str, Dict[Any, Any], Tuple[Any, ...], List[Any])
+Redacted = Union[Redactable, str]
 
 log = logging.getLogger(__name__)
 
-
 DEFAULT_SENSITIVE_FIELDS = frozenset(
     {
         'access_token',
@@ -91,14 +90,13 @@ def mask_secret(secret: Union[str, dict, Iterable], name: Optional[str] = None)
     _secrets_masker().add_mask(secret, name)
 
 
-def redact(value: "RedactableItem", name: Optional[str] = None) -> "RedactableItem":
+def redact(value: Redactable, name: Optional[str] = None) -> Redacted:
     """Redact any secrets found in ``value``."""
     return _secrets_masker().redact(value, name)
 
 
 @cache
 def _secrets_masker() -> "SecretsMasker":
-
     for flt in logging.getLogger('airflow.task').filters:
         if isinstance(flt, SecretsMasker):
             return flt
@@ -177,7 +175,7 @@ def filter(self, record) -> bool:
 
         return True
 
-    def _redact_all(self, item: "RedactableItem", depth: int) -> "RedactableItem":
+    def _redact_all(self, item: Redactable, depth: int) -> Redacted:
         if depth > self.MAX_RECURSION_DEPTH or isinstance(item, str):
             return '***'
         if isinstance(item, dict):
@@ -190,7 +188,7 @@ def _redact_all(self, item: "RedactableItem", depth: int) -> "RedactableItem":
         else:
             return item
 
-    def _redact(self, item: "RedactableItem", name: Optional[str], depth: int) -> "RedactableItem":
+    def _redact(self, item: Redactable, name: Optional[str], depth: int) -> Redacted:
         # Avoid spending too much effort on redacting on deeply nested
         # structures. This also avoid infinite recursion if a structure has
         # reference to self.
@@ -231,7 +229,7 @@ def _redact(self, item: "RedactableItem", name: Optional[str], depth: int) -> "R
             )
             return item
 
-    def redact(self, item: "RedactableItem", name: Optional[str] = None) -> "RedactableItem":
+    def redact(self, item: Redactable, name: Optional[str] = None) -> Redacted:
         """Redact an any secrets found in ``item``, if it is a string.
 
         If ``name`` is given, and it's a "sensitive" name (see
@@ -258,3 +256,23 @@ def add_mask(self, secret: Union[str, dict, Iterable], name: Optional[str] = Non
         elif isinstance(secret, collections.abc.Iterable):
             for v in secret:
                 self.add_mask(v, name)
+
+
+class RedactedIO(TextIO):
+    """IO class that redacts values going into stdout.
+
+    Expected usage::
+
+        with contextlib.redirect_stdout(RedactedIO()):
+            ...  # Writes to stdout will be redacted.
+    """
+
+    def __init__(self):
+        self.target = sys.stdout
+
+    def write(self, s: str) -> int:
+        s = redact(s)
+        return self.target.write(s)
+
+    def flush(self) -> None:
+        return self.target.flush()


Test patch:
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -58,7 +58,7 @@ def reset(dag_id):
 
 
 # TODO: Check if tests needs side effects - locally there's missing DAG
-class TestCliTasks(unittest.TestCase):
+class TestCliTasks:
     run_id = 'TEST_RUN_ID'
     dag_id = 'example_python_operator'
     parser: ArgumentParser
@@ -67,7 +67,7 @@ class TestCliTasks(unittest.TestCase):
     dag_run: DagRun
 
     @classmethod
-    def setUpClass(cls):
+    def setup_class(cls):
         cls.dagbag = DagBag(include_examples=True)
         cls.parser = cli_parser.get_parser()
         clear_db_runs()
@@ -78,7 +78,7 @@ def setUpClass(cls):
         )
 
     @classmethod
-    def tearDownClass(cls) -> None:
+    def teardown_class(cls) -> None:
         clear_db_runs()
 
     def test_cli_list_tasks(self):
@@ -103,20 +103,34 @@ def test_test(self):
         assert "'example_python_operator__print_the_context__20180101'" in stdout.getvalue()
 
     @pytest.mark.filterwarnings("ignore::airflow.utils.context.AirflowContextDeprecationWarning")
-    def test_test_with_existing_dag_run(self):
+    def test_test_with_existing_dag_run(self, caplog):
         """Test the `airflow test` command"""
         task_id = 'print_the_context'
-
         args = self.parser.parse_args(["tasks", "test", self.dag_id, task_id, DEFAULT_DATE.isoformat()])
+        with caplog.at_level("INFO", logger="airflow.task"):
+            task_command.task_test(args)
+        assert f"Marking task as SUCCESS. dag_id={self.dag_id}, task_id={task_id}" in caplog.text
+
+    @pytest.mark.filterwarnings("ignore::airflow.utils.context.AirflowContextDeprecationWarning")
+    def test_test_filters_secrets(self, capsys):
+        """Test ``airflow test`` does not print secrets to stdout.
+
+        Output should be filtered by SecretsMasker.
+        """
+        password = "somepassword1234!"
+        logging.getLogger("airflow.task").filters[0].add_mask(password)
+        args = self.parser.parse_args(
+            ["tasks", "test", "example_python_operator", "print_the_context", "2018-01-01"],
+        )
 
-        with self.assertLogs('airflow.task', level='INFO') as cm:
+        with mock.patch("airflow.models.TaskInstance.run", new=lambda *_, **__: print(password)):
             task_command.task_test(args)
-            assert any(
-                [
-                    f"Marking task as SUCCESS. dag_id={self.dag_id}, task_id={task_id}" in log
-                    for log in cm.output
-                ]
-            )
+        assert capsys.readouterr().out.endswith("***\n")
+
+        not_password = "!4321drowssapemos"
+        with mock.patch("airflow.models.TaskInstance.run", new=lambda *_, **__: print(not_password)):
+            task_command.task_test(args)
+        assert capsys.readouterr().out.endswith(f"{not_password}\n")
 
     @mock.patch("airflow.cli.commands.task_command.get_dag_by_deserialization")
     @mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
@@ -229,7 +243,7 @@ def test_run_raises_when_theres_no_dagrun(self, mock_local_job):
             task0_id,
             run_id,
         ]
-        with self.assertRaises(DagRunNotFound):
+        with pytest.raises(DagRunNotFound):
             task_command.task_run(self.parser.parse_args(args0), dag=dag)
 
     def test_cli_test_with_params(self):
diff --git a/tests/utils/log/test_secrets_masker.py b/tests/utils/log/test_secrets_masker.py
--- a/tests/utils/log/test_secrets_masker.py
+++ b/tests/utils/log/test_secrets_masker.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import contextlib
 import inspect
 import logging
 import logging.config
@@ -23,7 +24,7 @@
 import pytest
 
 from airflow import settings
-from airflow.utils.log.secrets_masker import SecretsMasker, should_hide_value_for_key
+from airflow.utils.log.secrets_masker import RedactedIO, SecretsMasker, should_hide_value_for_key
 from tests.test_utils.config import conf_vars
 
 settings.MASK_SECRETS_IN_LOGS = True
@@ -340,3 +341,23 @@ def formatException(self, exc_info):
 def lineno():
     """Returns the current line number in our program."""
     return inspect.currentframe().f_back.f_lineno
+
+
+class TestRedactedIO:
+    def test_redacts_from_print(self, capsys):
+        # Without redacting, password is printed.
+        print(p)
+        stdout = capsys.readouterr().out
+        assert stdout == f"{p}\n"
+        assert "***" not in stdout
+
+        # With context manager, password is redacted.
+        with contextlib.redirect_stdout(RedactedIO()):
+            print(p)
+        stdout = capsys.readouterr().out
+        assert stdout == "***\n"
+
+    def test_write(self, capsys):
+        RedactedIO().write(p)
+        stdout = capsys.readouterr().out
+        assert stdout == "***"