diff --git a/connect.py b/connect.py index 273edc052..8a4949662 100755 --- a/connect.py +++ b/connect.py @@ -188,8 +188,12 @@ def get_user_hostgroup_host(username, gid): def verify_connect(username, part_ip): ip_matched = [] - hosts_attr = get_user_host(username) - hosts = hosts_attr.values() + try: + hosts_attr = get_user_host(username) + hosts = hosts_attr.values() + except ServerError, e: + color_print(e, 'red') + return False for ip_info in hosts: for info in ip_info[1:]: if part_ip in info: @@ -217,7 +221,11 @@ def print_prompt(): def print_user_host(username): - hosts_attr = get_user_host(username) + try: + hosts_attr = get_user_host(username) + except ServerError, e: + color_print(e, 'red') + return hosts = hosts_attr.keys() hosts.sort() for ip in hosts: @@ -267,7 +275,6 @@ def connect(username, password, host, port, login_name): global channel win_size = get_win_size() channel = ssh.invoke_shell(height=win_size[0], width=win_size[1]) - #channel.resize_pty(height=win_size[0], width=win_size[1]) try: signal.signal(signal.SIGWINCH, set_win_size) except: diff --git a/jumpserver/api.py b/jumpserver/api.py index f7846f956..7c416925a 100644 --- a/jumpserver/api.py +++ b/jumpserver/api.py @@ -349,14 +349,18 @@ def get_user_host(username): """Get the hosts of under the user control.""" hosts_attr = {} asset_all = user_perm_asset_api(username) - user = User.objects.get(username=username) - for asset in asset_all: - alias = AssetAlias.objects.filter(user=user, host=asset) - if alias and alias[0].alias != '': - hosts_attr[asset.ip] = [asset.id, asset.ip, alias[0].alias] - else: - hosts_attr[asset.ip] = [asset.id, asset.ip, asset.comment] - return hosts_attr + user = User.objects.filter(username=username) + if user: + user = user[0] + for asset in asset_all: + alias = AssetAlias.objects.filter(user=user, host=asset) + if alias and alias[0].alias != '': + hosts_attr[asset.ip] = [asset.id, asset.ip, alias[0].alias] + else: + hosts_attr[asset.ip] = [asset.id, asset.ip, asset.comment] + return hosts_attr + else: + raise ServerError('User %s does not exit!' % username) def get_connect_item(username, ip):