import boto3
import json
from datetime import datetime, timedelta

class JITAccessManager:
    def __init__(self):
        self.sts = boto3.client('sts')
        self.iam = boto3.client('iam')
        self.sns = boto3.client('sns')
        
    def request_elevated_access(self, requestor, role_arn, duration_hours, justification):
        """Request temporary elevated access with approval workflow"""
        
        # Create temporary policy for specific duration
        temp_policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Principal": {"AWS": f"arn:aws:iam::{self.get_account_id()}:user/{requestor}"},
                    "Action": "sts:AssumeRole",
                    "Resource": role_arn,
                    "Condition": {
                        "DateLessThan": {
                            "aws:CurrentTime": (datetime.utcnow() + timedelta(hours=duration_hours)).isoformat()
                        }
                    }
                }
            ]
        }
        
        # Send approval request
        approval_request = {
            "requestor": requestor,
            "role_arn": role_arn,
            "duration": duration_hours,
            "justification": justification,
            "timestamp": datetime.utcnow().isoformat(),
            "policy": temp_policy
        }
        
        self.sns.publish(
            TopicArn='arn:aws:sns:region:account:jit-access-requests',
            Message=json.dumps(approval_request),
            Subject=f'JIT Access Request: {requestor} -> {role_arn}'
        )
        
        return approval_request
    
    def approve_access_request(self, request_id, approver):
        """Approve and implement temporary access"""
        # Implementation for approval workflow
        pass

import boto3
import json

def create_security_monitoring():
    events = boto3.client('events')
    sns = boto3.client('sns')
    
    # Rule for suspicious role assumptions
    suspicious_pattern = {
        "source": ["aws.sts"],
        "detail-type": ["AWS API Call via CloudTrail"],
        "detail": {
            "eventSource": ["sts.amazonaws.com"],
            "eventName": ["AssumeRole"],
            "errorCode": {"exists": False},
            "sourceIPAddress": {
                "anything-but": {
                    "prefix": ["10.", "172.16.", "192.168."]
                }
            }
        }
    }
    
    events.put_rule(
        Name='SuspiciousRoleAssumption',
        EventPattern=json.dumps(suspicious_pattern),
        State='ENABLED',
        Description='Alert on role assumptions from external IPs'
    )
    
    # Rule for failed role assumptions
    failure_pattern = {
        "source": ["aws.sts"],
        "detail-type": ["AWS API Call via CloudTrail"],
        "detail": {
            "eventSource": ["sts.amazonaws.com"],
            "eventName": ["AssumeRole"],
            "errorCode": {"exists": True}
        }
    }
    
    events.put_rule(
        Name='FailedRoleAssumption',
        EventPattern=json.dumps(failure_pattern),
        State='ENABLED',
        Description='Alert on failed role assumption attempts'
    )

def lambda_security_handler(event, context):
    """Lambda function to process security events and send alerts"""
    
    detail = event['detail']
    event_name = detail.get('eventName')
    source_ip = detail.get('sourceIPAddress')
    user_identity = detail.get('userIdentity', {})
    
    # Analyze the event
    risk_score = calculate_risk_score(detail)
    
    if risk_score > 70:  # High risk threshold
        send_security_alert(detail, risk_score)
        
        # Optionally, automatically revoke the session
        if risk_score > 90:
            revoke_active_sessions(user_identity)
    
    return {'statusCode': 200}

def calculate_risk_score(detail):
    """Calculate risk score based on various factors"""
    score = 0
    
    # Check source IP reputation
    if is_suspicious_ip(detail.get('sourceIPAddress')):
        score += 30
    
    # Check time of access
    if is_unusual_time(detail.get('eventTime')):
        score += 20
    
    # Check role sensitivity
    role_arn = detail.get('requestParameters', {}).get('roleArn', '')
    if 'Admin' in role_arn or 'Root' in role_arn:
        score += 25
    
    # Check user behavior anomaly
    if is_anomalous_behavior(detail):
        score += 25
    
    return score

import boto3

