From de770faf6a99d47be6ed1f657785ff8f8fd5a35d Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Fri, 16 Aug 2019 15:27:58 -0400
Subject: [PATCH] Add FileVar for file uploads

---
 netbox/extras/scripts.py            | 26 +++++++++++++++++++-------
 netbox/extras/views.py              |  4 ++--
 netbox/templates/extras/script.html | 10 ++++++----
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/netbox/extras/scripts.py b/netbox/extras/scripts.py
index c38a795d5..91e387e9b 100644
--- a/netbox/extras/scripts.py
+++ b/netbox/extras/scripts.py
@@ -21,13 +21,14 @@ from .forms import ScriptForm
 
 
 __all__ = [
+    'BooleanVar',
+    'FileVar',
+    'IntegerVar',
+    'IPNetworkVar',
+    'ObjectVar',
     'Script',
     'StringVar',
     'TextVar',
-    'IntegerVar',
-    'BooleanVar',
-    'ObjectVar',
-    'IPNetworkVar',
 ]
 
 
@@ -145,6 +146,13 @@ class ObjectVar(ScriptVariable):
             self.form_field = TreeNodeChoiceField
 
 
+class FileVar(ScriptVariable):
+    """
+    An uploaded file.
+    """
+    form_field = forms.FileField
+
+
 class IPNetworkVar(ScriptVariable):
     """
     An IPv4 or IPv6 prefix.
@@ -193,12 +201,12 @@ class Script:
     def run(self, data):
         raise NotImplementedError("The script must define a run() method.")
 
-    def as_form(self, data=None):
+    def as_form(self, data=None, files=None):
         """
         Return a Django form suitable for populating the context data required to run this Script.
         """
         vars = self._get_vars()
-        form = ScriptForm(vars, data)
+        form = ScriptForm(vars, data, files)
 
         return form
 
@@ -260,7 +268,7 @@ def is_variable(obj):
     return isinstance(obj, ScriptVariable)
 
 
-def run_script(script, data, commit=True):
+def run_script(script, data, files, commit=True):
     """
     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.
@@ -269,6 +277,10 @@ def run_script(script, data, commit=True):
     start_time = None
     end_time = None
 
+    # Add files to form data
+    for field_name, fileobj in files.items():
+        data[field_name] = fileobj
+
     try:
         with transaction.atomic():
             start_time = time.time()
diff --git a/netbox/extras/views.py b/netbox/extras/views.py
index 6ef68b995..e43e4a86d 100644
--- a/netbox/extras/views.py
+++ b/netbox/extras/views.py
@@ -400,13 +400,13 @@ class ScriptView(PermissionRequiredMixin, View):
             return HttpResponseForbidden()
 
         script = self._get_script(module, name)
-        form = script.as_form(request.POST)
+        form = script.as_form(request.POST, request.FILES)
         output = None
         execution_time = None
 
         if form.is_valid():
             commit = form.cleaned_data.pop('_commit')
-            output, execution_time = run_script(script, form.cleaned_data, commit)
+            output, execution_time = run_script(script, form.cleaned_data, request.FILES, commit)
 
         return render(request, 'extras/script.html', {
             'module': module,
diff --git a/netbox/templates/extras/script.html b/netbox/templates/extras/script.html
index ae1f89b49..e230fd974 100644
--- a/netbox/templates/extras/script.html
+++ b/netbox/templates/extras/script.html
@@ -57,9 +57,11 @@
                                     </tr>
                                 {% endfor %}
                             </table>
-                            <div class="panel-footer text-right text-muted">
-                                <small>Exec time: {{ execution_time|floatformat:3 }}s</small>
-                            </div>
+                            {% if execution_time %}
+                                <div class="panel-footer text-right text-muted">
+                                    <small>Exec time: {{ execution_time|floatformat:3 }}s</small>
+                                </div>
+                            {% endif %}
                         </div>
                     </div>
                 </div>
@@ -72,7 +74,7 @@
                             You do not have permission to run scripts.
                         </div>
                     {% endif %}
-                    <form action="" method="post" class="form form-horizontal">
+                    <form action="" method="post" enctype="multipart/form-data" class="form form-horizontal">
                     {% csrf_token %}
                         {% if not form.requires_input %}
                             <div class="alert alert-info">