#!/usr/bin/env python3
"""
Script to use STS to assume a role and list objects in a bucket
Compatible with the RustFS project's STS implementation
"""

import boto3
import requests
import json
from botocore.exceptions import ClientError
import xml.etree.ElementTree as ET
from urllib.parse import urlencode


def main_with_boto3_sts():
    """
    Alternative implementation using boto3's built-in STS support
    """
    
    # Configuration
    STS_ENDPOINT = "http://127.0.0.1:9010"
    S3_ENDPOINT = "http://127.0.0.1:9010"
    BUCKET_NAME = "test-bucket"
    
    PERMANENT_ACCESS_KEY = "user-2"           # Replace with your access key
    PERMANENT_SECRET_KEY = "rustfsadmin"           # Replace with your secret key
    
    
    ROLE_ARN = "arn:aws:iam::123456789012:role/YourRole"
    SESSION_NAME = "python-sts-session"
    
    print("=== RustFS STS Client Example (boto3) ===")
    
    try:
        # Create STS client
        sts_client = boto3.client(
            'sts',
            aws_access_key_id=PERMANENT_ACCESS_KEY,
            aws_secret_access_key=PERMANENT_SECRET_KEY,
            endpoint_url=STS_ENDPOINT,
            region_name='us-east-1'
        )
        
        # Assume role
        print("1. Assuming role via STS...")
        sts_response = sts_client.assume_role(
            RoleArn=ROLE_ARN,
            RoleSessionName=SESSION_NAME,
            DurationSeconds=3600
        )
        
        # Extract temporary credentials
        temp_credentials = sts_response['Credentials']
        print("Successfully assumed role")
        print(f"sts_response: ",sts_response)
        print()
        
        # Create S3 client with temporary credentials
        print("2. Listing objects in bucket with temporary credentials...")
        s3_client = boto3.client(
            's3',
            aws_access_key_id=temp_credentials['AccessKeyId'],
            aws_secret_access_key=temp_credentials['SecretAccessKey'],
            aws_session_token=temp_credentials['SessionToken'],
            endpoint_url=S3_ENDPOINT,
            region_name='us-east-1'
        )
        
        # List objects
        response = s3_client.list_objects_v2(Bucket=BUCKET_NAME, MaxKeys=1000)
        
        objects = []
        if 'Contents' in response:
            objects = [obj['Key'] for obj in response['Contents']]
        
        print(f"Found {len(objects)} objects:")
        for obj in objects:
            print(f"  - {obj}")
            
        print("\n=== Script completed successfully ===")
        
    except ClientError as e:
        print(f"AWS Client Error: {e}")
    except Exception as e:
        print(f"Error: {str(e)}")



def assume_role_sts(sts_endpoint, access_key, secret_key, role_arn, role_session_name, duration_seconds=3600):
    """
    Assume a role using STS compatible with RustFS implementation
    
    Args:
        sts_endpoint (str): STS endpoint URL
        access_key (str): Access key for authentication
        secret_key (str): Secret key for authentication
        role_arn (str): Role ARN to assume
        role_session_name (str): Session name
        duration_seconds (int): Duration of credentials in seconds
        
    Returns:
        dict:   credentials
    """
    
    # Prepare STS request parameters
    params = {
        'Action': 'AssumeRole',
        'Version': '2011-06-15',
        'RoleArn': role_arn,
        'RoleSessionName': role_session_name,
        'DurationSeconds': str(duration_seconds)
    }
    
    # Create session for making requests
    session = requests.Session()
    
    # Prepare request
    url = f"{sts_endpoint}/"
    headers = {
        'Content-Type': 'application/x-www-form-urlencoded',
        'User-Agent': 'rustfs-python-sts-client/1.0'
    }
    
    # Sign the request using AWS signature v4
    # For simplicity in this example, we're sending credentials directly
    # In production, you should properly sign the request
    
    # Add authentication parameters
    auth_params = {
        'AWSAccessKeyId': access_key,
        # Note: Proper implementation would require signature calculation
        # This is a simplified example for demonstration
    }
    
    # Combine parameters
    all_params = {**params, **auth_params}
    data = urlencode(all_params)
    
    try:
        # Make the STS request
        print(url)
        response = session.post(url, headers=headers, data=data, verify=False)
        
        if response.status_code != 200:
            print(f"STS request failed with status {response.status_code}")
            print(response.text)
            return None
            
        # Parse XML response
        root = ET.fromstring(response.text)
        
        # Extract credentials from XML
        credentials_elem = root.find('.//{*}Credentials')
        if credentials_elem is None:
            print("Could not find credentials in STS response")
            print(response.text)
            return None
            
        temp_credentials = {}
        for child in credentials_elem:
            tag_name = child.tag.split('}')[1] if '}' in child.tag else child.tag
            temp_credentials[tag_name] = child.text
            
        return temp_credentials
        
    except Exception as e:
        print(f"Error assuming role: {str(e)}")
        return None

def list_bucket_objects(s3_endpoint, bucket_name, access_key, secret_key, session_token=None):
    """
    List objects in a bucket using temporary credentials
    
    Args:
        s3_endpoint (str): S3 endpoint URL
        bucket_name (str): Name of the bucket
        access_key (str): Access key
        secret_key (str): Secret key
        session_token (str): Session token for temporary credentials
        
    Returns:
        list: List of object keys
    """
    
    try:
        # Configure S3 client
        s3_config = {
            'aws_access_key_id': access_key,
            'aws_secret_access_key': secret_key,
            'endpoint_url': s3_endpoint,
            'region_name': 'us-east-1'  # Default region
        }
        
        # Add session token if provided
        if session_token:
            s3_config['aws_session_token'] = session_token
            
        # Create S3 client
        s3_client = boto3.client('s3', **s3_config)
        
        # List objects in the bucket
        response = s3_client.list_objects_v2(Bucket=bucket_name, MaxKeys=1000)
        
        # Extract object keys
        objects = []
        if 'Contents' in response:
            objects = [obj['Key'] for obj in response['Contents']]
            
        return objects
        
    except ClientError as e:
        print(f"AWS Client Error: {e}")
        return None
    except Exception as e:
        print(f"Error listing bucket objects: {str(e)}")
        return None

# Alternative implementation using boto3 for STS (if supported by your RustFS setup)
if __name__ == "__main__":
    main_with_boto3_sts()  # boto3-based implementation`