diff_app_services_api_key_security.py

diff --git a/app/services/api_key_security.py b/app/services/api_key_security.py
index e8ab3984..c4282fdb 100644
--- a/app/services/api_key_security.py
+++ b/app/services/api_key_security.py
@@ -1,372 +1,401 @@
-"""
-API Key Security Service
-Handles key rotation, IP whitelisting, and usage tracking
-"""
-import logging
-from datetime import datetime, timedelta
-from typing import Dict, List, Optional, Tuple, Any
-from sqlalchemy.orm import Session
-from sqlalchemy import func, and_, desc
-
-from app.models import ApiKey, ApiKeyUsage, Tenant
-from app.config.api_keys import (
-    get_api_key_rotation_days,
-    get_rotation_warning_days,
-    is_rotation_enforced,
-    is_ip_whitelist_enabled,
-    is_usage_tracking_enabled
-)
-import uuid
-import ipaddress
-
-logger = logging.getLogger(__name__)
-
-
-class ApiKeySecurityService:
-    """
-    Service for managing API key security features
-
-    Features:
-    - Key rotation policy enforcement
-    - IP address whitelisting
-    - Usage tracking and analytics
-    - Security status monitoring
-    """
-
-    def __init__(self):
-        self.logger = logger
-
-    def check_rotation_status(self, api_key: ApiKey) -> Dict[str, Any]:
-        """
-        Check the rotation status of an API key
-
-        Args:
-            api_key: ApiKey instance
-
-        Returns:
-            Dict with rotation status information
-        """
-        if not api_key.last_rotated_at:
-            return {
-                "status": "expired",
-                "days_since_rotation": None,
-                "days_until_expiry": 0,
-                "message": "Key has never been rotated"
-            }
-
-        rotation_days = get_api_key_rotation_days()
-        warning_days = get_rotation_warning_days()
-
-        now = datetime.now()
-        days_since_rotation = (now - api_key.last_rotated_at).days
-        days_until_expiry = rotation_days - days_since_rotation
-
-        if days_until_expiry <= 0:
-            status = "expired"
-            message = f"Key expired {abs(days_until_expiry)} days ago"
-        elif days_until_expiry <= warning_days:
-            status = "warning"
-            message = f"Key expires in {days_until_expiry} days"
-        else:
-            status = "ok"
-            message = f"Key is valid for {days_until_expiry} days"
-
-        return {
-            "status": status,
-            "days_since_rotation": days_since_rotation,
-            "days_until_expiry": days_until_expiry,
-            "message": message,
-            "rotation_required": days_until_expiry <= 0 and is_rotation_enforced()
-        }
-
-    def is_key_rotation_required(self, api_key: ApiKey) -> bool:
-        """
-        Check if key rotation is required
-
-        Args:
-            api_key: ApiKey instance
-
-        Returns:
-            bool: True if rotation is required
-        """
-        if not is_rotation_enforced():
-            return False
-
-        status = self.check_rotation_status(api_key)
-        return status["days_until_expiry"] <= 0
-
-    def validate_ip_whitelist(self, ip_list: List[str]) -> Tuple[bool, List[str]]:
-        """
-        Validate a list of IP addresses for whitelisting
-
-        Args:
-            ip_list: List of IP addresses/CIDR ranges
-
-        Returns:
-            Tuple of (is_valid, error_messages)
-        """
-        if not ip_list:
-            return True, []
-
-        errors = []
-
-        for ip in ip_list:
-            try:
-                # Try to parse as IP address or network
-                ipaddress.ip_network(ip, strict=False)
-            except ValueError as e:
-                errors.append(f"Invalid IP/CIDR '{ip}': {str(e)}")
-
-        # Check for duplicates
-        if len(ip_list) != len(set(ip_list)):
-            errors.append("Duplicate IP addresses found")
-
-        # Check maximum allowed IPs
-        from app.config.api_keys import get_api_key_config
-        config = get_api_key_config()
-        if len(ip_list) > config["max_whitelist_ips"]:
-            errors.append(f"Too many IP addresses (max: {config['max_whitelist_ips']})")
-
-        return len(errors) == 0, errors
-
-    def is_ip_allowed(self, client_ip: str, api_key: ApiKey) -> bool:
-        """
-        Check if client IP is allowed for the API key
-
-        Args:
-            client_ip: Client IP address
-            api_key: ApiKey instance
-
-        Returns:
-            bool: True if IP is allowed
-        """
-        if not is_ip_whitelist_enabled():
-            return True
-
-        if not api_key.ip_whitelist or len(api_key.ip_whitelist) == 0:
-            # No whitelist means all IPs allowed
-            return True
-
-        try:
-            client_addr = ipaddress.ip_address(client_ip)
-
-            for allowed_ip in api_key.ip_whitelist:
-                try:
-                    allowed_network = ipaddress.ip_network(allowed_ip, strict=False)
-                    if client_addr in allowed_network:
-                        return True
-                except ValueError:
-                    self.logger.warning(f"Invalid IP in whitelist: {allowed_ip}")
-                    continue
-
-            return False
-
-        except ValueError:
-            self.logger.warning(f"Invalid client IP: {client_ip}")
-            return False
-
-    async def track_usage(
-        self,
-        db: Session,
-        api_key_id: str,
-        tenant_id: str,
-        endpoint: str
-    ) -> None:
-        """
-        Track API key usage for an endpoint
-
-        Args:
-            db: Database session
-            api_key_id: API key ID
-            tenant_id: Tenant ID
-            endpoint: API endpoint called
-        """
-        if not is_usage_tracking_enabled():
-            return
-
-        try:
-            # Find or create usage record
-            usage = db.query(ApiKeyUsage).filter(
-                and_(
-                    ApiKeyUsage.api_key_id == api_key_id,
-                    ApiKeyUsage.endpoint == endpoint
-                )
-            ).first()
-
-            if usage:
-                # Update existing record
-                usage.count += 1
-                usage.last_used = func.now()
-            else:
-                # Create new usage record
-                usage = ApiKeyUsage(
-                    id=str(uuid.uuid4()),
-                    api_key_id=api_key_id,
-                    tenant_id=tenant_id,
-                    endpoint=endpoint,
-                    count=1
-                )
-                db.add(usage)
-
-            db.commit()
-
-        except Exception as e:
-            self.logger.error(f"Failed to track API key usage: {str(e)}")
-            db.rollback()
-
-    def get_usage_stats(self, db: Session, api_key_id: str) -> Dict[str, Any]:
-        """
-        Get usage statistics for an API key
-
-        Args:
-            db: Database session
-            api_key_id: API key ID
-
-        Returns:
-            Dict with usage statistics
-        """
-        # Get all usage records for this API key
-        usage_records = db.query(ApiKeyUsage).filter(
-            ApiKeyUsage.api_key_id == api_key_id
-        ).all()
-
-        if not usage_records:
-            return {
-                "api_key_id": api_key_id,
-                "total_requests": 0,
-                "unique_endpoints": 0,
-                "last_used": None,
-                "most_used_endpoint": None,
-                "endpoint_breakdown": []
-            }
-
-        total_requests = sum(record.count for record in usage_records)
-        unique_endpoints = len(usage_records)
-        last_used = max(record.last_used for record in usage_records)
-
-        # Sort by usage count
-        sorted_records = sorted(usage_records, key=lambda r: r.count, reverse=True)
-        most_used_endpoint = sorted_records[0].endpoint if sorted_records else None
-
-        # Create endpoint breakdown
-        endpoint_breakdown = []
-        for record in sorted_records:
-            percentage = (record.count / total_requests * 100) if total_requests > 0 else 0
-            endpoint_breakdown.append({
-                "endpoint": record.endpoint,
-                "count": record.count,
-                "percentage": round(percentage, 2),
-                "last_used": record.last_used
-            })
-
-        return {
-            "api_key_id": api_key_id,
-            "total_requests": total_requests,
-            "unique_endpoints": unique_endpoints,
-            "last_used": last_used,
-            "most_used_endpoint": most_used_endpoint,
-            "endpoint_breakdown": endpoint_breakdown
-        }
-
-    def get_tenant_usage_summary(
-        self,
-        db: Session,
-        tenant_id: str,
-        days: int = 30
-    ) -> Dict[str, Any]:
-        """
-        Get usage summary for all API keys of a tenant
-
-        Args:
-            db: Database session
-            tenant_id: Tenant ID
-            days: Number of days to include in summary
-
-        Returns:
-            Dict with tenant usage summary
-        """
-        cutoff_date = datetime.now() - timedelta(days=days)
-
-        # Get all API keys for tenant
-        api_keys = db.query(ApiKey).filter(
-            ApiKey.tenant_id == tenant_id
-        ).all()
-
-        if not api_keys:
-            return {
-                "tenant_id": tenant_id,
-                "total_api_keys": 0,
-                "active_api_keys": 0,
-                "total_requests": 0,
-                "requests_last_30d": 0,
-                "most_active_key": None,
-                "key_breakdown": []
-            }
-
-        # Get usage stats for recent period
-        recent_usage = db.query(ApiKeyUsage).filter(
-            and_(
-                ApiKeyUsage.tenant_id == tenant_id,
-                ApiKeyUsage.last_used >= cutoff_date
-            )
-        ).all()
-
-        # Calculate totals
-        total_api_keys = len(api_keys)
-        active_api_keys = len([key for key in api_keys if key.active])
-        total_requests = sum(
-            sum(record.count for record in db.query(ApiKeyUsage).filter(
-                ApiKeyUsage.api_key_id == key.id
-            ).all())
-            for key in api_keys
-        )
-        requests_last_30d = sum(record.count for record in recent_usage)
-
-        # Find most active key
-        key_usage = {}
-        for key in api_keys:
-            key_total = sum(
-                record.count for record in db.query(ApiKeyUsage).filter(
-                    ApiKeyUsage.api_key_id == key.id
-                ).all()
-            )
-            key_usage[key.id] = {
-                "key": key,
-                "total_requests": key_total
-            }
-
-        most_active_key = None
-        if key_usage:
-            most_active_id = max(key_usage.keys(), key=lambda k: key_usage[k]["total_requests"])
-            most_active_key = {
-                "id": most_active_id,
-                "name": key_usage[most_active_id]["key"].name,
-                "requests": key_usage[most_active_id]["total_requests"]
-            }
-
-        # Create key breakdown
-        key_breakdown = [
-            {
-                "id": key_id,
-                "name": data["key"].name,
-                "prefix": data["key"].prefix,
-                "total_requests": data["total_requests"],
-                "active": data["key"].active,
-                "last_used": data["key"].last_used_at
-            }
-            for key_id, data in key_usage.items()
-        ]
-
-        return {
-            "tenant_id": tenant_id,
-            "total_api_keys": total_api_keys,
-            "active_api_keys": active_api_keys,
-            "total_requests": total_requests,
-            "requests_last_30d": requests_last_30d,
-            "most_active_key": most_active_key,
-            "key_breakdown": sorted(key_breakdown, key=lambda x: x["total_requests"], reverse=True)
-        }
-
-
-# Global service instance
+"""
+API Key Security Service
+Handles key rotation, IP whitelisting, and usage tracking
+"""
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Dict, List, Optional, Tuple, Any, Union
+from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import func, and_, desc, select
+
+from app.models import ApiKey, ApiKeyUsage, Tenant
+from app.config.api_keys import (
+    get_api_key_rotation_days,
+    get_rotation_warning_days,
+    is_rotation_enforced,
+    is_ip_whitelist_enabled,
+    is_usage_tracking_enabled
+)
+import uuid
+import ipaddress
+
+logger = logging.getLogger(__name__)
+
+
+class ApiKeySecurityService:
+    """
+    Service for managing API key security features
+
+    Features:
+    - Key rotation policy enforcement
+    - IP address whitelisting
+    - Usage tracking and analytics
+    - Security status monitoring
+    """
+
+    def __init__(self):
+        self.logger = logger
+
+    def check_rotation_status(self, api_key: ApiKey) -> Dict[str, Any]:
+        """
+        Check the rotation status of an API key
+
+        Args:
+            api_key: ApiKey instance
+
+        Returns:
+            Dict with rotation status information
+        """
+        if not api_key.last_rotated_at:
+            return {
+                "status": "expired",
+                "days_since_rotation": None,
+                "days_until_expiry": 0,
+                "message": "Key has never been rotated"
+            }
+
+        rotation_days = get_api_key_rotation_days()
+        warning_days = get_rotation_warning_days()
+
+        now = datetime.now(timezone.utc)
+        days_since_rotation = (now - api_key.last_rotated_at).days
+        days_until_expiry = rotation_days - days_since_rotation
+
+        if days_until_expiry <= 0:
+            status = "expired"
+            message = f"Key expired {abs(days_until_expiry)} days ago"
+        elif days_until_expiry <= warning_days:
+            status = "warning"
+            message = f"Key expires in {days_until_expiry} days"
+        else:
+            status = "ok"
+            message = f"Key is valid for {days_until_expiry} days"
+
+        return {
+            "status": status,
+            "days_since_rotation": days_since_rotation,
+            "days_until_expiry": days_until_expiry,
+            "message": message,
+            "rotation_required": days_until_expiry <= 0 and is_rotation_enforced()
+        }
+
+    def is_key_rotation_required(self, api_key: ApiKey) -> bool:
+        """
+        Check if key rotation is required
+
+        Args:
+            api_key: ApiKey instance
+
+        Returns:
+            bool: True if rotation is required
+        """
+        if not is_rotation_enforced():
+            return False
+
+        status = self.check_rotation_status(api_key)
+        return status["days_until_expiry"] <= 0
+
+    def validate_ip_whitelist(self, ip_list: List[str]) -> Tuple[bool, List[str]]:
+        """
+        Validate a list of IP addresses for whitelisting
+
+        Args:
+            ip_list: List of IP addresses/CIDR ranges
+
+        Returns:
+            Tuple of (is_valid, error_messages)
+        """
+        if not ip_list:
+            return True, []
+
+        errors = []
+
+        for ip in ip_list:
+            try:
+                # Try to parse as IP address or network
+                ipaddress.ip_network(ip, strict=False)
+            except ValueError as e:
+                errors.append(f"Invalid IP/CIDR '{ip}': {str(e)}")
+
+        # Check for duplicates
+        if len(ip_list) != len(set(ip_list)):
+            errors.append("Duplicate IP addresses found")
+
+        # Check maximum allowed IPs
+        from app.config.api_keys import get_api_key_config
+        config = get_api_key_config()
+        if len(ip_list) > config["max_whitelist_ips"]:
+            errors.append(f"Too many IP addresses (max: {config['max_whitelist_ips']})")
+
+        return len(errors) == 0, errors
+
+    def is_ip_allowed(self, client_ip: str, api_key: ApiKey) -> bool:
+        """
+        Check if client IP is allowed for the API key
+
+        Args:
+            client_ip: Client IP address
+            api_key: ApiKey instance
+
+        Returns:
+            bool: True if IP is allowed
+        """
+        if not is_ip_whitelist_enabled():
+            return True
+
+        if not api_key.ip_whitelist or len(api_key.ip_whitelist) == 0:
+            # No whitelist means all IPs allowed
+            return True
+
+        try:
+            client_addr = ipaddress.ip_address(client_ip)
+
+            for allowed_ip in api_key.ip_whitelist:
+                try:
+                    allowed_network = ipaddress.ip_network(allowed_ip, strict=False)
+                    if client_addr in allowed_network:
+                        return True
+                except ValueError:
+                    self.logger.warning(f"Invalid IP in whitelist: {allowed_ip}")
+                    continue
+
+            return False
+
+        except ValueError:
+            self.logger.warning(f"Invalid client IP: {client_ip}")
+            return False
+
+    async def track_usage(

+        self,

+        db: Union[Session, AsyncSession],

+        api_key_id: str,

+        tenant_id: str,

+        endpoint: str

+    ) -> None:

