gh-christianlouis-docuelevate/app/tasks/check_credentials.py

import logging
from app.celery_app import celery
from app.config import settings
from app.utils.notification import notify_credential_failure
import time
import os
import json
import asyncio
import inspect

# Import the test functions from API routes
from app.api.openai import test_openai_connection
from app.api.azure import test_azure_connection
from app.api.dropbox import test_dropbox_token
from app.api.google_drive import test_google_drive_token
from app.api.onedrive import test_onedrive_token

# Import config validation utilities
from app.utils.config_validator import validate_storage_configs, get_provider_status

# Create an enhanced mock Request object for API functions that expect it
class MockRequest:
    """Mock request object with session and other attributes needed for API functions"""
    def __init__(self):
        self.session = {"user": {"id": "credential_checker", "name": "System Credential Checker"}}
        self.app = None
        self.headers = {}
        self.query_params = {}
        self.path_params = {}

    async def json(self):
        return {}

    async def form(self):
        return {}

logger = logging.getLogger(__name__)

# Path to store failure counts
FAILURE_STATE_FILE = os.path.join(settings.workdir, 'credential_failures.json')

def get_failure_state():
    """Read the failure state from file"""
    try:
        if os.path.exists(FAILURE_STATE_FILE):
            with open(FAILURE_STATE_FILE, 'r') as f:
                return json.load(f)
    except Exception as e:
        logger.error(f"Error reading failure state file: {e}")
    
    # Default empty state
    return {}

def save_failure_state(state):
    """Save failure state to file"""
    try:
        with open(FAILURE_STATE_FILE, 'w') as f:
            json.dump(state, f)
    except Exception as e:
        logger.error(f"Error saving failure state file: {e}")

# Helper function to get the inner function without the decorator
def unwrap_decorated_function(func):
    """Get the original function from a decorated function"""
    if hasattr(func, "__wrapped__"):
        return unwrap_decorated_function(func.__wrapped__)
    return func

# Create synchronous versions of the test functions that bypass authentication
def sync_test_openai_connection():
    """Synchronous wrapper for the OpenAI test function that bypasses auth"""
    # Get the original function without the @require_login decorator
    inner_func = unwrap_decorated_function(test_openai_connection)
    request = MockRequest()
    if inspect.iscoroutinefunction(inner_func):
        return asyncio.run(inner_func(request))
    return inner_func(request)

def sync_test_azure_connection():
    """Synchronous wrapper for the Azure test function that bypasses auth"""
    inner_func = unwrap_decorated_function(test_azure_connection)
    request = MockRequest()
    if inspect.iscoroutinefunction(inner_func):
        return asyncio.run(inner_func(request))
    return inner_func(request)

def sync_test_dropbox_token():
    """Synchronous wrapper for the Dropbox test function that bypasses auth"""
    inner_func = unwrap_decorated_function(test_dropbox_token)
    request = MockRequest()
    if inspect.iscoroutinefunction(inner_func):
        return asyncio.run(inner_func(request))
    return inner_func(request)

def sync_test_google_drive_token():
    """Synchronous wrapper for the Google Drive test function that bypasses auth"""
    inner_func = unwrap_decorated_function(test_google_drive_token)
    request = MockRequest()
    if inspect.iscoroutinefunction(inner_func):
        return asyncio.run(inner_func(request))
    return inner_func(request)

def sync_test_onedrive_token():
    """Synchronous wrapper for the OneDrive test function that bypasses auth"""
    inner_func = unwrap_decorated_function(test_onedrive_token)
    request = MockRequest()
    if inspect.iscoroutinefunction(inner_func):
        return asyncio.run(inner_func(request))
    return inner_func(request)

