netbox_agent/cli.py

@@ -10,6 +10,7 @@ from netbox_agent.vendors.qct import QCTHost
 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 @@ MANUFACTURERS = {
     '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__':

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')
         )
     )