From 0d289d660d794835ee753da2a573e9ae3234a00f Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Mon, 12 Aug 2019 14:28:06 -0400
Subject: [PATCH] Add option to commit database changes

---
 netbox/extras/forms.py              |  9 +++++++++
 netbox/extras/scripts.py            | 28 ++++++++++++++++++++--------
 netbox/extras/views.py              |  3 ++-
 netbox/templates/extras/script.html |  4 ++--
 netbox/utilities/exceptions.py      |  5 +++++
 5 files changed, 38 insertions(+), 11 deletions(-)
 create mode 100644 netbox/utilities/exceptions.py

diff --git a/netbox/extras/forms.py b/netbox/extras/forms.py
index 15c91a880..32ffc1396 100644
--- a/netbox/extras/forms.py
+++ b/netbox/extras/forms.py
@@ -387,6 +387,12 @@ class ObjectChangeFilterForm(BootstrapMixin, forms.Form):
 #
 
 class ScriptForm(BootstrapMixin, forms.Form):
+    _commit = forms.BooleanField(
+        required=False,
+        initial=True,
+        label="Commit changes",
+        help_text="Commit changes to the database (uncheck for a dry-run)"
+    )
 
     def __init__(self, vars, *args, **kwargs):
 
@@ -395,3 +401,6 @@ class ScriptForm(BootstrapMixin, forms.Form):
         # Dynamically populate fields for variables
         for name, var in vars.items():
             self.fields[name] = var.as_field()
+
+        # Move _commit to the end of the form
+        self.fields.move_to_end('_commit', True)
diff --git a/netbox/extras/scripts.py b/netbox/extras/scripts.py
index 7814163b7..49ea43abb 100644
--- a/netbox/extras/scripts.py
+++ b/netbox/extras/scripts.py
@@ -7,6 +7,7 @@ from django.conf import settings
 from django.core.validators import RegexValidator
 from django.db import transaction
 
+from utilities.exceptions import AbortTransaction
 from .constants import LOG_DEFAULT, LOG_FAILURE, LOG_INFO, LOG_SUCCESS, LOG_WARNING
 from .forms import ScriptForm
 
@@ -197,21 +198,32 @@ def is_variable(obj):
     return isinstance(obj, ScriptVariable)
 
 
-def run_script(script, data=None):
+def run_script(script, data, commit=True):
     """
-    A wrapper for calling Script.run(). This performs error handling. It exists outside of the Script class to ensure
-    it cannot be overridden by a script author.
+    A wrapper for calling Script.run(). This performs error handling and provides a hook for committing changes. It
+    exists outside of the Script class to ensure it cannot be overridden by a script author.
     """
+    output = None
+
     try:
         with transaction.atomic():
-            return script.run(data)
+            output = script.run(data)
+            if not commit:
+                raise AbortTransaction()
+    except AbortTransaction:
+        pass
     except Exception as e:
         script.log_failure(
-            "An exception occurred: {}".format(e)
-        )
-        script.log_info(
-            "Database changes have been reverted automatically."
+            "An exception occurred. {}: {}".format(type(e).__name__, e)
         )
+        commit = False
+    finally:
+        if not commit:
+            script.log_info(
+                "Database changes have been reverted automatically."
+            )
+
+    return output
 
 
 def get_scripts():
diff --git a/netbox/extras/views.py b/netbox/extras/views.py
index b359077cf..6b87c27de 100644
--- a/netbox/extras/views.py
+++ b/netbox/extras/views.py
@@ -404,7 +404,8 @@ class ScriptView(PermissionRequiredMixin, View):
         output = None
 
         if form.is_valid():
-            run_script(script)
+            commit = form.cleaned_data.pop('_commit')
+            run_script(script, form.cleaned_data, commit)
 
         return render(request, 'extras/script.html', {
             'module': module,
diff --git a/netbox/templates/extras/script.html b/netbox/templates/extras/script.html
index 87f389c75..fcb2428f4 100644
--- a/netbox/templates/extras/script.html
+++ b/netbox/templates/extras/script.html
@@ -56,7 +56,7 @@
                 </div>
             {% endif %}
             <div class="row">
-                <div class="col-md-8 col-md-offset-2">
+                <div class="col-md-6 col-md-offset-3">
                     {% if not perms.extras.run_script %}
                         <div class="alert alert-warning">
                             <i class="fa fa-warning"></i>
@@ -82,7 +82,7 @@
             <pre>{{ output }}</pre>
         </div>
         <div role="tabpanel" class="tab-pane" id="source">
-            <strong>{{ script.filename }}</strong>
+            <p><code>{{ script.filename }}</code></p>
             <pre>{{ script.source }}</pre>
         </div>
     </div>
diff --git a/netbox/utilities/exceptions.py b/netbox/utilities/exceptions.py
new file mode 100644
index 000000000..5032aacee
--- /dev/null
+++ b/netbox/utilities/exceptions.py
@@ -0,0 +1,5 @@
+class AbortTransaction(Exception):
+    """
+    A dummy exception used to trigger a database transaction rollback.
+    """
+    pass