From 7a8164696bb913a75cf79cf6b57c9973530efefa Mon Sep 17 00:00:00 2001
From: rabbitali <wenxin32@foxmail.com>
Date: Sun, 15 Oct 2023 16:37:55 +0800
Subject: [PATCH 1/1] add a way about key authentication for add host api
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 database/zeus.sql            |  2 +-
 zeus/conf/constant.py        |  6 +--
 zeus/database/table.py       |  2 +-
 zeus/function/verify/host.py |  4 +-
 zeus/host_manager/view.py    | 99 ++++++++++++++++++++++++++++--------
 5 files changed, 85 insertions(+), 28 deletions(-)

diff --git a/database/zeus.sql b/database/zeus.sql
index 3dc9f3c..7db734e 100644
--- a/database/zeus.sql
+++ b/database/zeus.sql
@@ -42,7 +42,7 @@ CREATE TABLE IF NOT EXISTS `host`  (
   `os_version` varchar(40) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci NULL DEFAULT NULL,
   `ssh_user` varchar(40) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci NULL DEFAULT NULL,
   `ssh_port` int(11) NULL DEFAULT NULL,
-  `pkey` varchar(2048) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci NULL DEFAULT NULL,
+  `pkey` varchar(4096) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci NULL DEFAULT NULL,
   `status` int(11) NULL DEFAULT NULL,
   `user` varchar(40) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci NULL DEFAULT NULL,
   `host_group_id` int(11) NULL DEFAULT NULL,
diff --git a/zeus/conf/constant.py b/zeus/conf/constant.py
index 3175c65..bf8792a 100644
--- a/zeus/conf/constant.py
+++ b/zeus/conf/constant.py
@@ -90,9 +90,9 @@ CHECK_IDENTIFY_SCENE = "/check/scene/identify"
 CHECK_WORKFLOW_HOST_EXIST = '/check/workflow/host/exist'
 
 # host template file content
-HOST_TEMPLATE_FILE_CONTENT = """host_ip,ssh_port,ssh_user,password,host_name,host_group_name,management
-test_ip_1,22,root,password,test_host,test_host_group,False
-test_ip_2,22,root,password,test_host,test_host_group,False
+HOST_TEMPLATE_FILE_CONTENT = """host_ip,ssh_port,ssh_user,password,ssh_pkey,host_name,host_group_name,management
+test_ip_1,22,root,password,ssh_pkey,test_host,test_host_group,False
+test_ip_2,22,root,password,ssh_pkey,test_host,test_host_group,False
 """
 
 
diff --git a/zeus/database/table.py b/zeus/database/table.py
index 9596492..265eb45 100644
--- a/zeus/database/table.py
+++ b/zeus/database/table.py
@@ -59,7 +59,7 @@ class Host(Base, MyBase):  # pylint: disable=R0903
     os_version = Column(String(40))
     ssh_user = Column(String(40), default="root")
     ssh_port = Column(Integer(), default=22)
-    pkey = Column(String(2048))
+    pkey = Column(String(4096))
     status = Column(Integer(), default=2)
 
     user = Column(String(40), ForeignKey('user.username'))
diff --git a/zeus/function/verify/host.py b/zeus/function/verify/host.py
index b054d62..d09eedd 100644
--- a/zeus/function/verify/host.py
+++ b/zeus/function/verify/host.py
@@ -103,11 +103,12 @@ class AddHostSchema(Schema):
     """
 
     ssh_user = fields.String(required=True, validate=lambda s: len(s) > 0)
-    password = fields.String(required=True, validate=lambda s: len(s) > 0)
+    password = fields.String(required=True, allow_none=True, validate=lambda s: len(s) >= 0)
     host_name = fields.String(
         required=True, validate=[validate.Length(min=1, max=50), ValidateRules.space_character_check]
     )
     host_ip = fields.IP(required=True)
+    ssh_pkey = fields.String(required=True, allow_none=True, validate=lambda s: 4096 >= len(s) >= 0)
     ssh_port = fields.Integer(required=True, validate=lambda s: 65535 >= s > 0)
     host_group_name = fields.String(required=True, validate=lambda s: len(s) > 0)
     management = fields.Boolean(required=True)
@@ -133,3 +134,4 @@ class UpdateHostSchema(Schema):
     host_name = fields.String(required=False, validate=lambda s: len(s) > 0)
     host_group_name = fields.String(required=False, validate=lambda s: len(s) > 0)
     management = fields.Boolean(required=False)
+    ssh_pkey = fields.String(required=False, validate=lambda s: 4096 >= len(s) >= 0)
diff --git a/zeus/host_manager/view.py b/zeus/host_manager/view.py
index 768d2cd..95e1434 100644
--- a/zeus/host_manager/view.py
+++ b/zeus/host_manager/view.py
@@ -16,12 +16,13 @@ Author:
 Description: Restful APIs for host
 """
 import json
-from io import BytesIO
+from io import BytesIO, StringIO
 from typing import Iterable, List, Tuple, Union
 import socket
 
 import gevent
 import paramiko
+from paramiko.ssh_exception import SSHException
 from flask import request, send_file
 from marshmallow import Schema
 from marshmallow.fields import Boolean
@@ -333,7 +334,8 @@ class AddHost(BaseResponse):
                 "host_ip":"127.0.0.1",
                 "ssh_port":"22",
                 "management":false,
-                "username": "admin"
+                "username": "admin",
+                "ssh_pkey": "RSA key"
             }
 
         Returns:
@@ -363,6 +365,7 @@ class AddHost(BaseResponse):
                 "ssh_port": host_info.get("ssh_port"),
                 "user": host_info.get("username"),
                 "management": host_info.get("management"),
+                "pkey": host_info.get("ssh_pkey"),
             }
         )
         if host in hosts:
@@ -384,7 +387,8 @@ class AddHost(BaseResponse):
                 "host_ip":"127.0.0.1",
                 "ssh_port":"22",
                 "management":false,
-                "username": "admin"
+                "username": "admin",
+                "ssh_pkey": "RSA key"
             }
 
         Returns:
@@ -396,15 +400,55 @@ class AddHost(BaseResponse):
         if status != state.SUCCEED:
             return self.response(code=status)
 
-        status, private_key = save_ssh_public_key_to_client(
-            params.get('host_ip'), params.get('ssh_port'), params.get('ssh_user'), params.get('password')
-        )
-        if status == state.SUCCEED:
-            host.pkey = private_key
-            host.status = HostStatus.ONLINE
+        if params.get("ssh_pkey"):
+            status = verify_ssh_login_info(
+                ClientConnectArgs(
+                    params.get("host_ip"), params.get("ssh_port"), params.get("ssh_user"), params.get("ssh_pkey")
+                )
+            )
+            host.status = HostStatus.ONLINE if status == state.SUCCEED else HostStatus.UNESTABLISHED
+        else:
+            status, private_key = save_ssh_public_key_to_client(
+                params.get('host_ip'), params.get('ssh_port'), params.get('ssh_user'), params.get('password')
+            )
+            if status == state.SUCCEED:
+                host.pkey = private_key
+                host.status = HostStatus.ONLINE
         return self.response(code=self.proxy.add_host(host))
 
 
+def verify_ssh_login_info(ssh_login_info: ClientConnectArgs) -> str:
+    """
+    Verify that the ssh login information is correct
+
+    Args:
+        ssh_login_info(ClientConnectArgs): e.g
+            ClientConnectArgs(host_ip='127.0.0.1', ssh_port=22, ssh_user='root', pkey=RSAKey string)
+
+    Returns:
+        status code
+    """
+    try:
+        client = SSH(
+            ip=ssh_login_info.host_ip,
+            username=ssh_login_info.ssh_user,
+            port=ssh_login_info.ssh_port,
+            pkey=paramiko.RSAKey.from_private_key(StringIO(ssh_login_info.pkey)),
+        )
+        client.close()
+    except socket.error as error:
+        LOGGER.error(error)
+        return state.SSH_CONNECTION_ERROR
+    except SSHException as error:
+        LOGGER.error(error)
+        return state.SSH_AUTHENTICATION_ERROR
+    except Exception as error:
+        LOGGER.error(error)
+        return state.SSH_CONNECTION_ERROR
+
+    return state.SUCCEED
+
+
 def save_ssh_public_key_to_client(ip: str, port: int, username: str, password: str) -> tuple:
     """
     generate RSA key pair,save public key to the target host machine
@@ -465,7 +509,7 @@ class GetHostTemplateFile(BaseResponse):
         file = BytesIO()
         file.write(HOST_TEMPLATE_FILE_CONTENT.encode('utf-8'))
         file.seek(0)
-        response = send_file(file,mimetype="application/octet-stream")
+        response = send_file(file, mimetype="application/octet-stream")
         response.headers['Content-Disposition'] = 'attachment; filename=template.csv'
         return response
 
@@ -574,6 +618,7 @@ class AddHostBatch(BaseResponse):
                 continue
 
             password = host_info.pop("password")
+            pkey = host_info.pop("ssh_pkey", None)
             host_info.update(
                 {"host_group_id": group_id_info.get(host_info['host_group_name']), "user": data["username"]}
             )
@@ -585,7 +630,7 @@ class AddHostBatch(BaseResponse):
                 )
                 continue
 
