From 4e32546faf227a6497ce8b282fff7450cae6f665 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Mon, 12 Oct 2020 19:27:01 +0100
Subject: [PATCH] Mask Password in Log table when using the CLI (#11468)

---
 airflow/utils/cli.py         | 14 ++++++++++-
 tests/utils/test_cli_util.py | 48 +++++++++++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py
index dc0dbeb8ec403..6e0ea25826cd6 100644
--- a/airflow/utils/cli.py
+++ b/airflow/utils/cli.py
@@ -105,8 +105,20 @@ def _build_metrics(func_name, namespace):
     :param namespace: Namespace instance from argparse
     :return: dict with metrics
     """
+    sensitive_fields = {'-p', '--password', '--conn-password'}
+    full_command = list(sys.argv)
+    for idx, command in enumerate(full_command):  # pylint: disable=too-many-nested-blocks
+        if command in sensitive_fields:
+            # For cases when password is passed as "--password xyz" (with space between key and value)
+            full_command[idx + 1] = "*" * 8
+        else:
+            # For cases when password is passed as "--password=xyz" (with '=' between key and value)
+            for sensitive_field in sensitive_fields:
+                if command.startswith(f'{sensitive_field}='):
+                    full_command[idx] = f'{sensitive_field}={"*" * 8}'
+
     metrics = {'sub_command': func_name, 'start_datetime': datetime.utcnow(),
-               'full_command': '{}'.format(list(sys.argv)), 'user': getpass.getuser()}
+               'full_command': '{}'.format(full_command), 'user': getpass.getuser()}
 
     if not isinstance(namespace, Namespace):
         raise ValueError("namespace argument should be argparse.Namespace instance,"
diff --git a/tests/utils/test_cli_util.py b/tests/utils/test_cli_util.py
index ed17a4e4d1138..dfa71e10c7413 100644
--- a/tests/utils/test_cli_util.py
+++ b/tests/utils/test_cli_util.py
@@ -16,12 +16,16 @@
 # specific language governing permissions and limitations
 # under the License.
 #
-
+import json
 import os
+import sys
 import unittest
 from argparse import Namespace
 from contextlib import contextmanager
 from datetime import datetime
+from unittest import mock
+
+from parameterized import parameterized
 
 from airflow import settings
 from airflow.exceptions import AirflowException
@@ -85,6 +89,48 @@ def test_get_dags(self):
         with self.assertRaises(AirflowException):
             cli.get_dags(None, "foobar", True)
 
+    @parameterized.expand(
+        [
+            (
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin --password test",
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin --password ********"
+            ),
+            (
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin -p test",
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin -p ********"
+            ),
+            (
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin --password=test",
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin --password=********"
+            ),
+            (
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin -p=test",
+                "airflow users create -u test2 -l doe -f jon -e jdoe@apache.org -r admin -p=********"
+            ),
+            (
+                "airflow connections add dsfs --conn-login asd --conn-password test --conn-type google",
+                "airflow connections add dsfs --conn-login asd --conn-password ******** --conn-type google",
+            )
+        ]
+    )
+    def test_cli_create_user_supplied_password_is_masked(self, given_command, expected_masked_command):
+        args = given_command.split()
+
+        expected_command = expected_masked_command.split()
+
+        exec_date = datetime.utcnow()
+        namespace = Namespace(dag_id='foo', task_id='bar', subcommand='test', execution_date=exec_date)
+        with mock.patch.object(sys, "argv", args):
+            metrics = cli._build_metrics(args[1], namespace)
+
+        self.assertTrue(metrics.get('start_datetime') <= datetime.utcnow())
+
+        log = metrics.get('log')
+        command = json.loads(log.extra).get('full_command')  # type: str
+        # Replace single quotes to double quotes to avoid json decode error
+        command = json.loads(command.replace("'", '"'))
+        self.assertEqual(command, expected_command)
+
 
 @contextmanager
 def fail_action_logger_callback():