Initial push to public repo

netbox/extras/management/commands/run_inventory.py

@@ -0,0 +1,117 @@
+from Exscript.protocols.Exception import LoginFailure
+from getpass import getpass
+from ncclient.transport.errors import AuthenticationError
+
+from django.conf import settings
+from django.core.management.base import BaseCommand, CommandError
+from django.db import transaction
+
+from dcim.models import Device, Module, Site
+
+
+class Command(BaseCommand):
+    help = "Update inventory information for specified devices"
+    username = settings.NETBOX_USERNAME
+    password = settings.NETBOX_PASSWORD
+
+    def add_arguments(self, parser):
+        parser.add_argument('-u', '--username', dest='username', help="Specify the username to use")
+        parser.add_argument('-p', '--password', action='store_true', default=False, help="Prompt for password to use")
+        parser.add_argument('-s', '--site', dest='site', action='append', help="Filter devices by site (include argument once per site)")
+        parser.add_argument('-n', '--name', dest='name', help="Filter devices by name (regular expression)")
+        parser.add_argument('--full', action='store_true', default=False, help="For inventory update for all devices")
+        parser.add_argument('--fake', action='store_true', default=False, help="Do not actually update database")
+
+    def handle(self, *args, **options):
+
+        # Credentials
+        if options['username']:
+            self.username = options['username']
+        if options['password']:
+            self.password = getpass("Password: ")
+
+        device_list = Device.objects.filter()
+
+        # --site: Include only devices belonging to specified site(s)
+        if options['site']:
+            sites = Site.objects.filter(slug__in=options['site'])
+            if sites:
+                site_names = [s.name for s in sites]
+                self.stdout.write("Running inventory for these sites: {}".format(', '.join(site_names)))
+            else:
+                raise CommandError("One or more sites specified but none found.")
+            device_list = device_list.filter(rack__site__in=sites)
+
+        # --name: Filter devices by name matching a regex
+        if options['name']:
+            device_list = device_list.filter(name__iregex=options['name'])
+
+        # --full: Gather inventory data for *all* devices
+        if options['full']:
+            self.stdout.write("WARNING: Running inventory for all devices! Prior data will be overwritten. (--full)")
+
+        # --fake: Gathering data but not updating the database
+        if options['fake']:
+            self.stdout.write("WARNING: Inventory data will not be saved! (--fake)")
+
+        device_count = device_list.count()
+        self.stdout.write("** Found {} devices...".format(device_count))
+
+        for i, device in enumerate(device_list, start=1):
+
+            self.stdout.write("[{}/{}] {}: ".format(i, device_count, device.name), ending='')
+
+            # Skip inactive devices
+            if not device.status:
+                self.stdout.write("Skipped (inactive)")
+                continue
+
+            # Skip devices without primary_ip set
+            if not device.primary_ip:
+                self.stdout.write("Skipped (no primary IP set)")
+                continue
+
+            # Skip devices which have already been inventoried if not doing a full update
+            if device.serial and not options['full']:
+                self.stdout.write("Skipped (Serial: {})".format(device.serial))
+                continue
+
+            RPC = device.get_rpc_client()
+            if not RPC:
+                self.stdout.write("Skipped (no RPC client available for platform {})".format(device.platform))
+                continue
+
+            # Connect to device and retrieve inventory info
+            try:
+                with RPC(device, self.username, self.password) as rpc_client:
+                    inventory = rpc_client.get_inventory()
+            except KeyboardInterrupt:
+                raise
+            except (AuthenticationError, LoginFailure):
+                self.stdout.write("Authentication error!")
+                continue
+            except Exception as e:
+                self.stdout.write("Error for {} ({}): {}".format(device, device.primary_ip.address.ip, e))
+                continue
+
+            self.stdout.write("")
+            self.stdout.write("\tSerial: {}".format(inventory['chassis']['serial']))
+            self.stdout.write("\tDescription: {}".format(inventory['chassis']['description']))
+            for module in inventory['modules']:
+                self.stdout.write("\tModule: {} / {} ({})".format(module['name'], module['part_id'], module['serial']))
+
+            if not options['fake']:
+                with transaction.atomic():
+                    if inventory['chassis']['serial']:
+                        device.serial = inventory['chassis']['serial']
+                        device.save()
+                    Module.objects.filter(device=device).delete()
+                    modules = []
+                    for module in inventory['modules']:
+                        modules.append(Module(device=device,
+                                              name=module['name'],
+                                              part_id=module['part_id'],
+                                              serial=module['serial']))
+                    Module.objects.bulk_create(modules)
+
+        self.stdout.write("Finished!")