add host batch sync

spug_api/apps/host/urls.py

@@ -16,4 +16,5 @@ urlpatterns = [
     path('import/cloud/', cloud_import),
     path('import/region/', get_regions),
     path('parse/', post_parse),
+    path('valid/', batch_valid),
 ]

spug_api/apps/host/utils.py

@@ -1,14 +1,18 @@
 # Copyright: (c) OpenSpug Organization. https://github.com/openspug/spug
 # Copyright: (c) <spug.dev@gmail.com>
 # Released under the AGPL-3.0 License.
+from django_redis import get_redis_connection
 from libs.helper import make_ali_request, make_tencent_request
 from libs.ssh import SSH, AuthenticationException
 from libs.utils import AttrDict, human_datetime
 from apps.host.models import HostExtend
+from apps.setting.utils import AppSetting
 from collections import defaultdict
 from datetime import datetime, timezone
+from concurrent import futures
 import ipaddress
 import json
+import os
 
 
 def check_os_type(os_name):
@@ -176,22 +180,6 @@ def fetch_tencent_instances(ak, ac, region_id, page_number=1):
     return data
 
 
-def sync_host_extend(host, private_key, public_key, password=None):
-    kwargs = host.to_dict(selects=('hostname', 'port', 'username'))
-    ssh = _get_ssh(kwargs, host.pkey, private_key, public_key, password)
-    form = AttrDict(fetch_host_extend(ssh))
-    form.disk = json.dumps(form.disk)
-    form.public_ip_address = json.dumps(form.public_ip_address)
-    form.private_ip_address = json.dumps(form.private_ip_address)
-    form.updated_at = human_datetime()
-    form.os_type = check_os_type(form.os_name)
-    if hasattr(host, 'hostextend'):
-        extend = host.hostextend
-        extend.update_by_dict(form)
-    else:
-        HostExtend.objects.create(host=host, **form)
-
-
 def fetch_host_extend(ssh):
     commands = [
         "lscpu | grep '^CPU(s)' | awk '{print $2}'",
@@ -223,6 +211,40 @@ def fetch_host_extend(ssh):
     return response
 
 
+def batch_sync_host(token, hosts, password, ):
+    private_key, public_key = AppSetting.get_ssh_key()
+    threads, latest_exception, rds = [], None, get_redis_connection()
+    max_workers = min(10, os.cpu_count() * 4)
+    with futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        for host in hosts:
+            t = executor.submit(_sync_host_extend, host, private_key, public_key, password)
+            t.h_id = host.id
+            threads.append(t)
+        for t in futures.as_completed(threads):
+            exception = t.exception()
+            if exception:
+                rds.rpush(token, json.dumps({'key': t.h_id, 'status': 'fail', 'message': f'{exception}'}))
+            else:
+                rds.rpush(token, json.dumps({'key': t.h_id, 'status': 'ok'}))
+        rds.expire(token, 60)
+
+
+def _sync_host_extend(host, private_key, public_key, password=None):
+    kwargs = host.to_dict(selects=('hostname', 'port', 'username'))
+    ssh = _get_ssh(kwargs, host.pkey, private_key, public_key, password)
+    form = AttrDict(fetch_host_extend(ssh))
+    form.disk = json.dumps(form.disk)
+    form.public_ip_address = json.dumps(form.public_ip_address)
+    form.private_ip_address = json.dumps(form.private_ip_address)
+    form.updated_at = human_datetime()
+    form.os_type = check_os_type(form.os_name)
+    if hasattr(host, 'hostextend'):
+        extend = host.hostextend
+        extend.update_by_dict(form)
+    else:
+        HostExtend.objects.create(host=host, **form)
+
+
 def _get_ssh(kwargs, pkey=None, private_key=None, public_key=None, password=None):
     try:
         if pkey:
@@ -233,6 +255,7 @@ def _get_ssh(kwargs, pkey=None, private_key=None, public_key=None, password=None
             ssh = SSH(password=str(password), **kwargs)
             ssh.add_public_key(public_key)
             return _get_ssh(kwargs, private_key)
-    except AuthenticationException:
+    except AuthenticationException as e:
         if password:
             return _get_ssh(kwargs, None, public_key, public_key, password)
+        raise e

spug_api/apps/host/views.py

@@ -8,12 +8,15 @@ from libs import json_response, JsonParser, Argument, AttrDict
 from apps.setting.utils import AppSetting
 from apps.account.utils import get_host_perms
 from apps.host.models import Host, Group
+from apps.host.utils import batch_sync_host
 from apps.app.models import Deploy
 from apps.schedule.models import Task
 from apps.monitor.models import Detection
 from libs.ssh import SSH, AuthenticationException
 from paramiko.ssh_exception import BadAuthenticationType
 from openpyxl import load_workbook
+from threading import Thread
+import uuid
 
 
 class HostView(View):
@@ -149,3 +152,19 @@ def post_parse(request):
         return json_response(data.decode())
     else:
         return HttpResponseBadRequest()
+
+
+def batch_valid(request):
+    form, error = JsonParser(
+        Argument('password', required=False),
+        Argument('range', filter=lambda x: x in ('1', '2'), help='参数错误')
+    ).parse(request.body)
+    if error is None:
+        if form.range == '1':  # all hosts
+            hosts = Host.objects.all()
+        else:
+            hosts = Host.objects.filter(is_verified=False).all()
+        token = uuid.uuid4().hex
+        Thread(target=batch_sync_host, args=(token, hosts, form.password)).start()
+        return json_response({'token': token, 'hosts': {x.id: {'name': x.name} for x in hosts}})
+    return json_response(error=error)

spug_api/consumer/consumers.py

@@ -46,6 +46,8 @@ class ComConsumer(WebsocketConsumer):
             self.key = f'{settings.BUILD_KEY}:{token}'
         elif module == 'request':
             self.key = f'{settings.REQUEST_KEY}:{token}'
+        elif module == 'host':
+            self.key = token
         else:
             raise TypeError(f'unknown module for {module}')
         self.rds = get_redis_connection()

spug_web/src/pages/host/BatchSync.js

@@ -0,0 +1,96 @@
+/**
+ * Copyright (c) OpenSpug Organization. https://github.com/openspug/spug
+ * Copyright (c) <spug.dev@gmail.com>
+ * Released under the AGPL-3.0 License.
+ */
+import React, { useState, useEffect } from 'react';
+import { observer } from 'mobx-react';
+import { Modal, Form, Input, Button, Radio } from 'antd';
+import { LoadingOutlined } from '@ant-design/icons';
+import { http, X_TOKEN } from 'libs';
+import store from './store';
+
+export default observer(function () {
+  const [loading, setLoading] = useState(false);
+  const [password, setPassword] = useState();
+  const [range, setRange] = useState('2');
+  const [hosts, setHosts] = useState({});
+  const [token, setToken] = useState();
+
+  useEffect(() => {
+    if (token) {
+      let index = 0;
+      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+      const socket = new WebSocket(`${protocol}//${window.location.host}/api/ws/host/${token}/?x-token=${X_TOKEN}`);
+      socket.onopen = () => socket.send(String(index));
+      socket.onmessage = e => {
+        if (e.data === 'pong') {
+          socket.send(String(index))
+        } else {
+          index += 1;
+          const {key, status, message} = JSON.parse(e.data);
+          hosts[key]['status'] = status;
+          hosts[key]['message'] = message;
+          setHosts({...hosts})
+        }
+      }
+      return () => socket && socket.close()
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [token])
+
+  function handleSubmit() {
+    setLoading(true);
+    http.post('/api/host/valid/', {password, range})
+      .then(res => {
+        setHosts(res.hosts);
+        setToken(res.token);
+      })
+      .finally(() => setLoading(false))
+  }
+
+  function handleClose() {
+    store.showSync();
+    store.fetchRecords()
+  }
+
+  const unVerifiedLength = store.records.filter(x => !x.is_verified).length;
+  return (
+    <Modal
+      visible
+      maskClosable={false}
+      title="批量验证（同步）"
+      okText="导入"
+      onCancel={handleClose}
+      footer={null}>
+      <Form hidden={token} labelCol={{span: 6}} wrapperCol={{span: 14}}>
+        <Form.Item name="password" label="默认密码" tooltip="会被用于未验证主机的验证。">
+          <Input.Password value={password} onChange={e => setPassword(e.target.value)}/>
+        </Form.Item>
+        <Form.Item label="选择主机" tooltip="要批量验证/同步哪些主机，全部主机或仅未验证主机。" extra="将会覆盖已有的扩展信息（CPU、内存、磁盘等）。">
+          <Radio.Group
+            value={range}
+            onChange={e => setRange(e.target.value)}
+            options={[
+              {label: `全部（${store.records.length}）`, value: '1'},
+              {label: `未验证（${unVerifiedLength}）`, value: '2'}
+            ]}
+            optionType="button"/>
+        </Form.Item>
+        <Form.Item wrapperCol={{span: 14, offset: 6}}>
+          <Button loading={loading} type="primary" onClick={handleSubmit}>提交验证</Button>
+        </Form.Item>
+      </Form>
+
+      <Form hidden={!token} labelCol={{span: 8}} wrapperCol={{span: 14}}>
+        {Object.entries(hosts).map(([key, item]) => (
+          <Form.Item key={key} label={item.name} help={item.message}>
+            {item.status === 'ok' && <span style={{color: "#52c41a"}}>成功</span>}
+            {item.status === 'fail' && <span style={{color: "red"}}>失败</span>}
+            {item.status === undefined && <LoadingOutlined style={{fontSize: 20}}/>}
+          </Form.Item>
+        ))}
+      </Form>
+    </Modal>
+  );
+})

spug_web/src/pages/host/Table.js

@@ -6,7 +6,7 @@
 import React from 'react';
 import { observer } from 'mobx-react';
 import { Table, Modal, Dropdown, Button, Menu, Avatar, Tooltip, Space, Tag, Radio, message } from 'antd';
-import { PlusOutlined, DownOutlined } from '@ant-design/icons';
+import { PlusOutlined, DownOutlined, SyncOutlined } from '@ant-design/icons';
 import { Action, TableCard, AuthButton, AuthFragment } from 'components';
 import { http, hasPermission } from 'libs';
 import store from './store';
@@ -58,6 +58,11 @@ function ComTable() {
           type="primary"
           icon={<PlusOutlined/>}
           onClick={() => store.showForm()}>新建</AuthButton>,
+        <AuthButton
+          auth="host.host.add"
+          type="primary"
+          icon={<SyncOutlined/>}
+          onClick={() => store.showSync()}>批量验证</AuthButton>,
         <AuthFragment auth="host.host.import">
           <Dropdown overlay={(
             <Menu onClick={handleImport}>

spug_web/src/pages/host/index.js

@@ -13,6 +13,7 @@ import ComTable from './Table';
 import ComForm from './Form';
 import ComImport from './Import';
 import CloudImport from './CloudImport';
+import BatchSync from './BatchSync';
 import Detail from './Detail';
 import Selector from './Selector';
 import store from './store';
@@ -46,6 +47,7 @@ export default observer(function () {
       {store.formVisible && <ComForm/>}
       {store.importVisible && <ComImport/>}
       {store.cloudImport && <CloudImport/>}
+      {store.syncVisible && <BatchSync/>}
       {store.selectorVisible &&
       <Selector oneGroup={!store.addByCopy} onCancel={() => store.selectorVisible = false} onOk={store.updateGroup}/>}
     </AuthDiv>

spug_web/src/pages/host/store.js

@@ -21,6 +21,7 @@ class Store {
   @observable isFetching = false;
   @observable formVisible = false;
   @observable importVisible = false;
+  @observable syncVisible = false;
   @observable cloudImport = null;
   @observable detailVisible = false;
   @observable selectorVisible = false;
@@ -94,6 +95,10 @@ class Store {
     this.record = info
   }
 
+  showSync = () => {
+    this.syncVisible = !this.syncVisible
+  }
+
   showDetail = (info) => {
     this.record = info;
     this.detailVisible = true;