From ee8e8cb1bbc3bb1ba27ee6f0e8acfc663cf10c12 Mon Sep 17 00:00:00 2001
From: rearcher <123781007@qq.com>
Date: Tue, 12 Dec 2023 09:47:12 +0800
Subject: [PATCH] Add interface for detecting host status
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 zeus/conf/constant.py        |  2 +
 zeus/database/proxy/host.py  | 73 +++++++++++++++++++++++++++++++++++-
 zeus/function/verify/host.py |  8 ++++
 zeus/host_manager/view.py    | 63 +++++++++++++++++++++++++++++++
 zeus/url.py                  |  2 +
 5 files changed, 147 insertions(+), 1 deletion(-)

diff --git a/zeus/conf/constant.py b/zeus/conf/constant.py
index 9305604..994dd90 100644
--- a/zeus/conf/constant.py
+++ b/zeus/conf/constant.py
@@ -42,6 +42,7 @@ ADD_HOST_BATCH = "/manage/host/add/batch"
 GET_HOST_TEMPLATE_FILE = "/manage/host/file/template"
 DELETE_HOST = "/manage/host/delete"
 QUERY_HOST = "/manage/host/get"
+GET_HOST_STATUS = "/manage/host/status/get"
 GET_HOST_COUNT = "/manage/host/count"
 AUTH_REDIRECT_URL = "/manage/account/authredirecturl"
 BIND_AUTH_ACCOUNT = "/manage/account/bindaccount"
@@ -116,3 +117,4 @@ class HostStatus:
     ONLINE = 0
     OFFLINE = 1
     UNESTABLISHED = 2
+    SCANNING = 3
diff --git a/zeus/database/proxy/host.py b/zeus/database/proxy/host.py
index 1656c56..477c482 100644
--- a/zeus/database/proxy/host.py
+++ b/zeus/database/proxy/host.py
@@ -268,7 +268,6 @@ class HostProxy(MysqlProxy):
                 "host_group_name": host.host_group_name,
                 "host_ip": host.host_ip,
                 "management": host.management,
-                "status": host.status,
                 "scene": host.scene,
                 "os_version": host.os_version,
                 "ssh_port": host.ssh_port,
@@ -340,6 +339,52 @@ class HostProxy(MysqlProxy):
             LOGGER.error("query host %s basic info fail", host_list)
             return DATABASE_QUERY_ERROR, result
 
+    def get_host_ssh_info(self, data):
+        """
+        Get host ssh info according to host id from table
+
+        Args:
+            data(dict): parameter, e.g.
+                {
+                    "username": "admin"
+                    "host_list": ["id1", "id2"]
+                }
+
+        Returns:
+            int: status code
+            dict: query result
+        """
+        username = data.get('username')
+        host_list = data.get('host_list')
+        result = []
+        query_fields = [
+            Host.host_id,
+            Host.host_ip,
+            Host.ssh_port,
+            Host.pkey,
+            Host.ssh_user,
+        ]
+        filters = {Host.user == username}
+        if host_list:
+            filters.add(Host.host_id.in_(host_list))
+        try:
+            hosts = self.session.query(*query_fields).filter(*filters).all()
+            for host in hosts:
+                host_info = {
+                    "host_id": host.host_id,
+                    "host_ip": host.host_ip,
+                    "ssh_port": host.ssh_port,
+                    "pkey": host.pkey,
+                    "ssh_user": host.ssh_user,
+                }
+                result.append(host_info)
+            LOGGER.debug("query host %s ssh info succeed", host_list)
+            return SUCCEED, result
+        except sqlalchemy.exc.SQLAlchemyError as error:
+            LOGGER.error(error)
+            LOGGER.error("query host %s ssh info fail", host_list)
+            return DATABASE_QUERY_ERROR, result
+
     def get_total_host_info_by_user(self, data):
         """
         Get host basic info according to user from table
@@ -775,3 +820,29 @@ class HostProxy(MysqlProxy):
             LOGGER.error(error)
             self.session.rollback()
             return DATABASE_UPDATE_ERROR
+
+
+    def update_host_status(self, host_info: list) -> str:
+        """
+        update host status to host table
+
+        Args:
+            host_info(list): e.g
+                {
+                    "host_id": host_id,
+                    "status": status
+                }
+
+        Returns:
+            str: SUCCEED or DATABASE_UPDATE_ERROR
+        """
+        try:
+            for host in host_info:
+                self.session.query(Host).filter(Host.host_id == host.get('host_id')).update(
+                    {"status": host.get('status')})
+            self.session.commit()
+            return SUCCEED
+        except sqlalchemy.exc.SQLAlchemyError as error:
+            LOGGER.error(error)
+            self.session.rollback()
+            return DATABASE_UPDATE_ERROR
diff --git a/zeus/function/verify/host.py b/zeus/function/verify/host.py
index d09eedd..f746968 100644
--- a/zeus/function/verify/host.py
+++ b/zeus/function/verify/host.py
@@ -60,6 +60,14 @@ class GetHostSchema(Schema):
     per_page = fields.Integer(required=False, validate=lambda s: 50 > s > 0)
 
 
+class GetHostStatusSchema(Schema):
+    """
+    validators for parameter of /manage/host/getstatus
+    """
+
+    host_list = fields.List(fields.Integer(), required=True)
+
+
 class AddHostGroupSchema(Schema):
     """
     validators for parameter of /manage/host/add_host_group
