Update okta.py

okta.py

@@ -46,64 +46,125 @@ def create_session_zip(session_id):
         f.write(zip_file.read())
     return zip_filename
 
-def handle_okta_call(api_base_url, api_token, session_id, param_values, *checkbox_values):
-    """Handle Okta API calls with parameter support"""
-    if not session_id:
-        session_id = str(uuid.uuid4())
+def handle_okta_call(api_base_url, api_token, session_id, param_values, endpoints):
+    """
+    Make API calls to Okta endpoints.
     
-    if not api_base_url.startswith('https://'):
-        api_base_url = f"https://{api_base_url}"
+    Args:
+        api_base_url: The base URL for Okta API
+        api_token: The Okta API token for authentication
+        session_id: Current session ID
+        param_values: Dictionary of parameter values for the endpoints
+        endpoints: List of endpoints to call
+        
+    Returns:
+        Tuple of (responses, download_file, session_id, message)
+    """
+    # Strip trailing slashes for consistency
+    api_base_url = api_base_url.rstrip('/')
     
+    # Set up headers with authentication
     headers = {
+        "Authorization": f"SSWS {api_token}",
         "Accept": "application/json",
-        "Authorization": f"SSWS {api_token}"
+        "Content-Type": "application/json"
     }
     
     responses = {}
-    base_save_folder = os.path.join("sessions", session_id, "Okta")
-    os.makedirs(base_save_folder, exist_ok=True)
     
-    # Process selected endpoints
-    for selections in checkbox_values:
-        if isinstance(selections, list):
-            for selection in selections:
+    # Process each endpoint
+    for endpoint in endpoints:
+        try:
+            # Apply path parameters if present
+            formatted_endpoint = endpoint
+            for param_name, param_value in param_values.items():
+                if "{" + param_name + "}" in endpoint:
+                    formatted_endpoint = formatted_endpoint.replace("{" + param_name + "}", param_value)
+            
+            # Build full URL
+            url = f"{api_base_url}{formatted_endpoint}"
+            
+            # Extract query parameters for this endpoint
+            query_params = {}
+            for param_name, param_value in param_values.items():
+                if param_name not in endpoint and param_value:  # Not a path param and has value
+                    query_params[param_name] = param_value
+            
+            # Special handling for credential verification endpoint
+            if endpoint == "/api/v1/users/me":
                 try:
-                    if " | " in selection:
-                        endpoint, method_part = selection.split(" | ")
-                        method = method_part.split(" - ")[0].lower()
+                    r = requests.get(url, headers=headers, timeout=30)
+                    r.raise_for_status()  # This will raise an exception for 4xx/5xx status codes
+                    responses[endpoint] = r.json()
+                except requests.exceptions.HTTPError as e:
+                    # Handle authentication errors specifically
+                    if e.response.status_code in (401, 403):
+                        responses[endpoint] = {"error": f"Authentication failed: {str(e)}"}
                     else:
-                        endpoint = selection
-                        method = "get"
-                    
-                    # Ensure endpoint starts with /api/v1
-                    if not endpoint.startswith('/api/v1'):
-                        endpoint = f"/api/v1{endpoint}"
-                    
-                    # Handle parameter replacement if needed
-                    if any(char in endpoint for char in ['{', '}']):
-                        full_url, error = handle_path_parameters(endpoint, api_base_url, param_values)
-                        if error:
-                            responses[endpoint] = f"Error: {error}"
-                            continue
-                    else:
-                        full_url = f"{api_base_url.rstrip('/')}{endpoint}"
-                    
-                    print(f"Calling Okta endpoint: {full_url}")
-                    
-                    r = requests.get(full_url, headers=headers)
-                    r.raise_for_status()
-                    
-                    data = r.json() if r.headers.get('content-type', '').startswith('application/json') else r.text
-                    responses[endpoint] = data
-                    
-                    # Save response data
-                    save_response_data(data, endpoint, base_save_folder)
-                        
+                        responses[endpoint] = {"error": f"HTTP Error: {str(e)}"}
                 except Exception as e:
-                    responses[endpoint] = f"Error: {traceback.format_exc()}"
+                    responses[endpoint] = {"error": str(e)}
+                continue
+                
+            # Make the API request
+            r = requests.get(url, headers=headers, params=query_params, timeout=30)
+            
+            # Check for successful response
+            if r.status_code == 200:
+                try:
+                    responses[endpoint] = r.json()
+                except ValueError:
+                    responses[endpoint] = {"error": "Invalid JSON response"}
+            else:
+                responses[endpoint] = {"error": f"Error {r.status_code}: {r.text}"}
+                
+        except requests.exceptions.RequestException as e:
+            responses[endpoint] = {"error": f"Request failed: {str(e)}"}
+        except Exception as e:
+            responses[endpoint] = {"error": f"Error: {str(e)}"}
+    
+    # Create a ZIP file with the results
+    if responses and not all(isinstance(resp, dict) and "error" in resp for _, resp in responses.items()):
+        try:
+            import tempfile
+            import zipfile
+            import json
+            import os
+            from datetime import datetime
+            
+            # Create a temporary directory
+            temp_dir = tempfile.mkdtemp()
+            zip_path = os.path.join(temp_dir, "okta_data.zip")
+            
+            # Create ZIP file
+            with zipfile.ZipFile(zip_path, 'w') as zipf:
+                for endpoint, data in responses.items():
+                    # Clean endpoint name for filename
+                    safe_name = endpoint.replace('/', '_').replace('{', '').replace('}', '')
+                    if safe_name.startswith('_'):
+                        safe_name = safe_name[1:]
+                        
+                    filename = f"{safe_name}.json"
+                    json_data = json.dumps(data, indent=2)
+                    zipf.writestr(filename, json_data)
+                
+                # Add metadata
+                metadata = {
+                    "timestamp": datetime.now().isoformat(),
+                    "api_base_url": api_base_url,
+                    "endpoints": endpoints,
+                    "session_id": session_id or str(datetime.now().timestamp())
+                }
+                zipf.writestr("metadata.json", json.dumps(metadata, indent=2))
+                
+            return responses, zip_path, session_id, "✅ API calls completed successfully"
+        except Exception as e:
+            return responses, None, session_id, f"⚠️ API calls completed but export failed: {str(e)}"
+    
+    # Handle case where all responses are errors
+    if all(isinstance(resp.get(endpoint), dict) and "error" in resp.get(endpoint, {}) for resp in [responses]):
+        return responses, None, session_id, "❌ All API calls failed"
     
-    # Create and return session zip
-    zip_filename = create_session_zip(session_id)
-    return responses, zip_filename, session_id, "✅ Okta API calls complete!"
+    return responses, None, session_id, "✅ API calls completed"
 
 # Include save_response_data and create_session_zip functions (same as IdentityNow)