From 2bd159509f6d74710bf28ff50a08e9f20887c002 Mon Sep 17 00:00:00 2001
From: rabbitali <shusheng.wen@outlook.com>
Date: Tue, 25 Apr 2023 10:29:32 +0800
Subject: [PATCH] fix shell command return error and update register function
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ceres/function/command.py  | 10 +++++++++-
 ceres/function/register.py | 10 ++++------
 ceres/function/schema.py   |  2 +-
 ceres/function/util.py     |  8 ++------
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/ceres/function/command.py b/ceres/function/command.py
index df1c5ae..d84177e 100644
--- a/ceres/function/command.py
+++ b/ceres/function/command.py
@@ -132,6 +132,9 @@ def collect_command_manage(args):
         if not validate_data(data, STRING_ARRAY):
             exit(1)
         print(json.dumps(Collect.collect_file(data)))
+    else:
+        print("Please check the input parameters!")
+        exit(1)
 
 
 def plugin_command_manage(args):
@@ -152,7 +155,9 @@ def plugin_command_manage(args):
         print(json.dumps(change_collect_items(data)))
     elif args.info:
         print(json.dumps(Collect.get_plugin_info()))
-
+    else:
+        print("Please check the input parameters!")
+        exit(1)
 
 def cve_command_manage(args):
     if args.set_repo:
@@ -180,3 +185,6 @@ def cve_command_manage(args):
         status_code, cve_fix_result = VulnerabilityManage().cve_fix(data.get("cves"))
         res = StatusCode.make_response_body((status_code, {"result": cve_fix_result}))
         print(json.dumps(res))
+    else:
+        print("Please check the input parameters!")
+        exit(1)
\ No newline at end of file
diff --git a/ceres/function/register.py b/ceres/function/register.py
index 8176df8..f8ee397 100644
--- a/ceres/function/register.py
+++ b/ceres/function/register.py
@@ -77,12 +77,10 @@ def register(register_info: dict) -> int:
         LOGGER.error(e)
         return HTTP_CONNECT_ERROR
 
-    if ret.status_code != SUCCESS:
+    if ret.status_code != requests.codes["ok"]:
         LOGGER.warning(ret.text)
         return ret.status_code
 
-    ret_data = json.loads(ret.text)
-    if ret_data.get('code') == SUCCESS:
-        return SUCCESS
-    LOGGER.error(ret_data)
-    return int(ret_data.get('code'))
+    if ret.json().get('label') != SUCCESS:
+        LOGGER.error(ret.text)
+    return ret.json().get('label')
diff --git a/ceres/function/schema.py b/ceres/function/schema.py
index e7e4ce7..f8541aa 100644
--- a/ceres/function/schema.py
+++ b/ceres/function/schema.py
@@ -133,6 +133,6 @@ CVE_FIX_SCHEMA = {
 HOST_INFO_SCHEMA = {
     "type": "array",
     "items": {
-        "enum": ["os", "cpu", "memory"]
+        "enum": ["os", "cpu", "memory", "disk"]
     }
 }
diff --git a/ceres/function/util.py b/ceres/function/util.py
index 73a0014..42cebe2 100644
--- a/ceres/function/util.py
+++ b/ceres/function/util.py
@@ -159,12 +159,8 @@ def get_dict_from_file(file_path: str) -> dict:
     try:
         with open(file_path, "r") as f:
             data = json.load(f)
-    except FileNotFoundError:
-        LOGGER.error('file not found')
-        data = {}
-    except json.decoder.JSONDecodeError:
-        LOGGER.error('Json conversion error, the file content'
-                     ' structure is not json format.')
+    except (IOError, ValueError) as error:
+        LOGGER.error(error)
         data = {}
     if not isinstance(data, dict):
         data = {}
-- 
Gitee