@celery.task
def check_credentials():
    """Check all configured credentials and notify if any are invalid"""
    logger.info("Starting credential check task")
    
    # Load current failure state
    failure_state = get_failure_state()
    
    # Track failures
    failures = []
    
    # Get provider configurations from config_validator
    provider_status = get_provider_status()
    storage_configs = validate_storage_configs()
    
    # Define services with their test functions and configuration status
    services = [
        {
            "name": "OpenAI", 
            "check_func": sync_test_openai_connection,
            "configured": provider_status.get("OpenAI", {}).get("configured", False),
            "config_issues": []  # OpenAI isn't in storage_configs
        },
        {
            "name": "Azure Document Intelligence", 
            "check_func": sync_test_azure_connection,
            "configured": provider_status.get("Azure AI", {}).get("configured", False),
            "config_issues": []  # Azure isn't in storage_configs
        },
        {
            "name": "Dropbox", 
            "check_func": sync_test_dropbox_token,
            "configured": provider_status.get("Dropbox", {}).get("configured", False),
            "config_issues": storage_configs.get("dropbox", [])
        },
        {
            "name": "Google Drive", 
            "check_func": sync_test_google_drive_token,
            "configured": provider_status.get("Google Drive", {}).get("configured", False),
            "config_issues": storage_configs.get("google_drive", [])
        },
        {
            "name": "OneDrive", 
            "check_func": sync_test_onedrive_token,
            "configured": provider_status.get("OneDrive", {}).get("configured", False),
            "config_issues": storage_configs.get("onedrive", [])
        }
    ]
    
    # Check each service
    results = {}
    current_time = int(time.time())
    
    for service in services:
        service_name = service["name"]
        logger.info(f"Checking credentials for {service_name}")
        
        # Skip services that aren't configured
        if not service["configured"]:
            config_issues = service["config_issues"]
            issue_msg = f"Not properly configured" + (f": {', '.join(config_issues)}" if config_issues else "")
            logger.info(f"Skipping {service_name}: {issue_msg}")
            
            results[service_name] = {
                "status": "unconfigured",
                "message": issue_msg
            }
            continue
        
        try:
            # Call the synchronized test function and get the result
            result = service["check_func"]()
            
            # All test functions return a dict with "status" field
            is_valid = result.get("status") == "success"
            error_message = result.get("message", "Unknown error")
            
            # Store the result
            results[service_name] = {
                "status": "valid" if is_valid else "invalid",
                "message": error_message
            }
            
            if not is_valid:
                failures.append(service_name)
                
                # Get current failure count for this service
                service_state = failure_state.get(service_name, {"count": 0, "last_notified": 0})
                service_state["count"] = service_state.get("count", 0) + 1
                
                # Only notify if we haven't reached the notification threshold (3 failures)
                # or if this is the first failure after a recovery
                if service_state["count"] <= 3 or service_state.get("recovered", False):
                    notify_credential_failure(service_name, error_message)
                    service_state["last_notified"] = current_time
                    service_state["recovered"] = False
                    logger.warning(f"{service_name} credentials check failed ({service_state['count']} times): {error_message}")
                else:
                    # We're in cooldown mode
                    logger.warning(f"{service_name} credentials check failed ({service_state['count']} times): {error_message} - notification suppressed")
                
                # Update failure state
                failure_state[service_name] = service_state
            else:
                logger.info(f"{service_name} credentials are valid")
                
                # Check if this was previously failing and now recovered
                if service_name in failure_state and failure_state[service_name].get("count", 0) > 0:
                    logger.info(f"{service_name} has recovered after {failure_state[service_name]['count']} failures")
                    
                    # Mark it as recovered and reset count
                    failure_state[service_name] = {"count": 0, "recovered": True, "last_notified": 0}
                elif service_name in failure_state:
                    # Just make sure recovered flag is cleared if it was there
                    failure_state[service_name]["recovered"] = True
                
        except Exception as e:
            logger.error(f"Error checking {service_name} credentials: {e}", exc_info=True)
            failures.append(service_name)
            error_message = f"Exception during credential check: {str(e)}"
            
            # Get current failure count for this service
            service_state = failure_state.get(service_name, {"count": 0, "last_notified": 0})
            service_state["count"] = service_state.get("count", 0) + 1
            
            # Only notify if we haven't reached the notification threshold or if we just recovered
            if service_state["count"] <= 3 or service_state.get("recovered", False):
                notify_credential_failure(service_name, error_message)
                service_state["last_notified"] = current_time
                service_state["recovered"] = False
                
            # Update failure state
            failure_state[service_name] = service_state
            
            # Store the error result
            results[service_name] = {
                "status": "error",
                "message": error_message
            }
    
    # Save updated failure state
    save_failure_state(failure_state)
    
    # Count only services that were actually checked (configured services)
    configured_services = [s for s in services if s["configured"]]
    num_configured = len(configured_services)
    
    # Summarize results
    logger.info(f"Credential check completed. Configured services: {num_configured}, Valid: {num_configured - len(failures)}, Invalid: {len(failures)}")
    
    return {
        "checked": num_configured,
        "unconfigured": len(services) - num_configured,
        "failures": len(failures),
        "results": results,
        "failure_state": failure_state
    }