import traceback
from contextlib import contextmanager
from enum import IntEnum
from typing import List
import requests
from schrodinger import get_maestro
from schrodinger.models import parameters
from schrodinger.models.mappers import TargetSpec
from schrodinger.Qt import QtWidgets
from schrodinger.Qt.QtCore import pyqtSlot
from schrodinger.ui.qt.basewidgets import Panel
from schrodinger.utils import fileutils
from schrodinger.utils import mmutil
from schrodinger.utils import preferences
from schrodinger.utils.env import prepend_sys_path
from . import login
from . import login2_ui
from . import login_gui
maestro = get_maestro()
ldclient = None
TOKEN_PATH = '/livedesign/api/tokens/generate_refresh'
PREFKEY_MODE = 'mode'
PREFKEY_HOST = 'host'
PREFKEY_USERNAME = 'username'
[docs]class AuthenticationMode(IntEnum):
CREDENTIALS = 0
SINGLE = 1
[docs] @classmethod
def getItemFromInt(cls, value: int):
"""
:return: the enum value corresponding to the supplied integer, if any
"""
for item in cls:
if int(item) == value:
return item
[docs]class LDLoginPanelModel(parameters.CompoundParam):
ld_client: object = None
import_paths: List
authentication_mode: AuthenticationMode = AuthenticationMode.CREDENTIALS
credentials_host: str
single_sign_on_host: str
username: str
password: str
token_link: str
token: str
error_msg: str
[docs] def getActiveHost(self) -> str:
"""
:return: the host value associated with the selected authentication mode
"""
if self.authentication_mode == AuthenticationMode.CREDENTIALS:
return self.credentials_host
else:
return self.single_sign_on_host
[docs] def setActiveHost(self, host: str):
"""
:param host: assign this value to the host parameter corresponding to
the selected authentication mode
"""
if self.authentication_mode == AuthenticationMode.CREDENTIALS:
self.credentials_host = host
else:
self.single_sign_on_host = host
[docs]class LDLoginPanel(Panel):
"""
Panel for logging into LiveDesign using standard credentials or SAML token.
"""
ui_module = login2_ui
model_class = LDLoginPanelModel
APPLY_LEGACY_STYLESHEET = False
[docs] def initSetOptions(self):
super().initSetOptions()
self._pref_handler = preferences.Preferences(preferences.SCRIPTS)
self._pref_handler.beginGroup(type(self).__name__)
[docs] def initSetUp(self):
super().initSetUp()
self.setWindowTitle('Connect to LiveDesign')
ui = self.ui
self.spacer = QtWidgets.QSpacerItem(0, 0,
QtWidgets.QSizePolicy.Expanding,
QtWidgets.QSizePolicy.Fixed)
rb_type_tuples = [
(ui.ld_cred_rb, AuthenticationMode.CREDENTIALS),
(ui.single_rb, AuthenticationMode.SINGLE),
]
for rb, authentication_mode in rb_type_tuples:
ui.sign_in_group.setId(rb, authentication_mode)
ui.log_in_btn.clicked.connect(self._onLogInClicked)
ui.get_token_link_btn.clicked.connect(self._onGetTokenLinkClicked)
ui.back_btn.clicked.connect(self._onBackClicked)
ui.confirm_btn.clicked.connect(self._authenticate)
for le in [ui.ld_cred_host_le, ui.username_le, ui.password_le]:
le.returnPressed.connect(self._onCredentialsSubmitted)
ui.single_sign_on_host_le.returnPressed.connect(
ui.get_token_link_btn.click)
ui.token_le.returnPressed.connect(ui.confirm_btn.click)
if maestro:
# Add a direct connection as we need the global vars to be reset as
# soon as the project is closed, not when the login panel is
# reopened as it would be if the callback method decorator is used.
maestro.project_close_callback_add(self._onProjectClosed)
[docs] def initSetDefaults(self):
super().initSetDefaults()
model = self.model
self._setGetTokenModeActive(False)
mode_int = self._pref_handler.get(PREFKEY_MODE, default=None)
if mode_int is None:
return
model.authentication_mode = AuthenticationMode.getItemFromInt(mode_int)
host = self._pref_handler.get(PREFKEY_HOST, default=None)
if host is not None:
model.setActiveHost(host)
if model.authentication_mode == AuthenticationMode.CREDENTIALS:
model.username = self._pref_handler.get(PREFKEY_USERNAME,
default='')
[docs] def defineMappings(self):
ui = self.ui
M = self.model_class
def set_stack_widget(authentication_mode):
type_wdg_map = {
AuthenticationMode.CREDENTIALS: ui.credentials_wdg,
AuthenticationMode.SINGLE: ui.single_sign_on_wdg
}
wdg = type_wdg_map[authentication_mode]
ui.sign_in_stack.setCurrentWidget(wdg)
stack_target = TargetSpec(setter=set_stack_widget)
return super().defineMappings() + [
(ui.sign_in_group, M.authentication_mode),
(stack_target, M.authentication_mode),
(ui.ld_cred_host_le, M.credentials_host),
(ui.single_sign_on_host_le, M.single_sign_on_host),
(ui.single_sign_on_host_lbl, M.single_sign_on_host),
(ui.username_le, M.username),
(ui.password_le, M.password),
(ui.link_lbl, M.token_link),
(ui.token_le, M.token),
(ui.error_lbl, M.error_msg),
(self._updateLogInButtonState, M.credentials_host),
(self._updateLogInButtonState, M.username),
(self._updateLogInButtonState, M.password),
(self._onSingleSignInHostChanged, M.single_sign_on_host),
(self._updateConfirmButtonState, M.token),
]
def _updateHost(self):
"""
Update the host string to use the proper formatting.
"""
host = self.model.getActiveHost()
host = login.format_host(host)
self.model.setActiveHost(host)
def _onSingleSignInHostChanged(self):
url = self.model.getActiveHost() + TOKEN_PATH
self.model.token_link = f'<a href="{url}">{url}</a>'
self._updateConnectButtonState()
def _updateLogInButtonState(self):
model = self.model
enable = bool(model.credentials_host and model.username and
model.password)
self.ui.log_in_btn.setEnabled(enable)
def _updateConnectButtonState(self):
enable = bool(self.model.single_sign_on_host)
self.ui.get_token_link_btn.setEnabled(enable)
def _updateConfirmButtonState(self):
enable = bool(self.model.token)
self.ui.confirm_btn.setEnabled(enable)
@pyqtSlot()
def _onCredentialsSubmitted(self):
if self.ui.log_in_btn.isEnabled():
self.ui.log_in_btn.click()
@contextmanager
def _contactingServer(self):
"""
Context manager for communicating with LiveDesign.
Update the error message label with useful information in the event of
a failure. Clear the error message label otherwise.
"""
model = self.model
msg = ''
try:
yield
except requests.HTTPError as error:
str_error = str(error)
is_unauthorized = 'Unauthorized' in str_error
if (model.authentication_mode == AuthenticationMode.CREDENTIALS and
is_unauthorized):
msg = login.INVALID_CREDENTIALS_MSG
elif is_unauthorized:
msg = login.INVALID_TOKEN_MSG
else:
msg = str_error
except RuntimeError as error:
str_error = str(error)
if 'not supported' in str_error:
msg = login.VERSION_MISMATCH_MSG
else:
msg = str_error
msg += login.CONTACT_SUPPORT_MSG
self._disconnectFromHost()
except requests.exceptions.ConnectionError as error:
str_error = str(error)
if 'Operation timed out' in str_error:
# Show special message when the request timed out - which
# is the most common exception raised.
msg = login.TIMEOUT_MSG
elif 'Failed to establish a new connection' in str_error:
msg = login.NO_CONNECTION_MSG
else:
# Something else. Message will be shown on the 2nd line in the
# status label.
msg = f'{login.NO_LDCLIENT_MSG}:\n{error}'
self._disconnectFromHost()
except Exception as exc:
# Print the traceback to the terminal, so that we know what type
# of exception to catch. Still show something to the user as
# well (without crashing the GUI):
msg = str(exc)
traceback.print_exc()
self._disconnectFromHost()
finally:
model.error_msg = msg
@pyqtSlot()
def _onLogInClicked(self):
"""
Contact the LD server, download the LD client module, and authenticate.
"""
if not self._connectToHost():
self._clearGlobalVariables()
return
self._authenticate()
@pyqtSlot()
def _onGetTokenLinkClicked(self):
"""
Attempt to access the single sign-on token from the specified server.
If successful, advance the GUI state.
"""
if self._connectToHost():
self._setGetTokenModeActive(True)
else:
self._disconnectFromHost()
@pyqtSlot()
def _onBackClicked(self):
"""
Deactivate "get token mode".
"""
self._setGetTokenModeActive(False)
@pyqtSlot()
def _authenticate(self):
"""
Attempt to authenticate with the LD client.
If successful, close this panel. Otherwise, "disconnect" from the
server.
"""
model = self.model
self._setUpLDClient()
if not model.error_msg:
host = model.getActiveHost()
self._pref_handler.set(PREFKEY_MODE, int(model.authentication_mode))
self._pref_handler.set(PREFKEY_HOST, host)
self._pref_handler.set(PREFKEY_USERNAME, model.username)
self._updateGlobalVariables()
self.close()
def _connectToHost(self) -> bool:
"""
Attempt to access the `LDClient` module from the specified server.
:return: whether the connection was successful
"""
model = self.model
model.error_msg = ''
self._updateHost()
with self._contactingServer():
self._downloadLDClient()
return not model.error_msg
def _disconnectFromHost(self):
"""
Update the panel and global variables when there is a problem with a
connection to the host.
"""
self._setGetTokenModeActive(False)
self.model.ld_client = None
self.model.import_paths = []
self._clearGlobalVariables()
def _downloadLDClient(self):
"""
Attempt to download the LiveDesign `client` module.
"""
model = self.model
model.import_paths = []
client_path = login.download_ld_client(
url=model.getActiveHost() + login.LDCLIENT_PATH,
tmp_dir=fileutils.get_directory_path(fileutils.TEMP),
tar_filename='ldclient',
glob_path='ldclient-*')
model.import_paths = [client_path]
def _setUpLDClient(self):
"""
Import the LD client module and authenticate with the server.
"""
model = self.model
module_path = model.import_paths[0] if model.import_paths else None
try:
_import_client(module_path)
except ImportError:
model.error_msg = login.IMPORT_ERROR_MSG + login.CONTACT_SUPPORT_MSG
self._disconnectFromHost()
return
# Instantiate the `LDClient` instance using the specified authentication
# procedure
kwargs = {
'host': model.getActiveHost() + login.API_PATH,
login.COMPATIBILITY_MODE: login.LD_VERSION_COMPATIBILITY_MODE,
}
if model.authentication_mode == AuthenticationMode.CREDENTIALS:
kwargs['username'] = model.username
kwargs['password'] = model.password
else:
kwargs['refresh_token'] = model.token
with self._contactingServer():
model.ld_client = ldclient.client.LDClient(**kwargs)
def _setGetTokenModeActive(self, active: bool):
"""
Progress or revert the single sign on widget.
If "connected", lock in the host by hiding the line edit and exposing a
label. Also, expose the token-related widgets. Otherwise, hide token-
related widgets and once again expose the ability to edit the host.
:param active: whether the GUI should express that this panel is
"connected" to a LiveDesign server
"""
ui = self.ui
for wdg in [ui.ld_cred_rb, ui.single_rb, ui.single_sign_on_host_le]:
wdg.setEnabled(not active)
ui.single_sign_on_host_le.setVisible(not active)
ui.get_token_link_wdg.setVisible(not active)
ui.single_sign_on_host_lbl.setVisible(active)
ui.token_groupbox.setVisible(active)
ui.token_groupbox.setEnabled(active)
ui.confirm_wdg.setVisible(active)
layout = ui.single_sign_on_host_layout
if active:
ui.token_le.setFocus()
layout.addSpacerItem(self.spacer)
else:
ui.single_sign_on_host_le.setFocus()
layout.removeItem(self.spacer)
self.model.token = ''
def _onProjectClosed(self):
"""
If current project is closed, then all saved credentials will be reset.
"""
self.initSetDefaults()
self._clearGlobalVariables()
def _clearGlobalVariables(self):
"""
Clear global LiveDesign log in variables.
"""
login._SESSION_HOST = None
login._SESSION_USERNAME = None
login._SESSION_PWD = None
login._TOKEN = None
login._SESSION_IMPORT_PATHS = []
def _updateGlobalVariables(self):
"""
Update global LiveDesign log in variables using data from the active
authentication mode.
"""
model = self.model
login._SESSION_HOST = model.getActiveHost()
if model.authentication_mode == AuthenticationMode.CREDENTIALS:
login._SESSION_USERNAME = model.username
login._SESSION_PWD = model.password
else:
login._SESSION_TOKEN = model.token
login._SESSION_IMPORT_PATHS = model.import_paths
def _import_client(module_path: str):
"""
Import the `ldclient.client` module from the supplied path.
:param module_path: the import path
"""
global ldclient
with prepend_sys_path(module_path):
import ldclient.client
[docs]def log_into_livedesign() -> bool:
"""
Attempt to access cached LiveDesign login credentials. If unavailable,
launch the login panel and store valid responses.
:return: whether LiveDesign login credentials have been cached
"""
if mmutil.feature_flag_is_enabled(mmutil.SAML_LOGIN_LIVEDESIGN):
panel_class = LDLoginPanel
else:
panel_class = login_gui.LiveDesignLoginDialog
if not login.required_login_info_set():
panel_instance = panel_class()
panel_instance.run(blocking=True, modal=True)
return login.required_login_info_set()
[docs]def panel():
return LDLoginPanel.panel()
if __name__ == "__main__":
panel()