From 07952babe8c9b902bbfd481d33b1aefd5b6c2d6b Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <jstretch@digitalocean.com>
Date: Mon, 28 Mar 2016 17:56:00 -0400
Subject: [PATCH] Expanded inventory function to support submodules

---
 netbox/dcim/views.py                          |  2 +-
 .../management/commands/run_inventory.py      | 15 ++++-----
 netbox/extras/rpc.py                          | 31 ++++++++++++-------
 netbox/templates/dcim/device_inventory.html   | 15 ++++++---
 4 files changed, 39 insertions(+), 24 deletions(-)

diff --git a/netbox/dcim/views.py b/netbox/dcim/views.py
index 95ddc1c91..3542d388c 100644
--- a/netbox/dcim/views.py
+++ b/netbox/dcim/views.py
@@ -760,7 +760,7 @@ class DeviceBulkDeleteView(PermissionRequiredMixin, BulkDeleteView):
 def device_inventory(request, pk):
 
     device = get_object_or_404(Device, pk=pk)
-    modules = Module.objects.filter(device=device)
+    modules = Module.objects.filter(device=device, parent=None).prefetch_related('submodules')
 
     return render(request, 'dcim/device_inventory.html', {
         'device': device,
diff --git a/netbox/extras/management/commands/run_inventory.py b/netbox/extras/management/commands/run_inventory.py
index ed85dd305..66dcb6125 100644
--- a/netbox/extras/management/commands/run_inventory.py
+++ b/netbox/extras/management/commands/run_inventory.py
@@ -24,6 +24,13 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
 
+        def create_modules(modules, parent=None):
+            for module in modules:
+                m = Module(device=device, parent=parent, name=module['name'], part_id=module['part_id'],
+                           serial=module['serial'])
+                m.save()
+                create_modules(module.get('modules', []), parent=m)
+
         # Credentials
         if options['username']:
             self.username = options['username']
@@ -106,12 +113,6 @@ class Command(BaseCommand):
                         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)
+                    create_modules(inventory.get('modules', []))
 
         self.stdout.write("Finished!")
diff --git a/netbox/extras/rpc.py b/netbox/extras/rpc.py
index b8fa798f6..59b076af7 100644
--- a/netbox/extras/rpc.py
+++ b/netbox/extras/rpc.py
@@ -95,6 +95,24 @@ class JunosNC(RPCClient):
 
     def get_inventory(self):
 
+        def glean_modules(node, depth=0):
+            modules = []
+            modules_list = node.get('chassis{}-module'.format('-sub' * depth), [])
+            # Junos like to return single children directly instead of as a single-item list
+            if hasattr(modules_list, 'items'):
+                modules_list = [modules_list]
+            for module in modules_list:
+                m = {
+                    'name': module['name'],
+                    'part_id': module.get('model-number', ''),
+                    'serial': module.get('serial-number', ''),
+                }
+                submodules = glean_modules(module, depth + 1)
+                if submodules:
+                    m['modules'] = submodules
+                modules.append(m)
+            return modules
+
         rpc_reply = self.manager.dispatch('get-chassis-inventory')
         inventory_raw = xmltodict.parse(rpc_reply.xml)['rpc-reply']['chassis-inventory']['chassis']
 
@@ -107,18 +125,7 @@ class JunosNC(RPCClient):
         }
 
         # Gather modules
-        result['modules'] = []
-        for module in inventory_raw['chassis-module']:
-            try:
-                # Skip built-in modules
-                if module['name'] and module['serial-number'] != inventory_raw['serial-number']:
-                    result['modules'].append({
-                        'name': module['name'],
-                        'part_id': module['model-number'] or '',
-                        'serial': module['serial-number'] or '',
-                    })
-            except KeyError:
-                pass
+        result['modules'] = glean_modules(inventory_raw)
 
         return result
 
diff --git a/netbox/templates/dcim/device_inventory.html b/netbox/templates/dcim/device_inventory.html
index 3ead33f04..47f432186 100644
--- a/netbox/templates/dcim/device_inventory.html
+++ b/netbox/templates/dcim/device_inventory.html
@@ -51,12 +51,19 @@
                     </tr>
                 </thead>
                 <tbody>
-                    {% for module in modules %}
+                    {% for m in modules %}
                         <tr>
-                            <td>{{ module.name }}</td>
-                            <td>{{ module.part_id }}</td>
-                            <td>{{ module.serial }}</td>
+                            <td>{{ m.name }}</td>
+                            <td>{{ m.part_id }}</td>
+                            <td>{{ m.serial }}</td>
                         </tr>
+                        {% for m2 in m.submodules.all %}
+                        <tr>
+                            <td style="padding-left: 20px">{{ m2.name }}</td>
+                            <td>{{ m2.part_id }}</td>
+                            <td>{{ m2.serial }}</td>
+                        </tr>
+                        {% endfor %}
                     {% endfor %}
                 </tbody>
             </table>