From cb89ec46f000c6cc21be9d5d1102d1ae594369f9 Mon Sep 17 00:00:00 2001
From: vapao <live1989@foxmail.com>
Date: Sat, 20 Jun 2020 19:26:02 +0800
Subject: [PATCH] =?UTF-8?q?U=20=E4=BC=98=E5=8C=96=E4=B8=BB=E6=9C=BA?=
 =?UTF-8?q?=E9=AA=8C=E8=AF=81=E4=BD=BF=E6=9B=B4=E5=87=86=E7=A1=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 spug_api/apps/host/views.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/spug_api/apps/host/views.py b/spug_api/apps/host/views.py
index 9d2c192..ad286c8 100644
--- a/spug_api/apps/host/views.py
+++ b/spug_api/apps/host/views.py
@@ -11,6 +11,7 @@ from apps.schedule.models import Task
 from apps.monitor.models import Detection
 from apps.account.models import Role
 from libs.ssh import SSH, AuthenticationException
+from paramiko.ssh_exception import BadAuthenticationType
 from libs import human_datetime, AttrDict
 from openpyxl import load_workbook
 import socket
@@ -118,7 +119,7 @@ def post_import(request):
             summary['skip'].append(i)
             continue
         try:
-            if valid_ssh(data.hostname, data.port, data.username, data.pop('password') or password) is False:
+            if valid_ssh(data.hostname, data.port, data.username, data.pop('password') or password, False) is False:
                 summary['fail'].append(i)
                 continue
         except AuthenticationException:
@@ -137,7 +138,7 @@ def post_import(request):
     return json_response(summary)
 
 
-def valid_ssh(hostname, port, username, password):
+def valid_ssh(hostname, port, username, password, with_expect=True):
     try:
         private_key = AppSetting.get('private_key')
         public_key = AppSetting.get('public_key')
@@ -145,18 +146,20 @@ def valid_ssh(hostname, port, username, password):
         private_key, public_key = SSH.generate_key()
         AppSetting.set('private_key', private_key, 'ssh private key')
         AppSetting.set('public_key', public_key, 'ssh public key')
+    cli = SSH(hostname, port, username, private_key)
     if password:
-        cli = SSH(hostname, port, username, password=str(password))
-        code, out = cli.exec_command('mkdir -p -m 700 ~/.ssh && \
+        _cli = SSH(hostname, port, username, password=str(password))
+        code, out = _cli.exec_command('mkdir -p -m 700 ~/.ssh && \
                 echo %r >> ~/.ssh/authorized_keys && \
                 chmod 600 ~/.ssh/authorized_keys' % public_key)
         if code != 0:
             raise Exception(f'add public key error: {out!r}')
-    else:
-        cli = SSH(hostname, port, username, private_key)
-
     try:
         cli.ping()
+    except BadAuthenticationType:
+        if with_expect:
+            raise TypeError('该主机不支持密钥认证，请参考官方文档，错误代码：E01')
+        return False
     except AuthenticationException:
         return False
     return True