+        """

+        Track API key usage for an endpoint

+

+        Args:

+            db: Database session (sync or async)

+            api_key_id: API key ID

+            tenant_id: Tenant ID

+            endpoint: API endpoint called

+        """

+        if not is_usage_tracking_enabled():

+            return

+

+        try:

+            if isinstance(db, AsyncSession):

+                usage_result = await db.execute(

+                    select(ApiKeyUsage).where(

+                        and_(

+                            ApiKeyUsage.api_key_id == api_key_id,

+                            ApiKeyUsage.endpoint == endpoint

+                        )

+                    )

+                )

+                usage = usage_result.scalar_one_or_none()

+

+                if usage:

+                    usage.count += 1

+                    usage.last_used = datetime.now(timezone.utc)

+                else:

+                    usage = ApiKeyUsage(

+                        id=str(uuid.uuid4()),

+                        api_key_id=api_key_id,

+                        tenant_id=tenant_id,

+                        endpoint=endpoint,

+                        count=1,

+                        last_used=datetime.now(timezone.utc)

+                    )

+                    db.add(usage)

+

+                await db.flush()

+            else:

+                usage = db.query(ApiKeyUsage).filter(

+                    and_(

+                        ApiKeyUsage.api_key_id == api_key_id,

+                        ApiKeyUsage.endpoint == endpoint

+                    )

+                ).first()

