Merge pull request #622 from bensonhome/main

🎨更新start命令参数和quickscan拉工具

Author: Lingghh

client/node/common/cmdarg.py

@@ -75,6 +75,8 @@ def parse_args():
         start_parser = subparsers.add_parser('start', help="启动节点")
         start_parser.add_argument("-t", "--token", dest='token', type=str, help="个人token,在代码分析网站获取", required=True)
         start_parser.add_argument("--org-sid", dest="org_sid", help="团队编号,在代码分析网站获取。指定注册为团队的节点，不指定则为公共节点")
+        start_parser.add_argument("--tag", dest="tag", help="机器标签")
+        start_parser.add_argument("--create-from", dest="create_from", type=str, help="客户端节点启动渠道")
 
         # quickscan命令
         quickscan_parser = subparsers.add_parser('quickscan', help="执行快速分析")

client/node/quicktask/toolloader.py

@@ -53,8 +53,20 @@ def load_tools(args):
         else:
             tool_tasks = QuickScan.get_scan_tasks(languages, labels, {})
         tool_names = [task['task_name'] for task in tool_tasks]
+
+        custom_tools = []
+        for tool_name in tool_names:
+            try:
+                __import__("tool." + tool_name)
+            except ModuleNotFoundError:
+                # 记录自定义工具列表
+                custom_tools.append(tool_name)
+            except:
+                LogPrinter.exception("encounter error.")
+                pass
+
         LogPrinter.info("Initing other tools ...")
-        ToolLoader(tool_names=tool_names, task_list=tool_tasks, include_common=False).git_load_tools(print_enable=False)
+        ToolLoader(tool_names=tool_names, task_list=tool_tasks, custom_tools=custom_tools, include_common=False).git_load_tools(print_enable=False)
 
 
 if __name__ == '__main__':

client/node/servertask/looprunner.py

@@ -40,7 +40,9 @@ def __init__(self, args):
         TaskRunner.__init__(self)
 
         self._token = args.token
+        self._tag = args.tag
         self._org_sid = args.org_sid
+        self._create_from = args.create_from
         self._server_url = LocalConfig.get_server_url()
         # 打印连接的sever地址
         LogPrinter.info("using server: %s" % self._server_url)
@@ -140,7 +142,7 @@ def _send_result(self, task):
     def run(self):
         """looprunner主函数"""
         # 向server注册节点
-        NodeMgr().register_node(self._server, self._org_sid)
+        NodeMgr().register_node(self._server, self._tag, self._org_sid, self._create_from)
 
         # 启动心跳上报线程
         HeartBeat(self._server).start()

client/node/servertask/nodemgr.py

@@ -14,6 +14,7 @@
 import psutil
 import uuid
 import sys
+import socket
 
 from node import app
 from platform import platform
@@ -24,23 +25,34 @@
 class NodeMgr(object):
     """本地节点管理"""
 
-    def register_node(self, server, org_sid=None):
+    def get_docker_uuid(self, create_from, tag):
+        """如果从docker创建，获取到docker的主机名，与标签一起拼接成为节点唯一标识NODE_UUID"""
+        if create_from and "docker" == create_from:
+            host_name = socket.gethostname()
+            return f"{tag}-{host_name}"
+        return None
+
+    def register_node(self, server, tag=None, org_sid=None, create_from=None):
         '''用本地node_uuid向server注册，获取server给的node_id。
         如果node_id和本地存储node_id不一致，则抛出异常。
-
-        :param server: server的node rpc接口
         '''
+        if not tag:
+            tag = app.settings.OS_TAG_MAP[sys.platform]
         node_uuid = app.persist_data.get('NODE_UUID')
         if not node_uuid:
-            node_uuid = uuid.uuid1().hex
+            node_uuid = self.get_docker_uuid(create_from, tag)
+            if not node_uuid:
+                node_uuid = uuid.uuid1().hex
             app.persist_data['NODE_UUID'] = node_uuid
-        tag = app.settings.OS_TAG_MAP[sys.platform]
+
         data = {
             "uuid": node_uuid,
             "tag": tag,
             "os_info": app.settings.PLATFORMS[sys.platform],
             "org_sid": org_sid  # 为空时，表示为公共节点，不为空时，表示指定团队的节点
         }
+        if create_from:
+            data["create_from"] = create_from
         node_id = server.register(data)
         LogPrinter.info('node(%s) registered in server node id:%s', node_uuid, node_id)
         app.persist_data['NODE_ID'] = node_id