diff --git a/zeus/host_manager/view.py b/zeus/host_manager/view.py
index 10418d1..7ad133d 100644
--- a/zeus/host_manager/view.py
+++ b/zeus/host_manager/view.py
@@ -46,6 +46,7 @@ from zeus.function.verify.host import (
     GetHostGroupSchema,
     GetHostInfoSchema,
     GetHostSchema,
+    GetHostStatusSchema,
     UpdateHostSchema,
 )
 from zeus.host_manager.ssh import SSH, execute_command_and_parse_its_result, generate_key
@@ -118,6 +119,68 @@ class GetHostCount(BaseResponse):
         return self.response(code=status_code, data=result)
 
 
+class GetHostStatus(BaseResponse):
+    """
+    Interface for get host status.
+    Restful API: POST
+    """
+
+    @BaseResponse.handle(schema=GetHostStatusSchema, proxy=HostProxy)
+    def post(self, callback: HostProxy, **params):
+        """
+        get host status
+
+        Args:
+            host_list (list): host id list
+            username: "admin"
+
+        Returns:
+            list: response body
+        """
+        status_code, host_infos = callback.get_host_ssh_info(params)
+
+        multi_thread_handler = MultiThreadHandler(lambda p: self.get_host_status(p), host_infos, None)
+        multi_thread_handler.create_thread()
+        result_list = multi_thread_handler.get_result()
+
+        callback.update_host_status(result_list)
+
+        return self.response(code=status_code, data=result_list)
+
+    @staticmethod
+    def get_host_status(host: dict) -> dict:
+        """
+        Get host status
+
+        Args:
+            host (dict): e.g
+            {
+                "host_id":"host id",
+                "ssh_user":"root",
+                "pkey":"pkey",
+                "host_ip":"host_ip",
+                "ssh_port":"port"
+            }
+
+        Returns:
+        """
+        status = verify_ssh_login_info(
+            ClientConnectArgs(
+                host.get("host_ip"), host.get("ssh_port"), host.get("ssh_user"), host.get("pkey")
+            )
+        )
+        if status == state.SUCCEED:
+            if status != HostStatus.SCANNING:
+                host['status'] = HostStatus.ONLINE
+        elif status == state.SSH_AUTHENTICATION_ERROR:
+            host['status'] = HostStatus.UNESTABLISHED
+        else:
+            host['status'] = HostStatus.OFFLINE
+
+        result = {"host_id": host.get("host_id"), "status": host.get("status")}
+        return result
+
+
 class AddHostGroup(BaseResponse):
     """
     Interface for add host group.
diff --git a/zeus/url.py b/zeus/url.py
index eb8a189..ad8cec9 100644
--- a/zeus/url.py
+++ b/zeus/url.py
@@ -52,6 +52,7 @@ from zeus.conf.constant import (
     USER_LOGIN,
     SYNC_CONFIG,
     OBJECT_FILE_CONFIG,
+    GET_HOST_STATUS,
 )
 from zeus.config_manager import view as config_view
 from zeus.host_manager import view as host_view
@@ -77,6 +78,7 @@ SPECIFIC_URLS = {
         (host_view.DeleteHost, DELETE_HOST),
         (host_view.UpdateHost, UPDATE_HOST),
         (host_view.GetHost, QUERY_HOST),
+        (host_view.GetHostStatus, GET_HOST_STATUS),
         (host_view.GetHostInfo, QUERY_HOST_DETAIL),
         (host_view.GetHostCount, GET_HOST_COUNT),
         (host_view.GetHostTemplateFile, GET_HOST_TEMPLATE_FILE),
-- 
2.33.0

