search_personalization/ubi_proxy.py

"""
UBI Proxy Module

This module provides backend proxy functions for forwarding UBI data to OpenSearch Ingestion Service
with AWS SigV4 authentication. This approach keeps AWS credentials secure on the server side
rather than exposing them in client-side JavaScript.
"""

import os
import json
import logging
from typing import Dict, Any, Optional
from dotenv import load_dotenv
import boto3
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
import requests

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_aws_credentials() -> tuple[Optional[str], Optional[str], Optional[str]]:
    """
    Get AWS credentials from environment or boto3 session.
    
    Returns:
        Tuple of (access_key_id, secret_access_key, session_token)
        Returns (None, None, None) if credentials cannot be retrieved
    """
    try:
        # Try to get credentials from environment variables first
        access_key = os.getenv('AWS_ACCESS_KEY_ID')
        secret_key = os.getenv('AWS_SECRET_ACCESS_KEY')
        session_token = os.getenv('AWS_SESSION_TOKEN')
        
        if access_key and secret_key:
            return access_key, secret_key, session_token
        
        # Fall back to boto3 session (will use IAM role if available)
        aws_profile = os.getenv("AWS_PROFILE")
        session = boto3.Session(profile_name=aws_profile) if aws_profile else boto3.Session()
        credentials = session.get_credentials()
        
        if credentials:
            return (
                credentials.access_key,
                credentials.secret_key,
                credentials.token
            )
        
        return None, None, None
    
    except Exception as e:
        logger.error(f"Error retrieving AWS credentials: {e}")
        return None, None, None


def sign_request(
    method: str,
    url: str,
    region: str,
    service: str,
    payload: str,
    access_key: str,
    secret_key: str,
    session_token: Optional[str] = None
) -> Dict[str, str]:
    """
    Sign an HTTP request using AWS SigV4.
    
    Args:
        method: HTTP method (GET, POST, etc.)
        url: Full URL to sign
        region: AWS region
        service: AWS service name (e.g., 'osis' for OpenSearch Ingestion Service)
        payload: Request body as string
        access_key: AWS access key ID
        secret_key: AWS secret access key
        session_token: Optional AWS session token
    
    Returns:
        Dictionary of headers including Authorization header
    """
    # Create AWS request
    request = AWSRequest(
        method=method,
        url=url,
        data=payload,
        headers={
            'Content-Type': 'application/json',
            'Host': url.split('/')[2]  # Extract host from URL
        }
    )
    
    # Create SigV4 auth
    credentials = boto3.Session(
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        aws_session_token=session_token
    ).get_credentials()
    
    signer = SigV4Auth(credentials, service, region)
    
    # Sign the request
    signer.add_auth(request)
    
    # Return signed headers
    return dict(request.headers)


