Issue #355 Use setproctitle to hide command line passwords

3 files changed, 34 insertions, 3 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 1ee56836..ad877b99 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -3,6 +3,7 @@ from __future__ import unicode_literals
 from __future__ import print_function
 
 import os
+import re
 import sys
 import traceback
 import logging
@@ -12,6 +13,7 @@ from time import time
 from codecs import open
 
 import click
+import setproctitle
 import sqlparse
 from prompt_toolkit import CommandLineInterface, Application, AbortAction
 from prompt_toolkit.enums import DEFAULT_BUFFER
@@ -520,8 +522,18 @@ def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
             '\thost: %r'
             '\tport: %r', database, user, host, port)
 
+    obscure_process_password()
     pgcli.run_cli()
 
+def obscure_process_password():
+    process_title = setproctitle.getproctitle()
+    if '://' in process_title:
+        process_title = re.sub(r":(.*):(.*)@", r":\1:xxxx@", process_title)
+    elif "=" in process_title:
+        process_title = re.sub(r"password=\S+", "password=xxxx", process_title)
+
+    setproctitle.setproctitle(process_title)
+
 def format_output(title, cur, headers, status, table_format, expanded=False, max_width=None):
     output = []
     if title:  # Only print the title if it's not None.
diff --git a/setup.py b/setup.py
index e69be836..835c5bea 100644
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,8 @@ setup(
             'prompt_toolkit==0.46',
             'psycopg2 >= 2.5.4',
             'sqlparse == 0.1.16',
-            'configobj >= 5.0.6'
+            'configobj >= 5.0.6',
+            'setproctitle >= 1.1.9'
             ],
         entry_points='''
             [console_scripts]
diff --git a/tests/test_main.py b/tests/test_main.py
index e25f7f08..6e8d14b0 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,5 +1,5 @@
 import pytest
-from pgcli.main import need_completion_refresh
+from pgcli.main import need_completion_refresh, obscure_process_password
 
 
 @pytest.mark.parametrize('sql', [
@@ -7,4 +7,22 @@ from pgcli.main import need_completion_refresh
     'SELECT * FROM foo; DROP TABLE foo',
 ])
 def test_need_completion_refresh(sql):
-    assert need_completion_refresh(sql)
\ No newline at end of file
+    assert need_completion_refresh(sql)
+
+def test_obscure_process_password():
+    import setproctitle
+    original_title = setproctitle.getproctitle()
+
+    setproctitle.setproctitle("pgcli user=root password=secret host=localhost")
+    obscure_process_password()
+    title = setproctitle.getproctitle()
+    expected = "pgcli user=root password=xxxx host=localhost"
+    assert title == expected
+
+    setproctitle.setproctitle("pgcli postgres://root:secret@localhost/db")
+    obscure_process_password()
+    title = setproctitle.getproctitle()
+    expected = "pgcli postgres://root:xxxx@localhost/db"
+    assert title == expected
+
+    setproctitle.setproctitle(original_title)