From 7a8164696bb913a75cf79cf6b57c9973530efefa Mon Sep 17 00:00:00 2001 From: rabbitali 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