From 6756ae3c286b2d9ba1fc97e2423f3fd7bf860ded Mon Sep 17 00:00:00 2001
From: Clay Sweetser <clay.sweetser@gmail.com>
Date: Mon, 27 Feb 2023 18:56:03 -0500
Subject: [PATCH] Simplify startup logic.

---
 netbox_agent/cli.py            | 48 +++++++++++++++++++---------------
 netbox_agent/virtualmachine.py | 23 +++++++---------
 2 files changed, 37 insertions(+), 34 deletions(-)

diff --git a/netbox_agent/cli.py b/netbox_agent/cli.py
index e112469a..ef007b4e 100644
--- a/netbox_agent/cli.py
+++ b/netbox_agent/cli.py
@@ -10,6 +10,7 @@
 from netbox_agent.vendors.supermicro import SupermicroHost
 from netbox_agent.virtualmachine import VirtualMachine, is_vm
 
+
 MANUFACTURERS = {
     'Dell Inc.': DellHost,
     'HP': HPHost,
@@ -17,37 +18,42 @@
     'Supermicro': SupermicroHost,
     'Quanta Cloud Technology Inc.': QCTHost,
     'Generic': GenericHost,
+    'Virtual Machine': VirtualMachine,
 }
 
 
-def run(config):
+def main():
     dmi = dmidecode.parse()
 
-    if config.virtual.enabled or is_vm(dmi):
-        if not config.virtual.cluster_name:
-            raise Exception('virtual.cluster_name parameter is mandatory because it\'s a VM')
-        server = VirtualMachine(dmi=dmi)
-    else:
-        manufacturer = dmidecode.get_by_type(dmi, 'Chassis')[0].get('Manufacturer')
-        try:
-            server = MANUFACTURERS[manufacturer](dmi=dmi)
-        except KeyError:
-            server = GenericHost(dmi=dmi)
-
     if version.parse(nb.version) < version.parse('2.9'):
-        print('netbox-agent is not compatible with Netbox prior to verison 2.9')
-        return False
+        raise SystemExit('netbox-agent requires Netbox version 2.9 or higher')
 
-    if config.register or config.update_all or config.update_network or \
-       config.update_location or config.update_inventory or config.update_psu:
+    if is_vm(dmi):
+        if not config.virtual.cluster_name:
+            raise SystemExit(
+                'The `virtual.cluster_name` parameter/configuration value must'
+                'be set for virtualized systems.'
+            )
+        kind = 'Virtual'
+    else:
+        chassis = dmidecode.get_by_type(dmi, 'Chassis')
+        kind = chassis[0].get('Manufacturer')
+
+    server_class = MANUFACTURERS.get(kind, GenericHost)
+    server = server_class(dmi)
+
+    if (
+        config.register or
+        config.update_all or
+        config.update_network or
+        config.update_location or
+        config.update_inventory or
+        config.update_psu
+    ):
         server.netbox_create_or_update(config)
+
     if config.debug:
         server.print_debug()
-    return True
-
-
-def main():
-    return run(config)
 
 
 if __name__ == '__main__':
diff --git a/netbox_agent/virtualmachine.py b/netbox_agent/virtualmachine.py
index b32d3bd9..49635710 100644
--- a/netbox_agent/virtualmachine.py
+++ b/netbox_agent/virtualmachine.py
@@ -14,20 +14,17 @@ def is_vm(dmi):
     system = dmidecode.get_by_type(dmi, 'System')[0]
 
     return (
+        config.virtual.enabled or
+        'Google Compute Engine' in system['Product Name'] or
+        'Hyper-V' in bios['Version'] or
+        'QEMU' in system['Manufacturer'] or
+        'RHEV Hypervisor' in system['Product Name'] or
+        'VMware' in system['Manufacturer'] or
+        'VirtualBox' in bios['Version'] or
+        'Xen' in bios['Version'] or
         (
-            'Hyper-V' in bios['Version'] or
-            'Xen' in bios['Version'] or
-            'Google Compute Engine' in system['Product Name']
-        ) or
-        (
-            (
-                'Amazon EC2' in system['Manufacturer'] and
-                not system['Product Name'].endswith('.metal')
-            ) or
-            'RHEV Hypervisor' in system['Product Name'] or
-            'QEMU' in system['Manufacturer'] or
-            'VirtualBox' in bios['Version'] or
-            'VMware' in system['Manufacturer']
+            'Amazon EC2' in system['Manufacturer'] and
+            not system['Product Name'].endswith('.metal')
         )
     )