[+] _ssh_connect

[+] Checks after reboot

src/main.py

@@ -14,17 +14,17 @@ from getpass import getpass
 from pathlib import Path
 from typing import Callable
 
-import netmiko
 from cryptography.fernet import Fernet, InvalidToken
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from loguru import logger
-from netmiko import ConnectHandler
+from netmiko import ConnectHandler, NetmikoAuthenticationException
 from pyVim.connect import SmartConnect, Disconnect
 from pyVmomi import vim
 from ruamel.yaml import YAML
 
+timeout = 120
 chars = string.ascii_letters + string.digits
 yaml = YAML()
 workdir = Path.home() / ".vmker"
@@ -64,7 +64,7 @@ class VMKer:
         self.si = None
         self.content = None
         self.config = None
-        self.vm = set()
+        self.vm_list = set()
         self.template = None
 
     @staticmethod
@@ -292,58 +292,78 @@ class VMKer:
 
     def _configure(self, vm):
         _name = vm.name
+        # TODO: Default user
+        # TODO: key file
         user = 'root'
         key_file = 'rsa'
 
-        logger.info("Waiting VirtualMachine..")
+        logger.info(f"[{_name}] Waiting VirtualMachine..")
         t = 0
         while True:
             ipv4, ipv6 = self._get_vm_ip(vm)
             if len(ipv4) > 0:
                 break
-            time.sleep(1)
+            time.sleep(0.99)
             t += 1
-            if t > 120:
+            if t > timeout:
                 return logger.error("No IPv4 address found.")
 
         logger.info(f"[{_name}] IPv4: {', '.join(ipv4)}")
         logger.info(f"[{_name}] IPv6: {', '.join(ipv6)}")
 
+        _ssh_connect = {
+            'device_type': 'linux',
+            'host': ipv4[0],
+            'username': user,
+            'use_keys': True,
+            'key_file': key_file
+        }
+
         try:
-            # TODO: Default user
-            # TODO: key file
-            ssh = ConnectHandler(
-                device_type='linux',
-                host=ipv4[0],
-                username=user,
-                use_keys=True,
-                key_file=key_file
-            )
-        except netmiko.exceptions.NetmikoAuthenticationException as e:
-            return logger.error(f"SSH Error: Authentication. ({e})")
-        logger.success(f"[SSH] [{_name}] Connected to {ipv4[0]} as {user} ({key_file=})")
-        # TODO: length
-        _pwd = ''.join(random.choice(chars) for _ in range(12))
-        ssh.send_command(f"echo '{user}:{_pwd}' | chpasswd")
-        logger.info(f"Password for root: {_pwd}")
-        ssh.send_command(f'echo {_name} > /etc/hostname')
+            ssh = ConnectHandler(**_ssh_connect)
+            logger.success(f"[{_name}] [SSH] Connected to {ipv4[0]} as {user} ({key_file=})")
+            # TODO: length
+            _pwd = ''.join(random.choice(chars) for _ in range(12))
+            ssh.send_command(f"echo '{user}:{_pwd}' | chpasswd")
+            logger.info(f"Password for root: {_pwd}")
+            ssh.send_command(f'echo {_name} > /etc/hostname')
 
-        # TODO: yml
-        logger.info(f"[{vm.name}] Executing %.yml")
+            # TODO: yml
+            logger.info(f"[{_name}] [YML] Executing %.yml")
 
-        ssh.disconnect()
+            logger.info(f"[{_name}] Rebooting")
+            ssh.send_command('reboot')
+            ssh.disconnect()
+
+            end_time = time.time() + timeout
+            while time.time() < end_time:
+                try:
+                    ssh = ConnectHandler(**_ssh_connect)
+                except Exception:
+                    time.sleep(1)
+            if not ssh.is_alive():
+                return logger.error(f"[{_name}] Host startup TimeOut.")
+
+            logger.debug(f'[{_name}] uname {ssh.send_command("uname -a")}')
+            logger.info(f"[{_name}] [YML] Checks %.yml")
+
+        except NetmikoAuthenticationException as e:
+            return logger.error(f"[{_name}] [SSH] Error: Authentication. ({e})")
+        except Exception as e:
+            logger.error(f"[{_name}] [SSH] Unhandled error ")
+            return logger.exception(e)
 
     def _cache_vms(self, template_name, datacenter):
         logger.debug("Checking VMs.")
-        self.vm = set()
+        self.vm_list = set()
         self.template = None
         container_view = self.content.viewManager.CreateContainerView(datacenter.vmFolder, [vim.VirtualMachine], True)
         for _vm in container_view.view:
             if _vm.name == template_name:
                 self.template = _vm
-            self.vm.add(_vm.name)
+            self.vm_list.add(_vm.name)
         container_view.Destroy()
-        logger.debug(f"{self.vm=}")
+        logger.debug(f"{self.vm_list=}")
         logger.debug(f"{self.template=}")
 
     def _create(self):
@@ -375,7 +395,7 @@ class VMKer:
             else:
                 new_names = [new_name]
             for n in new_names:
-                if n in self.vm:
+                if n in self.vm_list:
                     return logger.error(f"Error: {n!r} already exist.")
             if not self.template:
                 return logger.error(f"Template: {template_name} - not found")