# Create an STS client
sts_client = boto3.client('sts')

# Assume the role
assumed_role = sts_client.assume_role(
    RoleArn='arn:aws:iam::111122223333:role/EKSAdminRole',
    RoleSessionName='EKSAdminSession',
    DurationSeconds=3600  # 1 hour
)

# Extract the temporary credentials
credentials = assumed_role['Credentials']

# Create an EKS client using the temporary credentials
eks_client = boto3.client(
    'eks',
    region_name='us-west-2',
    aws_access_key_id=credentials['AccessKeyId'],
    aws_secret_access_key=credentials['SecretAccessKey'],
    aws_session_token=credentials['SessionToken']
)

# Now you can make EKS API calls
clusters = eks_client.list_clusters()
print(clusters)

def create_emergency_access_role():
    """Create break-glass role with temporary elevated permissions"""
    
    trust_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"AWS": f"arn:aws:iam::{ACCOUNT_ID}:role/EmergencyAccessRole"},
                "Action": "sts:AssumeRole",
                "Condition": {
                    "Bool": {"aws:MultiFactorAuthPresent": "true"},
                    "StringEquals": {"aws:PrincipalTag/EmergencyAccess": "approved"},
                    "NumericLessThan": {"aws:MultiFactorAuthAge": "900"}  # 15 minutes
                }
            }
        ]
    }
    
    # Create role with temporary permissions
    # Automatically expire after 2 hours

# DON'T DO THIS
AWS_ACCESS_KEY_ID = "AKIA..."
AWS_SECRET_ACCESS_KEY = "secret..."

# DO THIS INSTEAD
credentials = sts_client.assume_role(...)

import boto3
import json
import time
from datetime import datetime, timedelta
from botocore.exceptions import ClientError, BotoCoreError
import logging
from typing import Dict, Optional, Any
import threading
from dataclasses import dataclass

@dataclass
class AssumeRoleConfig:
    role_arn: str
    session_name: str
    duration_seconds: int = 3600
    external_id: Optional[str] = None
    session_policy: Optional[str] = None
    mfa_serial: Optional[str] = None
    region: str = 'us-east-1'

