improve allowed site check

companion.py

@@ -58,7 +58,7 @@ def askAllowSite(sitename):
         return True; # report success
     return False;
 
-async def handleJson(websocket, requestjson):
+async def handleJson(websocket, requestjson, allowed):
     global FILES
     global DOWNLOAD_DIR
     global ALLOWED_SITES
@@ -75,27 +75,46 @@ async def handleJson(websocket, requestjson):
                     "requestID": requestjson["requestID"],
                     "type": "authentication-status",
                     "payload": "ACCEPTED"
-                    }
+                }
+                await send(websocket, json.dumps(responsejson))
+                return True
             else:
                 print("-> REJECTED SITE: " + requestjson["payload"]["payload"]["siteTitle"])
                 responsejson = {
                     "requestID": requestjson["requestID"],
                     "type": "authentication-status",
                     "payload": "REJECTED"
-                    }
-            await send(websocket, json.dumps(responsejson))
+                }
+                await send(websocket, json.dumps(responsejson))
+                print("-> Close websocket")
+                websocket.close()
+                return False
 
         # cloud hosted
         elif(provider == "jwt"):
-            print("-> ACCEPTED JWT: " + requestjson["payload"]["payload"])
-            responsejson = {
-                "requestID": requestjson["requestID"],
-                "type": "authentication-status",
-                "payload": "ACCEPTED"
-            }
-            await send(websocket, json.dumps(responsejson))
-
-    elif(requestjson["type"] == "new-transaction" and requestjson["payload"]["transactionType"] == "file"):
+            currenSiteOrigin = json.loads(base64ToString(requestjson["payload"]["payload"].split(".")[1]))["confluence_origin"]
+            if(currenSiteOrigin in ALLOWED_SITES or askAllowSite(currenSiteOrigin)):
+                print("-> ACCEPTED JWT: " + requestjson["payload"]["payload"])
+                responsejson = {
+                    "requestID": requestjson["requestID"],
+                    "type": "authentication-status",
+                    "payload": "ACCEPTED"
+                }
+                await send(websocket, json.dumps(responsejson))
+                return True
+            else:
+                print("-> REJECTED JWT: " + requestjson["payload"]["payload"])
+                responsejson = {
+                    "requestID": requestjson["requestID"],
+                    "type": "authentication-status",
+                    "payload": "REJECTED"
+                }
+                await send(websocket, json.dumps(responsejson))
+                print("-> Close websocket")
+                websocket.close()
+                return False
+
+    elif(allowed and requestjson["type"] == "new-transaction" and requestjson["payload"]["transactionType"] == "file"):
         newUuid = str(uuid.uuid4())
         print("-> Start new transaction with uuid: "+newUuid)
         responsejson = {
@@ -104,7 +123,7 @@ async def handleJson(websocket, requestjson):
         }
         await send(websocket, json.dumps(responsejson))
 
-    elif(requestjson["type"] == "list-apps"):
+    elif(allowed and requestjson["type"] == "list-apps"):
         responsejson = {
             "requestID": requestjson["requestID"],
             "payload": [{
@@ -116,7 +135,7 @@ async def handleJson(websocket, requestjson):
         }
         await send(websocket, json.dumps(responsejson))
 
-    elif(requestjson["type"] == "launch-file-in-app"):
+    elif(allowed and requestjson["type"] == "launch-file-in-app"):
         appId = requestjson["payload"]["applicationID"]
         transId = requestjson["transactionID"]
         fileUrl = requestjson["payload"]["fileURL"]
@@ -166,7 +185,7 @@ async def handleJson(websocket, requestjson):
         subprocess.call(["xdg-open", filePath])
 
     # upload handler for self-hosted instances
-    elif(requestjson["type"] == "upload-file-in-app"):
+    elif(allowed and requestjson["type"] == "upload-file-in-app"):
         transId = requestjson["transactionID"]
         fileUrl = requestjson["payload"]["uploadUrl"]
 
@@ -240,7 +259,7 @@ async def handleJson(websocket, requestjson):
         await send(websocket, json.dumps(responsejson))
 
     # upload handler for cloud-hosted instances
-    elif(requestjson["type"] == "request-upload-token"):
+    elif(allowed and requestjson["type"] == "request-upload-token"):
         # yummy, we got an upload token!
         transId = requestjson["transactionID"]
         uploadToken = requestjson["payload"]
@@ -433,10 +452,13 @@ def process_IN_MODIFY(self, event):
                     self._loop.run_until_complete(task)
 
 async def companionHandler(websocket, path):
+    allowed = False # indicates if the current connection is authenticated
     while(True):
         request = await websocket.recv()
         print(f"< {request}")
-        await handleJson( websocket, json.loads(request) )
+        result = await handleJson( websocket, json.loads(request), allowed )
+        if(result == True or result == False): # result can also be None - no change in this case
+            allowed = result
 
 async def send(websocket, response):
     await websocket.send(response)