def send_ubi_query(query_data: Dict[str, Any]) -> tuple[bool, Optional[str]]:
    """
    Forward UBI query data to OpenSearch Ingestion Service with AWS SigV4 authentication.
    
    Args:
        query_data: UBI query data dictionary
    
    Returns:
        Tuple of (success, error_message)
        - success: True if data was sent successfully, False otherwise
        - error_message: Error message if failed, None if successful
    
    Requirements: 9.3, 9.4, 9.5, 9.6
    """
    try:
        # Get configuration
        queries_endpoint = os.getenv('UBI_QUERIES_ENDPOINT')
        aws_region = os.getenv('AWS_REGION', 'us-east-1')
        
        if not queries_endpoint:
            error_msg = "UBI_QUERIES_ENDPOINT not configured"
            logger.error(error_msg)
            return False, error_msg
        
        # Get AWS credentials
        access_key, secret_key, session_token = get_aws_credentials()
        
        if not access_key or not secret_key:
            error_msg = "AWS credentials not available"
            logger.error(error_msg)
            return False, error_msg
        
        # Wrap data in array format as required by ingestion service
        payload = json.dumps([query_data])
        
        # Construct full URL - OpenSearch Ingestion Service requires /ubi/queries path
        url = f"{queries_endpoint}/ubi/queries"
        
        # Sign the request
        headers = sign_request(
            method='POST',
            url=url,
            region=aws_region,
            service='osis',  # OpenSearch Ingestion Service
            payload=payload,
            access_key=access_key,
            secret_key=secret_key,
            session_token=session_token
        )
        
        # Send request
        logger.info(f"Sending UBI query to {url}")
        response = requests.post(
            url,
            data=payload,
            headers=headers,
            timeout=10
        )
        
        # Check response
        if response.status_code >= 200 and response.status_code < 300:
            logger.info(f"UBI query sent successfully: {response.status_code}")
            return True, None
        else:
            error_msg = f"UBI query failed: {response.status_code} {response.text}"
            logger.error(error_msg)
            return False, error_msg
    
    except requests.exceptions.Timeout:
        error_msg = "Request timeout while sending UBI query"
        logger.error(error_msg)
        return False, error_msg
    
    except requests.exceptions.ConnectionError as e:
        error_msg = f"Connection error while sending UBI query: {str(e)}"
        logger.error(error_msg)
        return False, error_msg
    
    except Exception as e:
        error_msg = f"Unexpected error sending UBI query: {str(e)}"
        logger.exception(error_msg)
        return False, error_msg


def send_ubi_event(event_data: Dict[str, Any]) -> tuple[bool, Optional[str]]:
    """
    Forward UBI event data to OpenSearch Ingestion Service with AWS SigV4 authentication.
    
    Args:
        event_data: UBI event data dictionary
    
    Returns:
        Tuple of (success, error_message)
        - success: True if data was sent successfully, False otherwise
        - error_message: Error message if failed, None if successful
    
    Requirements: 9.3, 9.4, 9.5, 9.6
    """
    try:
        # Get configuration
        events_endpoint = os.getenv('UBI_EVENTS_ENDPOINT')
        aws_region = os.getenv('AWS_REGION', 'us-east-1')
        
        if not events_endpoint:
            error_msg = "UBI_EVENTS_ENDPOINT not configured"
            logger.error(error_msg)
            return False, error_msg
        
        # Get AWS credentials
        access_key, secret_key, session_token = get_aws_credentials()
        
        if not access_key or not secret_key:
            error_msg = "AWS credentials not available"
            logger.error(error_msg)
            return False, error_msg
        
        # Wrap data in array format as required by ingestion service
        payload = json.dumps([event_data])
        
        # Construct full URL - OpenSearch Ingestion Service requires /ubi/events path
        url = f"{events_endpoint}/ubi/events"
        
        # Sign the request
        headers = sign_request(
            method='POST',
            url=url,
            region=aws_region,
            service='osis',  # OpenSearch Ingestion Service
            payload=payload,
            access_key=access_key,
            secret_key=secret_key,
            session_token=session_token
        )
        
        # Send request
        logger.info(f"Sending UBI event to {url}")
        response = requests.post(
            url,
            data=payload,
            headers=headers,
            timeout=10
        )
        
        # Check response
        if response.status_code >= 200 and response.status_code < 300:
            logger.info(f"UBI event sent successfully: {response.status_code}")
            return True, None
        else:
            error_msg = f"UBI event failed: {response.status_code} {response.text}"
            logger.error(error_msg)
            return False, error_msg
    
    except requests.exceptions.Timeout:
        error_msg = "Request timeout while sending UBI event"
        logger.error(error_msg)
        return False, error_msg
    
    except requests.exceptions.ConnectionError as e:
        error_msg = f"Connection error while sending UBI event: {str(e)}"
        logger.error(error_msg)
        return False, error_msg
    
    except Exception as e:
        error_msg = f"Unexpected error sending UBI event: {str(e)}"
        logger.exception(error_msg)
        return False, error_msg