class EnhancedSTSManager:
    """Enterprise-grade STS session manager with caching and error handling"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self._sessions: Dict[str, Dict[str, Any]] = {}
        self._lock = threading.RLock()
        self.sts_client = boto3.client('sts')
        
    def assume_role(self, config: AssumeRoleConfig, mfa_token: Optional[str] = None) -> Dict[str, Any]:
        """Assume role with comprehensive error handling and caching"""
        
        cache_key = self._generate_cache_key(config)
        
        with self._lock:
            # Check cache first
            if self._is_cached_session_valid(cache_key):
                self.logger.info(f"Using cached credentials for {config.role_arn}")
                return self._sessions[cache_key]['credentials']
            
            try:
                credentials = self._perform_assume_role(config, mfa_token)
                self._cache_session(cache_key, credentials)
                return credentials
                
            except ClientError as e:
                self._handle_client_error(e, config)
                raise
            except Exception as e:
                self.logger.error(f"Unexpected error assuming role: {str(e)}")
                raise
    
    def _perform_assume_role(self, config: AssumeRoleConfig, mfa_token: Optional[str] = None) -> Dict[str, Any]:
        """Perform the actual assume role operation with retry logic"""
        
        assume_role_params = {
            'RoleArn': config.role_arn,
            'RoleSessionName': config.session_name,
            'DurationSeconds': config.duration_seconds
        }
        
        # Add optional parameters
        if config.external_id:
            assume_role_params['ExternalId'] = config.external_id
        
        if config.session_policy:
            assume_role_params['Policy'] = config.session_policy
        
        if config.mfa_serial and mfa_token:
            assume_role_params['SerialNumber'] = config.mfa_serial
            assume_role_params['TokenCode'] = mfa_token
        
        # Retry logic with exponential backoff
        max_retries = 3
        base_delay = 1
        
        for attempt in range(max_retries):
            try:
                response = self.sts_client.assume_role(**assume_role_params)
                
                # Enhance credentials with metadata
                credentials = response['Credentials']
                credentials['AssumedRoleUser'] = response['AssumedRoleUser']
                credentials['PackedPolicySize'] = response.get('PackedPolicySize', 0)
                
                self.logger.info(f"Successfully assumed role: {config.role_arn}")
                return credentials
                
            except ClientError as e:
                error_code = e.response['Error']['Code']
                
                if error_code in ['Throttling', 'RequestLimitExceeded'] and attempt < max_retries - 1:
                    delay = base_delay * (2 ** attempt)
                    self.logger.warning(f"Rate limited, retrying in {delay} seconds...")
                    time.sleep(delay)
                    continue
                else:
                    raise
    
    def _handle_client_error(self, error: ClientError, config: AssumeRoleConfig):
        """Provide detailed error analysis and suggestions"""
        
        error_code = error.response['Error']['Code']
        error_message = error.response['Error']['Message']
        
        troubleshooting_guide = {
            'AccessDenied': "Check trust policy, ensure calling identity has sts:AssumeRole permission",
            'InvalidParameterValue': "Verify role ARN format and account ID",
            'MalformedPolicyDocument': "Review session policy JSON syntax",
            'TokenRefreshRequired': "MFA token expired, obtain new token",
            'RegionDisabledException': "STS not available in requested region"
        }
        
        suggestion = troubleshooting_guide.get(error_code, "Check AWS documentation for error code")
        
        self.logger.error(f"AssumeRole failed: {error_code} - {error_message}")
        self.logger.error(f"Suggestion: {suggestion}")
        self.logger.error(f"Role ARN: {config.role_arn}")
    
    def _generate_cache_key(self, config: AssumeRoleConfig) -> str:
        """Generate unique cache key for session"""
        return f"{config.role_arn}:{config.session_name}:{config.duration_seconds}"
    
    def _is_cached_session_valid(self, cache_key: str) -> bool:
        """Check if cached session is still valid"""
        if cache_key not in self._sessions:
            return False
        
        session = self._sessions[cache_key]
        expiration = datetime.fromisoformat(session['credentials']['Expiration'].replace('Z', '+00:00'))
        
        # Consider session valid if more than 5 minutes remaining
        return datetime.now(expiration.tzinfo) < (expiration - timedelta(minutes=5))
    
    def _cache_session(self, cache_key: str, credentials: Dict[str, Any]):
        """Cache session credentials"""
        self._sessions[cache_key] = {
            'credentials': credentials,
            'cached_at': datetime.utcnow()
        }
    
    def get_session_info(self, cache_key: str) -> Optional[Dict[str, Any]]:
        """Get information about cached session"""
        with self._lock:
            if cache_key in self._sessions:
                session = self._sessions[cache_key]
                return {
                    'cached_at': session['cached_at'],
                    'expires_at': session['credentials']['Expiration'],
                    'assumed_role_arn': session['credentials']['AssumedRoleUser']['Arn']
                }
        return None
    
    def revoke_session(self, cache_key: str):
        """Manually revoke cached session"""
        with self._lock:
            if cache_key in self._sessions:
                del self._sessions[cache_key]
                self.logger.info(f"Revoked cached session: {cache_key}")

# Usage example
def main():
    manager = EnhancedSTSManager()
    
    config = AssumeRoleConfig(
        role_arn='arn:aws:iam::111122223333:role/EKSAdminRole',
        session_name=f'enhanced-session-{int(time.time())}',
        duration_seconds=3600,
        external_id='unique-external-id-2024'
    )
    
    try:
        credentials = manager.assume_role(config, mfa_token='123456')
        
        # Use credentials to create service clients
        eks_client = boto3.client(
            'eks',
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken'],
            region_name='us-west-2'
        )
        
        # Your EKS operations here
        clusters = eks_client.list_clusters()
        print(f"Found {len(clusters['clusters'])} EKS clusters")
        
    except Exception as e:
        logging.error(f"Failed to assume role: {str(e)}")

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()