1 files changed, 22 insertions, 8 deletions
diff --git a/app/dispatch/views.py b/app/dispatch/views.py
index 63a6a09..a03868f 100644
--- a/app/dispatch/views.py
+++ b/app/dispatch/views.py
@@ -93,12 +93,14 @@ class DriverList(UserPassesTestMixin, ListView):
     def test_func(self):
         return self.request.user.is_superuser
 
-class DriverDetail(LoadDateSort):
+class DriverDetail(UserPassesTestMixin, LoadDateSort):
     template_name = "dispatch/drivers/detail.html"
     model = User
 
     def test_func(self):
-        return self.request.user.is_superuser
+        # Seems a little hacky at first but it works!
+        return self.request.user.is_superuser or \
+            self.get_object().pk is self.request.user.id
 
     def get_context_data(self, **kwargs):
         # Shit gets fucky with super() really fast, but this seems to work
@@ -126,14 +128,16 @@ class DriverDetail(LoadDateSort):
 
         return context
 
-class DriverUpdate(UpdateView): 
+class DriverUpdate(UserPassesTestMixin, UpdateView): 
     template_name = "dispatch/drivers/edit.html"
     model = User
     success_url = reverse_lazy('driver_list')
     fields = ['username', 'first_name','last_name','email','groups', 'is_active']
 
     def test_func(self):
-        return self.request.user.is_superuser
+        # Seems a little hacky at first but it works!
+        return self.request.user.is_superuser or \
+            self.get_object().pk is self.request.user.id
 
 # Company CRUD
 
@@ -141,23 +145,30 @@ class CompanyList(ListView):
     template_name = "dispatch/companies/list.html"
     model = Company
 
-class CompanyCreate(CreateView):
+class CompanyCreate(UserPassesTestMixin, CreateView):
     template_name = "dispatch/companies/create.html"
     model = Company
     success_url = reverse_lazy('company_list')
     fields = ['name', 'address', 'phone_number','email_address']
 
+    def test_func(self):
+        return self.request.user.is_superuser
+
 class CompanyDetail(LoadDateSort):
     template_name = "dispatch/companies/detail.html"
     model = Company
 
-class CompanyUpdate(UpdateView):
+class CompanyUpdate(UserPassesTestMixin, UpdateView):
     template_name = "dispatch/companies/edit.html"
     model = Company
     success_url = reverse_lazy('company_list')
     fields = ['name', 'address', 'phone_number','email_address']
 
-class CompanyDelete(DeleteView):
+    def test_func(self):
+        return self.request.user.is_superuser
+
+
+class CompanyDelete(UserPassesTestMixin, DeleteView):
     template_name = "dispatch/companies/delete.html"
     model = Company
     success_url = reverse_lazy('company_list')
@@ -227,11 +238,14 @@ class LoadUpdate(FilteredUpdateView):
             load.user = self.request.user
         return super(LoadUpdate, self).form_valid(form)
 
-class LoadDelete(FilteredDeleteView):
+class LoadDelete(UserPassesTestMixin, FilteredDeleteView):
     template_name = "dispatch/loads/delete.html"
     model = Load
     success_url = reverse_lazy('load_list')
 
+    def test_func(self):
+        return self.request.user.is_superuser
+
 
 # Paperwork Uploads