improve host sync

spug_api/apps/host/extend.py

@@ -4,8 +4,7 @@
 from django.views.generic import View
 from libs import json_response, JsonParser, Argument, human_datetime
 from apps.host.models import Host, HostExtend
-from apps.host.utils import check_os_type
-import ipaddress
+from apps.host.utils import check_os_type, fetch_host_extend
 import json
 
 
@@ -20,33 +19,8 @@ class ExtendView(View):
                 return json_response(error='未找到指定主机')
             if not host.is_verified:
                 return json_response(error='该主机还未验证')
-            cli = host.get_ssh()
-            commands = [
-                "lscpu | grep '^CPU(s)' | awk '{print $2}'",
-                "free -m | awk 'NR==2{print $2}'",
-                "hostname -I",
-                "cat /etc/os-release | grep PRETTY_NAME | awk -F \\\" '{print $2}'",
-                "fdisk -l | grep '^Disk /' | awk '{print $5}'"
-            ]
-            code, out = cli.exec_command(';'.join(commands))
-            if code != 0:
-                return json_response(error=f'Exception: {out}')
-            response = {'disk': [], 'public_ip_address': [], 'private_ip_address': []}
-            for index, line in enumerate(out.strip().split('\n')):
-                if index == 0:
-                    response['cpu'] = int(line)
-                elif index == 1:
-                    response['memory'] = round(int(line) / 1000, 1)
-                elif index == 2:
-                    for ip in line.split():
-                        if ipaddress.ip_address(ip).is_global:
-                            response['public_ip_address'].append(ip)
-                        else:
-                            response['private_ip_address'].append(ip)
-                elif index == 3:
-                    response['os_name'] = line
-                else:
-                    response['disk'].append(round(int(line) / 1024 / 1024 / 1024, 0))
+            ssh = host.get_ssh()
+            response = fetch_host_extend(ssh)
             return json_response(response)
         return json_response(error=error)
 

spug_api/apps/host/utils.py

@@ -2,8 +2,13 @@
 # Copyright: (c) <spug.dev@gmail.com>
 # Released under the AGPL-3.0 License.
 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 collections import defaultdict
 from datetime import datetime, timezone
+import ipaddress
+import json
 
 
 def check_os_type(os_name):
@@ -169,3 +174,65 @@ def fetch_tencent_instances(ak, ac, region_id, page_number=1):
         new_data = fetch_tencent_instances(ak, ac, region_id, page_number)
         data.extend(new_data)
     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}'",
+        "free -m | awk 'NR==2{print $2}'",
+        "hostname -I",
+        "cat /etc/os-release | grep PRETTY_NAME | awk -F \\\" '{print $2}'",
+        "fdisk -l | grep '^Disk /' | awk '{print $5}'",
+        "fdisk -l | grep '^磁盘 /' | awk '{print $5}'"
+    ]
+    code, out = ssh.exec_command(';'.join(commands))
+    if code != 0:
+        raise Exception(out)
+    response = {'disk': [], 'public_ip_address': [], 'private_ip_address': []}
+    for index, line in enumerate(out.strip().split('\n')):
+        if index == 0:
+            response['cpu'] = int(line)
+        elif index == 1:
+            response['memory'] = round(int(line) / 1000, 1)
+        elif index == 2:
+            for ip in line.split():
+                if ipaddress.ip_address(ip).is_global:
+                    response['public_ip_address'].append(ip)
+                else:
+                    response['private_ip_address'].append(ip)
+        elif index == 3:
+            response['os_name'] = line
+        else:
+            response['disk'].append(round(int(line) / 1024 / 1024 / 1024, 0))
+    return response
+
+
+def _get_ssh(kwargs, pkey=None, private_key=None, public_key=None, password=None):
+    try:
+        if pkey:
+            ssh = SSH(pkey=pkey, **kwargs)
+            ssh.ping()
+            return ssh
+        else:
+            ssh = SSH(password=str(password), **kwargs)
+            ssh.add_public_key(public_key)
+            return _get_ssh(kwargs, private_key)
+    except AuthenticationException:
+        if password:
+            return _get_ssh(kwargs, None, public_key, public_key, password)