1
0
Fork 0

Merge pull request #1347 from dbcli/j-bennet/1338-keyring-error

Fix exception when getting password from keyring
This commit is contained in:
Amjith Ramanujam 2022-06-14 11:28:16 -07:00 committed by GitHub
commit 0f21f86838
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 110 additions and 52 deletions

View File

@ -1,6 +1,11 @@
Upcoming:
=========
Bug fixes:
----------
* Fix exception when retrieving password from keyring ([issue 1338](https://github.com/dbcli/pgcli/issues/1338)).
Internal:
---------

58
pgcli/auth.py Normal file
View File

@ -0,0 +1,58 @@
import click
from textwrap import dedent
keyring = None # keyring will be loaded later
keyring_error_message = dedent(
"""\
{}
{}
To remove this message do one of the following:
- prepare keyring as described at: https://keyring.readthedocs.io/en/stable/
- uninstall keyring: pip uninstall keyring
- disable keyring in our configuration: add keyring = False to [main]"""
)
def keyring_initialize(keyring_enabled, *, logger):
"""Initialize keyring only if explicitly enabled"""
global keyring
if keyring_enabled:
# Try best to load keyring (issue #1041).
import importlib
try:
keyring = importlib.import_module("keyring")
except Exception as e: # ImportError for Python 2, ModuleNotFoundError for Python 3
logger.warning("import keyring failed: %r.", e)
def keyring_get_password(key):
"""Attempt to get password from keyring"""
# Find password from store
passwd = ""
try:
passwd = keyring.get_password("pgcli", key) or ""
except Exception as e:
click.secho(
keyring_error_message.format(
"Load your password from keyring returned:", str(e)
),
err=True,
fg="red",
)
return passwd
def keyring_set_password(key, passwd):
try:
keyring.set_password("pgcli", key, passwd)
except Exception as e:
click.secho(
keyring_error_message.format("Set password in keyring returned:", str(e)),
err=True,
fg="red",
)

View File

@ -18,8 +18,6 @@ import platform
from time import time, sleep
from typing import Optional
keyring = None # keyring will be loaded later
from cli_helpers.tabular_output import TabularOutputFormatter
from cli_helpers.tabular_output.preprocessors import align_decimals, format_numbers
from cli_helpers.utils import strip_ansi
@ -49,6 +47,7 @@ from pygments.lexers.sql import PostgresLexer
from pgspecial.main import PGSpecial, NO_QUERY, PAGER_OFF, PAGER_LONG_OUTPUT
import pgspecial as special
from . import auth
from .pgcompleter import PGCompleter
from .pgtoolbar import create_toolbar_tokens_func
from .pgstyle import style_factory, style_factory_output
@ -80,8 +79,6 @@ from psycopg.conninfo import make_conninfo, conninfo_to_dict
from collections import namedtuple
from textwrap import dedent
try:
import sshtunnel
@ -242,7 +239,7 @@ class PGCli:
self.on_error = c["main"]["on_error"].upper()
self.decimal_format = c["data_formats"]["decimal"]
self.float_format = c["data_formats"]["float"]
self.initialize_keyring()
auth.keyring_initialize(c["main"].as_bool("keyring"), logger=self.logger)
self.show_bottom_toolbar = c["main"].as_bool("show_bottom_toolbar")
self.pgspecial.pset_pager(
@ -497,19 +494,6 @@ class PGCli:
pgspecial_logger.addHandler(handler)
pgspecial_logger.setLevel(log_level)
def initialize_keyring(self):
global keyring
keyring_enabled = self.config["main"].as_bool("keyring")
if keyring_enabled:
# Try best to load keyring (issue #1041).
import importlib
try:
keyring = importlib.import_module("keyring")
except Exception as e: # ImportError for Python 2, ModuleNotFoundError for Python 3
self.logger.warning("import keyring failed: %r.", e)
def connect_dsn(self, dsn, **kwargs):
self.connect(dsn=dsn, **kwargs)
@ -552,18 +536,6 @@ class PGCli:
if not self.force_passwd_prompt and not passwd:
passwd = os.environ.get("PGPASSWORD", "")
# Find password from store
key = f"{user}@{host}"
keyring_error_message = dedent(
"""\
{}
{}
To remove this message do one of the following:
- prepare keyring as described at: https://keyring.readthedocs.io/en/stable/
- uninstall keyring: pip uninstall keyring
- disable keyring in our configuration: add keyring = False to [main]"""
)
# Prompt for a password immediately if requested via the -W flag. This
# avoids wasting time trying to connect to the database and catching a
# no-password exception.
@ -574,18 +546,10 @@ class PGCli:
"Password for %s" % user, hide_input=True, show_default=False, type=str
)
if not passwd and keyring:
key = f"{user}@{host}"
try:
passwd = keyring.get_password("pgcli", key) or ""
except (RuntimeError, keyring.errors.InitError) as e:
click.secho(
keyring_error_message.format(
"Load your password from keyring returned:", str(e)
),
err=True,
fg="red",
)
if not passwd and auth.keyring:
passwd = auth.keyring_get_password(key)
def should_ask_for_password(exc):
# Prompt for a password after 1st attempt to connect
@ -669,17 +633,8 @@ class PGCli:
)
else:
raise e
if passwd and keyring:
try:
keyring.set_password("pgcli", key, passwd)
except (RuntimeError, keyring.errors.KeyringError) as e:
click.secho(
keyring_error_message.format(
"Set password in keyring returned:", str(e)
),
err=True,
fg="red",
)
if passwd and auth.keyring:
auth.keyring_set_password(key, passwd)
except Exception as e: # Connecting to a database could fail.
self.logger.debug("Database connection failed: %r.", e)

40
tests/test_auth.py Normal file
View File

@ -0,0 +1,40 @@
import pytest
from unittest import mock
from pgcli import auth
@pytest.mark.parametrize("enabled,call_count", [(True, 1), (False, 0)])
def test_keyring_initialize(enabled, call_count):
logger = mock.MagicMock()
with mock.patch("importlib.import_module", return_value=True) as import_method:
auth.keyring_initialize(enabled, logger=logger)
assert import_method.call_count == call_count
def test_keyring_get_password_ok():
with mock.patch("pgcli.auth.keyring", return_value=mock.MagicMock()):
with mock.patch("pgcli.auth.keyring.get_password", return_value="abc123"):
assert auth.keyring_get_password("test") == "abc123"
def test_keyring_get_password_exception():
with mock.patch("pgcli.auth.keyring", return_value=mock.MagicMock()):
with mock.patch(
"pgcli.auth.keyring.get_password", side_effect=Exception("Boom!")
):
assert auth.keyring_get_password("test") == ""
def test_keyring_set_password_ok():
with mock.patch("pgcli.auth.keyring", return_value=mock.MagicMock()):
with mock.patch("pgcli.auth.keyring.set_password"):
auth.keyring_set_password("test", "abc123")
def test_keyring_set_password_exception():
with mock.patch("pgcli.auth.keyring", return_value=mock.MagicMock()):
with mock.patch(
"pgcli.auth.keyring.set_password", side_effect=Exception("Boom!")
):
auth.keyring_set_password("test", "abc123")