Fix up permissions throughout the system.

4 files changed, 30 insertions, 10 deletions

diff --git a/app/dispatch/templates/dispatch/companies/list.html b/app/dispatch/templates/dispatch/companies/list.html
index 2654587..e0a128f 100644
--- a/app/dispatch/templates/dispatch/companies/list.html
+++ b/app/dispatch/templates/dispatch/companies/list.html
@@ -28,7 +28,9 @@
                 <td>{{ company.phone_number }}</td>
                 <td>{{ company.email_address }}</td>
                 <td class="right-align">
+                    {% if user.is_superuser %}
                     <a href="{% url 'company_edit' company.id %}" class="btn orange">Edit</a>
+                    {% endif %}
                     <a href="{% url 'company_detail' company.id %}" class="btn blue">View</a>
                 </td>
             </tr>
diff --git a/app/dispatch/templates/dispatch/drivers/list.html b/app/dispatch/templates/dispatch/drivers/list.html
index ccfb50d..4153fe9 100644
--- a/app/dispatch/templates/dispatch/drivers/list.html
+++ b/app/dispatch/templates/dispatch/drivers/list.html
@@ -23,7 +23,9 @@
                 <td>{{ driver.last_name }}</td>
                 <td>{{ driver.email }}</td>
                 <td class="right-align">
+                    {% if user.is_superuser %}
                     <a href="{% url 'driver_edit' driver.id %}" class="btn orange">Edit</a>
+                    {% endif %}
                     <a href="{% url 'driver_detail' driver.id %}" class="btn blue">View</a>
                 </td>
             </tr>
diff --git a/app/dispatch/templates/dispatch/nav.html b/app/dispatch/templates/dispatch/nav.html
index 5ba40d9..d40224c 100644
--- a/app/dispatch/templates/dispatch/nav.html
+++ b/app/dispatch/templates/dispatch/nav.html
@@ -1,8 +1,10 @@
-{% if user.is_authenticated %}
+{% if request.user.is_authenticated %}
   <li><a href="{% url 'load_list' %}">Loads</a></li>
-  {% if user.is_superuser %}
+  {% if request.user.is_superuser %}
   <li><a href="{% url 'driver_list' %}">Drivers</a></li>
   <li><a href="{% url 'company_list' %}">Companies</a></li>
+  {% else %}
+  <li><a href="{% url 'driver_edit' request.user.id %}">My Account</a></li>
   {% endif %}
   <li><a href="{% url 'logout' %}">Logout</a></li>
 {% else %}
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