fix pcu site check

diff --git a/PLC/Methods/DeletePCU.py b/PLC/Methods/DeletePCU.py
index a6320ed..ae52b43 100644 (file)
--- a/PLC/Methods/DeletePCU.py
+++ b/PLC/Methods/DeletePCU.py
@@ -30,14 +30,8 @@ class DeletePCU(Method):
         pcu = pcus[0]
 
         if 'admin' not in self.caller['roles']:
-            ok = False
-            sites = Sites(self.api, self.caller['site_ids']).values()
-            for site in sites:
-                if pcu['pcu_id'] in site['pcu_ids']:
-                    ok = True
-                    break
-            if not ok:
-                raise PLCPermissionDenied, "Not allowed to delete that PCU"
+            if pcu['site_id'] not in self.caller['site_ids']:
+                raise PLCPermissionDenied, "Not allowed to update that PCU"
 
         pcu.delete()
 

diff --git a/PLC/Methods/GetPCUs.py b/PLC/Methods/GetPCUs.py
index 17b9670..6da58a0 100644 (file)
--- a/PLC/Methods/GetPCUs.py
+++ b/PLC/Methods/GetPCUs.py
@@ -25,15 +25,20 @@ class GetPCUs(Method):
     def call(self, auth, pcu_ids = None):
        # If we are not admin, make sure to only return our own PCUs
         if 'admin' not in self.caller['roles']:
+            # Get list of PCUs that we are able to view
+            sites = Sites(self.api, self.caller['site_ids']).values()
+
+            valid_pcu_ids = set()
+            for site in sites:
+                valid_pcu_ids = valid_pcu_ids.union(site['pcu_ids'])
+
             if not pcu_ids:
-                pcu_ids = []
-                sites = Sites(self.api, self.caller['site_ids']).values()
-                for site in sites:
-                    pcu_ids = set(pcu_ids).union(site['pcu_ids'])
+                pcu_ids = valid_pcu_ids
+            else:
+                pcu_ids = valid_pcu_ids.intersection(pcu_ids)
 
         pcus = PCUs(self.api, pcu_ids).values()
 
-       # turn each pcu into a real dict
        pcus = [dict(pcu) for pcu in pcus]
 
        return pcus

diff --git a/PLC/Methods/UpdatePCU.py b/PLC/Methods/UpdatePCU.py
index 4459b80..f133af9 100644 (file)
--- a/PLC/Methods/UpdatePCU.py
+++ b/PLC/Methods/UpdatePCU.py
@@ -39,13 +39,7 @@ class UpdatePCU(Method):
         pcu = pcus[0]
 
         if 'admin' not in self.caller['roles']:
-            ok = False
-            sites = Sites(self.api, self.caller['site_ids']).values()
-            for site in sites:
-                if pcu['pcu_id'] in site['pcu_ids']:
-                    ok = True
-                    break
-            if not ok:
+            if pcu['site_id'] not in self.caller['site_ids']:
                 raise PLCPermissionDenied, "Not allowed to update that PCU"
 
         pcu.update(pcu_fields)