minor security improvements

tutorials/apps/testcontainer/app.py

 import panel as pn
 from dashboard import *
-import time, sys, threading
+from verification import *
 token = None
 container_ip = None
 is_valid = None
 err_msg = None
 
+pn.extension()
 
 def build_app():
 
     token = pn.state.session_args.get('token', [None])[0]
     container_ip = pn.state.session_args.get('ip', [None])[0]
-
     if isinstance(token, bytes):
         token = token.decode('utf-8')
     if isinstance(container_ip, bytes):
         container_ip = container_ip.decode('utf-8')
+    
+    if not token or not container_ip:
+        return pn.Column("## Invalid session", "Missing token or container IP.")
+
+    is_valid, err_msg = is_token_valid(token)
+    if not is_valid:
+        return pn.Column("# Access Denied", f"Authentication failed: {err_msg}")
+
+   
+    decoded = verify_cognito_jwt(token, audience=CLIENT_ID)
+    pn.state.cache['user'] = decoded
+
     return run_dashboard(token, container_ip)
 
 
-pn.extension()
 main_layout = build_app()
 main_layout.servable()
     

tutorials/apps/testcontainer/dashboard.py

@@ -17,10 +17,11 @@ import plotly.graph_objs as go
 import holoviews as hv
 
 import jwt, requests
-from verification import is_token_valid, update_cookie, logout
+from verification import is_token_valid, logout
 
 
-def run_dashboard(token, container_ip):    
+def run_dashboard(token, container_ip):
+    
     pn.extension(
             "plotly",
             disconnect_notification='Connection lost.',
@@ -28,7 +29,10 @@ def run_dashboard(token, container_ip):
             template='bootstrap'
         )
     hv.extension('bokeh')
-   
+    
+    user = pn.state.cache.get('user')
+    if not user:
+        return pn.Column("## Unauthorized", "Missing or invalid user session.")
 
 
     is_valid, err_msg = is_token_valid(token)
@@ -79,8 +83,7 @@ def run_dashboard(token, container_ip):
         
         
     else:
-        cookie = pn.pane.HTML()
-        update_cookie(cookie, token)    
+           
         selector_country = pn.widgets.Select(
             description="Select country",
             value='Latvia',
@@ -152,7 +155,7 @@ def run_dashboard(token, container_ip):
         logout_button.on_click(on_logout_click)
         layout = pn.Row(
             pn.Column(
-                selector_country, selector_parameter,ifunc_timeseries, selector_timestep,logout_button,logout_pane, cookie), 
+                selector_country, selector_parameter,ifunc_timeseries, selector_timestep,logout_button,logout_pane), 
             pn.Column(ifunc_polygons, ifunc_rasters))
         
         return layout

tutorials/apps/testcontainer/setup_callbacks.py

@@ -3,7 +3,7 @@ import panel as pn
 import requests
 import logging
 import sys, time, threading, os
-import tornado.ioloop
+
 
 FORMAT = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
 
@@ -58,7 +58,7 @@ def session_destroyed(session_context):
     logger.info("Fired session_destroyed")
     if info and info["ip"]:
         ip = info["ip"]
-        url = f"http://localhost:5000/logout?ip={ip}"
+        url = f"http://atlas.kliima.ut.ee/logout?ip={ip}"
         try:
             r = requests.get(url, timeout=5)
             logger.info(f"Stop container response: {r.status_code}, {r.text}")

tutorials/apps/testcontainer/verification.py

-import jwt, requests, cryptography
-import panel as pn
-from panel.io.server import get_server
+import jwt
+import requests
+from jwt import algorithms
+from datetime import datetime
 from dotenv import load_dotenv
 import os
+
 load_dotenv()
+
 COGNITO_REGION = os.getenv("COGNITO_REGION")
 USER_POOL_ID = os.getenv("COGNITO_POOL_ID")
+CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
 COGNITO_ISS = f"https://cognito-idp.{COGNITO_REGION}.amazonaws.com/{USER_POOL_ID}"
 JWKS_URL = f"{COGNITO_ISS}/.well-known/jwks.json"
-CLIENT_ID = os.getenv("COGNITO_CLIENT_ID")
-response = requests.get(JWKS_URL)
-jwks = response.json()["keys"]
 
-# Helper: find the correct public key for the token's 'kid'
-def get_signing_key(token_kid):
-    for key in jwks:
-        if key["kid"] == token_kid:
+# Cache JWKS for performance
+_jwks_cache = None
+
+def get_jwks():
+    global _jwks_cache
+    if _jwks_cache is None:
+        resp = requests.get(JWKS_URL)
+        _jwks_cache = resp.json()['keys']
+    return _jwks_cache
+
+def get_signing_key(kid):
+    for key in get_jwks():
+        if key['kid'] == kid:
             return key
     return None
 
 def verify_cognito_jwt(token, audience=None):
-    """
-    Verify the given JWT against the Cognito JWKS. 
-    Optionally check the 'audience' claim if verifying ID tokens (client_id).
-    Returns the decoded payload if valid, or raises an exception if invalid.
-    """
-    # 1. Decode the header first to get 'kid'
     unverified_header = jwt.get_unverified_header(token)
     kid = unverified_header.get("kid")
     if not kid:
         raise ValueError("No 'kid' in token header")
 
-    # 2. Find the matching JWKS entry
     key_entry = get_signing_key(kid)
     if not key_entry:
-        raise ValueError("No matching JWKS key found for kid={}".format(kid))
+        raise ValueError(f"No matching JWKS key found for kid={kid}")
 
-    # 3. Construct a public key object
-    public_key = jwt.algorithms.RSAAlgorithm.from_jwk(key_entry)
-    
-    # 4. Decode + verify signature, issuer, audience, etc.
-    options = {
-        "verify_exp": True,
-        "verify_aud": (audience is not None),
-        "verify_iss": True
-    }
+    public_key = algorithms.RSAAlgorithm.from_jwk(key_entry)
     decoded = jwt.decode(
         token,
         public_key,
-        algorithms=[key_entry["alg"]],  # 'RS256'
-        audience=audience,  
+        algorithms=[key_entry['alg']],
+        audience=audience,
         issuer=COGNITO_ISS,
-        options=options
+        options={
+            "verify_exp": True,
+            "verify_iss": True,
+            "verify_aud": audience is not None
+        }
     )
     return decoded
 
-def is_token_valid(tok):
-    """
-    Attempt to verify the JWT with Cognito. 
-    Returns (True, None) if valid, or (False, error_message) if invalid.
-    """
-    if not tok:
+def is_token_valid(token):
+    if not token:
         return (False, "No token provided.")
-
-    if isinstance(tok, bytes):
-        tok = tok.decode("utf-8")
-
     try:
-        _decoded = verify_cognito_jwt(tok, audience=CLIENT_ID)
-        # If no exception, the token is valid
-        
+        verify_cognito_jwt(token, audience=CLIENT_ID)
         return (True, None)
     except Exception as e:
-        # Return False + the exception message
-        return (False, str(e) + jwt.__version__ + " "+jwt.__file__)
-    
-def update_cookie(html,token):
-    html.object = f"""
-    <script>
-      var t = encodeURIComponent("{token}");
-      document.cookie = "token=" + t + "; path=/; Max-Age=900;";
-    </script>
-    """
+        return (False, str(e))
 
 def logout(container_ip):
-    # Returns to login screen and sends a shutdown signal to ECS
-
+    # Provide a client-side redirect to trigger logout
     return f"""
     <script>
-      document.cookie = "token=; path=/; Max-Age=0";
-      window.location.href = "http://localhost:5000/logout?ip={container_ip}";
+      window.location.href = "https://atlas.kliima.ut.ee/logout?ip={container_ip}";
     </script>
     """
-