From bf6cfd22b41e5e224cec7b7ea9764d89298c94cc Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@netboxlabs.com>
Date: Fri, 24 Mar 2023 15:10:52 -0400
Subject: [PATCH] Support automatic migration for submodules

---
 netbox/core/models/files.py                   |  3 --
 .../migrations/0091_create_managedfiles.py    | 35 +++++++++++++------
 netbox/extras/models/models.py                | 25 ++++++++-----
 3 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/netbox/core/models/files.py b/netbox/core/models/files.py
index f17d58cb4..3cb9e5ced 100644
--- a/netbox/core/models/files.py
+++ b/netbox/core/models/files.py
@@ -56,9 +56,6 @@ class ManagedFile(SyncedDataMixin, models.Model):
             models.Index(fields=('file_root', 'file_path'), name='core_managedfile_root_path'),
         ]
 
-    def __str__(self):
-        return f'{self.get_file_root_display()}: {self.file_path}'
-
     def get_absolute_url(self):
         return reverse('core:managedfile', args=[self.pk])
 
diff --git a/netbox/extras/migrations/0091_create_managedfiles.py b/netbox/extras/migrations/0091_create_managedfiles.py
index 0bd569dd8..71581614a 100644
--- a/netbox/extras/migrations/0091_create_managedfiles.py
+++ b/netbox/extras/migrations/0091_create_managedfiles.py
@@ -1,3 +1,4 @@
+import os
 import pkgutil
 
 from django.conf import settings
@@ -5,28 +6,40 @@ from django.db import migrations, models
 import extras.models.models
 
 
-def create_files(cls, root_name, path):
+def create_files(cls, root_name, root_path):
 
-    modules = list(pkgutil.iter_modules([path]))
-    filenames = [f'{m.name}.py' for m in modules]
+    path_tree = [
+        path for path, _, _ in os.walk(root_path)
+        if os.path.basename(path)[0] not in ('_', '.')
+    ]
+
+    modules = list(pkgutil.iter_modules(path_tree))
+    filenames = []
+    for importer, module_name, is_pkg in modules:
+        if is_pkg:
+            continue
+        try:
+            module = importer.find_module(module_name).load_module(module_name)
+            rel_path = os.path.relpath(module.__file__, root_path)
+            filenames.append(rel_path)
+        except ImportError:
+            pass
 
     managed_files = [
-        cls(
-            file_root=root_name,
-            file_path=filename
-        ) for filename in filenames
+        cls(file_root=root_name, file_path=filename)
+        for filename in filenames
     ]
     cls.objects.bulk_create(managed_files)
 
 
 def replicate_scripts(apps, schema_editor):
-    ManagedFile = apps.get_model('core', 'ManagedFile')
-    create_files(ManagedFile, 'scripts', settings.SCRIPTS_ROOT)
+    ScriptModule = apps.get_model('extras', 'ScriptModule')
+    create_files(ScriptModule, 'scripts', settings.SCRIPTS_ROOT)
 
 
 def replicate_reports(apps, schema_editor):
-    ManagedFile = apps.get_model('core', 'ManagedFile')
-    create_files(ManagedFile, 'reports', settings.REPORTS_ROOT)
+    ReportModule = apps.get_model('extras', 'ReportModule')
+    create_files(ReportModule, 'reports', settings.REPORTS_ROOT)
 
 
 class Migration(migrations.Migration):
diff --git a/netbox/extras/models/models.py b/netbox/extras/models/models.py
index 30fb14462..a24a473e2 100644
--- a/netbox/extras/models/models.py
+++ b/netbox/extras/models/models.py
@@ -1,5 +1,6 @@
 import inspect
 import json
+import os
 import uuid
 from functools import cached_property
 from pkgutil import ModuleInfo, get_importer
@@ -824,9 +825,11 @@ class ConfigRevision(models.Model):
 class PythonModuleMixin:
 
     def get_module_info(self):
+        path = os.path.dirname(self.full_path)
+        module_name = os.path.splitext(os.path.basename(self.file_path))[0]
         return ModuleInfo(
-            module_finder=get_importer(self.file_root),
-            name=self.file_path.split('.py')[0],
+            module_finder=get_importer(path),
+            name=module_name,
             ispkg=False
         )
 
@@ -859,9 +862,8 @@ class ScriptModule(PythonModuleMixin, ManagedFile):
     class Meta:
         proxy = True
 
-    def save(self, *args, **kwargs):
-        self.file_root = SCRIPTS_ROOT_NAME
-        return super().save(*args, **kwargs)
+    def __str__(self):
+        return self.file_path
 
     def get_absolute_url(self):
         return reverse('extras:script_list')
@@ -882,6 +884,10 @@ class ScriptModule(PythonModuleMixin, ManagedFile):
 
         return scripts
 
+    def save(self, *args, **kwargs):
+        self.file_root = SCRIPTS_ROOT_NAME
+        return super().save(*args, **kwargs)
+
 
 #
 # Reports
@@ -910,9 +916,8 @@ class ReportModule(PythonModuleMixin, ManagedFile):
     class Meta:
         proxy = True
 
-    def save(self, *args, **kwargs):
-        self.file_root = REPORTS_ROOT_NAME
-        return super().save(*args, **kwargs)
+    def __str__(self):
+        return self.file_path
 
     def get_absolute_url(self):
         return reverse('extras:report_list')
@@ -932,3 +937,7 @@ class ReportModule(PythonModuleMixin, ManagedFile):
             reports[child_name] = cls
 
         return reports
+
+    def save(self, *args, **kwargs):
+        self.file_root = REPORTS_ROOT_NAME
+        return super().save(*args, **kwargs)