-            valid_host.append((host, password))
+            valid_host.append((host, password, pkey))
         return valid_host
 
     def save_key_to_client(self, host_connect_infos: List[tuple]) -> list:
@@ -598,8 +643,8 @@ class AddHostBatch(BaseResponse):
         Returns:
             host object list
         """
-        # 30 connections are created at a time.
-        tasks = [host_connect_infos[index : index + 30] for index in range(0, len(host_connect_infos), 30)]
+        # 100 connections are created at a time.
+        tasks = [host_connect_infos[index : index + 100] for index in range(0, len(host_connect_infos), 100)]
         result = []
 
         for task in tasks:
@@ -612,18 +657,23 @@ class AddHostBatch(BaseResponse):
         return result
 
     @staticmethod
-    def update_rsa_key_to_host(host: Host, password: str) -> Host:
+    def update_rsa_key_to_host(host: Host, password: str = None, pkey: str = None) -> Host:
         """
         save ssh public key to client and update its private key in host
 
         Args:
             host(Host): host object
             password(str): password for ssh login
+            pkey(str): rsa key for ssh login
 
         Returns:
             host object
         """
-        status, pkey = save_ssh_public_key_to_client(host.host_ip, host.ssh_port, host.ssh_user, password)
+        if pkey:
+            status = verify_ssh_login_info(ClientConnectArgs(host.host_ip, host.ssh_port, host.ssh_user, pkey))
+        else:
+            status, pkey = save_ssh_public_key_to_client(host.host_ip, host.ssh_port, host.ssh_user, password)
+
         if status == state.SUCCEED:
             host.status = HostStatus.ONLINE
             host.pkey = pkey
@@ -654,7 +704,7 @@ class AddHostBatch(BaseResponse):
                 new_host.update(update_info)
                 self.add_result.append(new_host)
         else:
-            for host, _ in hosts:
+            for host, _, _ in hosts:
                 new_host = {
                     "host_ip": host.host_ip,
                     "ssh_port": host.ssh_port,
@@ -789,9 +839,14 @@ class UpdateHost(BaseResponse):
         """
         ssh_user = params.get("ssh_user") or self.host.ssh_user
         ssh_port = params.get("ssh_port") or self.host.ssh_port