+

+                if usage:

+                    usage.count += 1

+                    usage.last_used = datetime.now(timezone.utc)

+                else:

+                    usage = ApiKeyUsage(

+                        id=str(uuid.uuid4()),

+                        api_key_id=api_key_id,

+                        tenant_id=tenant_id,

+                        endpoint=endpoint,

+                        count=1,

+                        last_used=datetime.now(timezone.utc)

+                    )

+                    db.add(usage)

+

+                db.commit()

+        except Exception as e:

+            self.logger.error(f"Failed to track API key usage: {str(e)}")

+            if isinstance(db, AsyncSession):

+                await db.rollback()

+            else:

+                db.rollback()

+

+

+    def get_usage_stats(self, db: Session, api_key_id: str) -> Dict[str, Any]:
+        """
+        Get usage statistics for an API key
+
+        Args:
+            db: Database session
+            api_key_id: API key ID
+
+        Returns:
+            Dict with usage statistics
+        """
+        # Get all usage records for this API key
+        usage_records = db.query(ApiKeyUsage).filter(
+            ApiKeyUsage.api_key_id == api_key_id
+        ).all()
+
+        if not usage_records:
+            return {
+                "api_key_id": api_key_id,
+                "total_requests": 0,
+                "unique_endpoints": 0,
+                "last_used": None,
+                "most_used_endpoint": None,
+                "endpoint_breakdown": []
+            }
+
+        total_requests = sum(record.count for record in usage_records)
+        unique_endpoints = len(usage_records)
+        last_used = max(record.last_used for record in usage_records)
+
+        # Sort by usage count
+        sorted_records = sorted(usage_records, key=lambda r: r.count, reverse=True)
+        most_used_endpoint = sorted_records[0].endpoint if sorted_records else None
+
+        # Create endpoint breakdown
+        endpoint_breakdown = []
+        for record in sorted_records:
+            percentage = (record.count / total_requests * 100) if total_requests > 0 else 0
+            endpoint_breakdown.append({
+                "endpoint": record.endpoint,
+                "count": record.count,
+                "percentage": round(percentage, 2),
+                "last_used": record.last_used
+            })
+
+        return {
+            "api_key_id": api_key_id,
+            "total_requests": total_requests,
+            "unique_endpoints": unique_endpoints,
+            "last_used": last_used,
+            "most_used_endpoint": most_used_endpoint,
+            "endpoint_breakdown": endpoint_breakdown
+        }
+
+    def get_tenant_usage_summary(
+        self,
+        db: Session,
+        tenant_id: str,
+        days: int = 30
+    ) -> Dict[str, Any]:
+        """
+        Get usage summary for all API keys of a tenant
+
+        Args:
+            db: Database session
+            tenant_id: Tenant ID
+            days: Number of days to include in summary
+
+        Returns:
+            Dict with tenant usage summary
+        """
+        cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+
+        # Get all API keys for tenant
+        api_keys = db.query(ApiKey).filter(
+            ApiKey.tenant_id == tenant_id
+        ).all()
+
+        if not api_keys:
+            return {
+                "tenant_id": tenant_id,
+                "total_api_keys": 0,
+                "active_api_keys": 0,
+                "total_requests": 0,
+                "requests_last_30d": 0,
+                "most_active_key": None,
+                "key_breakdown": []
+            }
+
+        # Get usage stats for recent period
+        recent_usage = db.query(ApiKeyUsage).filter(
+            and_(
+                ApiKeyUsage.tenant_id == tenant_id,
+                ApiKeyUsage.last_used >= cutoff_date
+            )
+        ).all()
+
+        # Calculate totals
+        total_api_keys = len(api_keys)
+        active_api_keys = len([key for key in api_keys if key.active])
+        total_requests = sum(
+            sum(record.count for record in db.query(ApiKeyUsage).filter(
+                ApiKeyUsage.api_key_id == key.id
+            ).all())
+            for key in api_keys
+        )
+        requests_last_30d = sum(record.count for record in recent_usage)
+
+        # Find most active key
+        key_usage = {}
+        for key in api_keys:
+            key_total = sum(
+                record.count for record in db.query(ApiKeyUsage).filter(
+                    ApiKeyUsage.api_key_id == key.id
+                ).all()
+            )
+            key_usage[key.id] = {
+                "key": key,
+                "total_requests": key_total
+            }
+
+        most_active_key = None
+        if key_usage:
+            most_active_id = max(key_usage.keys(), key=lambda k: key_usage[k]["total_requests"])
+            most_active_key = {
+                "id": most_active_id,
+                "name": key_usage[most_active_id]["key"].name,
+                "requests": key_usage[most_active_id]["total_requests"]
+            }
+
+        # Create key breakdown
+        key_breakdown = [
+            {
+                "id": key_id,
+                "name": data["key"].name,
+                "prefix": data["key"].prefix,
+                "total_requests": data["total_requests"],
+                "active": data["key"].active,
+                "last_used": data["key"].last_used_at
+            }
+            for key_id, data in key_usage.items()
+        ]
+
+        return {
+            "tenant_id": tenant_id,
+            "total_api_keys": total_api_keys,
+            "active_api_keys": active_api_keys,
+            "total_requests": total_requests,
+            "requests_last_30d": requests_last_30d,
+            "most_active_key": most_active_key,
+            "key_breakdown": sorted(key_breakdown, key=lambda x: x["total_requests"], reverse=True)
+        }
+
+
+# Global service instance
 api_key_security_service = ApiKeySecurityService()
\ No newline at end of file