From e9f28dcd4db1a1399d9702c175c3fcf8f418d455 Mon Sep 17 00:00:00 2001
From: Arthur <worldnomad@gmail.com>
Date: Tue, 13 Feb 2024 15:00:38 -0800
Subject: [PATCH] 14438 update migration

---
 .../extras/migrations/0108_script_models.py   | 60 +++++++++++++++----
 1 file changed, 47 insertions(+), 13 deletions(-)

diff --git a/netbox/extras/migrations/0108_script_models.py b/netbox/extras/migrations/0108_script_models.py
index 5f5a6e13d..7fc9f3873 100644
--- a/netbox/extras/migrations/0108_script_models.py
+++ b/netbox/extras/migrations/0108_script_models.py
@@ -1,10 +1,19 @@
 # Generated by Django 4.2.9 on 2024-02-05 21:37
 import inspect
+import os
 
+from django.conf import settings
 from django.db import migrations, models
 import django.db.models.deletion
+from importlib.machinery import SourceFileLoader
 
 
+#
+# Note: This has a couple dependencies on the codebase if doing future modifications:
+# There are imports from extras.scripts and extras.reports as well as expecting
+# settings.SCRIPTS_ROOT and settings.REPORTS_ROOT to be in settings
+#
+
 def update_event_rules(apps, schema_editor):
     Script = apps.get_model('extras', 'Script')
     ContentType = apps.get_model('contenttypes', 'ContentType')
@@ -25,10 +34,25 @@ def update_event_rules(apps, schema_editor):
         rule.save()
 
 
+def full_path(instance):
+    root_path = {
+        'scripts': settings.SCRIPTS_ROOT,
+        'reports': settings.REPORTS_ROOT,
+    }[instance.file_root]
+    return os.path.join(root_path, instance.file_path)
+
+
+def python_name(instance):
+    path, filename = os.path.split(full_path(instance))
+    name = os.path.splitext(filename)[0]
+    if name == '__init__':
+        # File is a package
+        return os.path.basename(path)
+    else:
+        return name
+
+
 def is_script(obj):
-    """
-    Returns True if the object is a Script or Report.
-    """
     from extras.scripts import Script
     from extras.reports import Report
 
@@ -38,6 +62,12 @@ def is_script(obj):
         return False
 
 
+def get_python_module(instance):
+    loader = SourceFileLoader(python_name(instance), full_path(instance))
+    module = loader.load_module()
+    return module
+
+
 def get_module_scripts(apps, instance):
 
     def _get_name(cls):
@@ -45,9 +75,10 @@ def get_module_scripts(apps, instance):
         return cls.full_name.split(".", maxsplit=1)[1]
 
     try:
-        module = instance.get_module()
+        module = get_python_module(instance)
     except Exception as e:
         module = None
+        return
 
     scripts = {}
     ordered = getattr(module, 'script_order', [])
@@ -55,7 +86,6 @@ def get_module_scripts(apps, instance):
     for cls in ordered:
         scripts[_get_name(cls)] = cls
     for name, cls in inspect.getmembers(module, is_script):
-        breakpoint()
         if cls not in ordered:
             scripts[_get_name(cls)] = cls
 
@@ -66,18 +96,22 @@ def update_scripts(apps, schema_editor):
     ScriptModule = apps.get_model('extras', 'ScriptModule')
     Script = apps.get_model('extras', 'Script')
     ContentType = apps.get_model('contenttypes', 'ContentType')
+    Job = apps.get_model('core', 'Job')
     ct = ContentType.objects.filter(app_label='extras', model='script').first()
+    ct_module = ContentType.objects.filter(app_label='extras', model='scriptmodule').first()
 
     for module in ScriptModule.objects.all():
-        for script in get_module_scripts(apps, module).keys():
-            obj = Script.objects.create(
-                name=script,
-                module=ScriptModule.objects.get(file_root=module.file_root, file_path=module.file_path),
-            )
+        module_scripts = get_module_scripts(apps, module)
+        if module_scripts:
+            for script in module_scripts.keys():
+                obj = Script.objects.create(
+                    name=script,
+                    module=ScriptModule.objects.get(file_root=module.file_root, file_path=module.file_path),
+                )
 
-            # update all jobs associated with this module/name to point to the new script obj
-            if ct:
-                module.jobs.filter(name=script).update(object_type=ct, object_id=obj.id)
+                # update all jobs associated with this module/name to point to the new script obj
+                if ct:
+                    Job.objects.filter(object_type=ct_module, object_id=module.id, name=script).update(object_type=ct, object_id=obj.id)
 
 
 class Migration(migrations.Migration):