-        status, private_key = save_ssh_public_key_to_client(
-            self.host.host_ip, ssh_port, ssh_user, params.pop("password", None)
-        )
+        private_key = params.pop("ssh_pkey", None)
+        if private_key:
+            status = verify_ssh_login_info(ClientConnectArgs(self.host.host_ip, ssh_port, ssh_user, private_key))
+        else:
+            status, private_key = save_ssh_public_key_to_client(
+                self.host.host_ip, ssh_port, ssh_user, params.pop("password", None)
+            )
+
         params.update(
             {
                 "ssh_user": ssh_user,
@@ -876,10 +931,10 @@ class UpdateHost(BaseResponse):
             return self.response(code=state.PARAM_ERROR, message="there is a duplicate host ssh address in database!")
 
         if params.get("ssh_user") or params.get("ssh_port"):
-            if not params.get("password"):
-                return self.response(code=state.PARAM_ERROR, message="please update password")
+            if not params.get("password") or not params.get("ssh_pkey"):
+                return self.response(code=state.PARAM_ERROR, message="please update password or authentication key.")
             self._save_ssh_key(params)
-        elif params.get("password"):
+        elif params.get("password") or params.get("ssh_pkey"):
             self._save_ssh_key(params)
 
         return self.response(callback.update_host_info(params.pop("host_id"), params))
-- 
2.33.0

