From 6caf246fca72910f26d48628ed13000859555824 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= <[email protected]>
Date: Sat, 29 Aug 2020 04:37:38 +0200
Subject: Add sensitive data masking to Connection.get_uri

---
 airflow/cli/commands/connection_command.py | 31 +++++++---------------
 airflow/models/connection.py               |  9 ++++---
 2 files changed, 14 insertions(+), 26 deletions(-)

diff --git a/airflow/cli/commands/connection_command.py b/airflow/cli/commands/connection_command.py
index 618f0079c..2a76b55e3 100644
--- a/airflow/cli/commands/connection_command.py
+++ b/airflow/cli/commands/connection_command.py
@@ -20,7 +20,6 @@ import json
 import os
 import sys
 from typing import List
-from urllib.parse import urlunparse
 
 import pygments
 import yaml
@@ -39,21 +38,6 @@ from airflow.utils.code_utils import get_terminal_formatter
 from airflow.utils.session import create_session
 
 
-def _prep_msg(msg, conn):
-    """Prepare status messages for connections"""
-
-    msg = msg.format(conn_id=conn.conn_id,
-                     uri=conn.get_uri() or
-                     urlunparse((conn.conn_type,
-                                '{login}:{password}@{host}:{port}'
-                                 .format(login=conn.conn_login or '',
-                                         password='******' if conn.conn_password else '',
-                                         host=conn.conn_host or '',
-                                         port=conn.conn_port or ''),
-                                 conn.conn_schema or '', '', '', '')))
-    return msg
-
-
 def _tabulate_connection(conns: List[Connection], tablefmt: str):
     tabulate_data = [
         {
@@ -224,12 +208,13 @@ def connections_add(args):
         if not (session.query(Connection)
                 .filter(Connection.conn_id == new_conn.conn_id).first()):
             session.add(new_conn)
-            msg = '\n\tSuccessfully added `conn_id`={conn_id} : {uri}\n'
-            msg = _prep_msg(msg, new_conn)
+            msg = '\n\tSuccessfully added `conn_id`={conn_id} : {uri}\n'.format(
+                conn_id=new_conn.conn_id,
+                uri=new_conn.get_uri(display_sensitive=False)
+            )
             print(msg)
         else:
-            msg = '\n\tA connection with `conn_id`={conn_id} already exists\n'
-            msg = msg.format(conn_id=new_conn.conn_id)
+            msg = f'\n\tA connection with `conn_id`={new_conn.conn_id} already exists\n'
             print(msg)
 
 
@@ -278,8 +263,10 @@ def _prep_import_status_msgs(conn_status_map):
 
         msg = msg + status + " : \n\t"
         for conn in conn_list:
-            msg = msg + '\n\t`conn_id`={conn_id} : {uri}\n'
-            msg = _prep_msg(msg, conn)
+            msg = msg + '\n\t`conn_id`={conn_id} : {uri}\n'.format(
+                conn_id=conn.conn_id,
+                uri=conn.get_uri(display_sensitive=False)
+            )
     return msg
 
 
diff --git a/airflow/models/connection.py b/airflow/models/connection.py
index 44e8e576a..efd17a02e 100644
--- a/airflow/models/connection.py
+++ b/airflow/models/connection.py
@@ -219,7 +219,7 @@ class Connection(Base, LoggingMixin):
         if uri_parts.query:
             self.extra = json.dumps(dict(parse_qsl(uri_parts.query, keep_blank_values=True)))
 
-    def get_uri(self) -> str:
+    def get_uri(self, display_sensitive=True) -> str:
         """Return connection in URI format"""
         uri = '{}://'.format(str(self.conn_type).lower().replace('_', '-'))
 
@@ -228,7 +228,7 @@ class Connection(Base, LoggingMixin):
             authority_block += quote(self.login, safe='')
 
         if self.password is not None:
-            authority_block += ':' + quote(self.password, safe='')
+            authority_block += ':' + quote(self.password, safe='') if display_sensitive else ":******"
 
         if authority_block > '':
             authority_block += '@'
@@ -250,8 +250,9 @@ class Connection(Base, LoggingMixin):
 
         uri += host_block
 
-        if self.extra_dejson:
-            uri += '?{}'.format(urlencode(self.extra_dejson))
+        extra_dejson = self.extra_dejson
+        if extra_dejson:
+            uri += '?{}'.format(urlencode(extra_dejson)) if display_sensitive else "?******"
 
         return uri
 
-- 
2.28.0