Blame for scripts/test_s3.py

#!/usr/bin/env python3
"""
S3 Standalone Test Script for Attestation Store

This script tests S3 operations with presigned URLs:
1. Generates a presigned URL for uploading a file
2. Uploads a test file to the S3 bucket using the presigned URL
3. Generates a presigned URL for downloading the file
4. Downloads the file using the presigned URL
5. Compares the original and downloaded files

The script supports both AWS S3 and MinIO storage backends.
"""

import os
import sys
import uuid
import json
import boto3
import hashlib
import logging
import requests
import tempfile
import argparse
from datetime import timedelta
from botocore.config import Config as BotoConfig
from urllib.parse import urlparse

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("s3_test")

# Default settings
DEFAULT_EXPIRATION = 3600  # 1 hour
DEFAULT_REGION = "il-central-1"
TEST_CONTENT = json.dumps({"test": "data", "timestamp": str(uuid.uuid4())})


class S3StorageBackend:
    """
    AWS S3 storage backend (compatible with the actual implementation in the project)
    """
    
    def __init__(self, access_key, secret_key, region, bucket, session_token=None):
        self.bucket = bucket
        
        # Configure boto3 client
        cfg = BotoConfig(signature_version="s3v4", s3={"addressing_style": "virtual"})
        
        # If keys are empty, allow the default provider chain (env vars, IAM role, etc.)
        client_kwargs = dict(service_name="s3", region_name=region, config=cfg)
        if access_key and secret_key:
            client_kwargs.update(
                aws_access_key_id=access_key,
                aws_secret_access_key=secret_key,
            )
            if session_token:  # required for ASIA... keys
                client_kwargs.update(aws_session_token=session_token)
        
        self.s3_client = boto3.client(**client_kwargs)
        logger.info(f"Initialized S3 client for region {region}, bucket {bucket}")
    
    def _normalize_content_type(self, content_type):
        """Normalize content type"""
        if not content_type:
            return "application/json"
        
        if content_type == "cyclonedx-json":
            return "application/cyclonedx+json"
        
        return "application/json"
    
    def generate_upload_url(self, file_name, content_type=None, session_id=None, expiration=DEFAULT_EXPIRATION):
        """Generate presigned URL for S3 upload"""
        file_id = session_id or str(uuid.uuid4())
        storage_key = f"attestations/{file_id}/{file_name}"
        ct = self._normalize_content_type(content_type)
        
        params = {
            "Bucket": self.bucket,
            "Key": storage_key,
            "ContentType": ct,
        }
        
        url = self.s3_client.generate_presigned_url(
            "put_object",
            Params=params,
            ExpiresIn=expiration,
            HttpMethod="PUT",
        )
        
        logger.info(f"S3 presigned PUT: {url}")
        logger.info(f"S3 region: {self.s3_client.meta.region_name}, bucket: {self.bucket}, key: {storage_key}")
        return {"upload_url": url, "storage_path": storage_key}
    
    def generate_download_url(self, storage_path, expiration=DEFAULT_EXPIRATION):
        """Generate presigned URL for S3 download"""
        url = self.s3_client.generate_presigned_url(
            "get_object",
            Params={"Bucket": self.bucket, "Key": storage_path},
            ExpiresIn=expiration,
        )
        logger.info(f"S3 presigned GET: {url}")
        return url
    
    def delete_file(self, storage_path):
        """Delete S3 object"""
        try:
            self.s3_client.delete_object(Bucket=self.bucket, Key=storage_path)
            logger.info(f"Successfully deleted S3 object: {storage_path}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete S3 object {storage_path}: {e}")
            return False


class MinIOStorageBackend:
    """
    MinIO storage backend (compatible with the actual implementation in the project)
    """
    
    def __init__(self, endpoint, access_key, secret_key, bucket, secure=True):
        try:
            from minio import Minio
            
            self.bucket = bucket
            self.client = Minio(
                endpoint,
                access_key=access_key,
                secret_key=secret_key,
                secure=secure
            )
            
            # Create bucket if it doesn't exist
            if not self.client.bucket_exists(bucket):
                self.client.make_bucket(bucket)
                logger.info(f"Created MinIO bucket: {bucket}")
            else:
                logger.info(f"Using existing MinIO bucket: {bucket}")
                
        except ImportError:
            logger.error("MinIO Python client not found. Please install with: pip install minio")
            sys.exit(1)
        except Exception as e:
            logger.error(f"Failed to initialize MinIO client: {e}")
            sys.exit(1)
    
    def generate_upload_url(self, file_name, content_type=None, session_id=None, expiration=DEFAULT_EXPIRATION):
        """Generate presigned URL for MinIO upload"""
        file_id = session_id or str(uuid.uuid4())
        storage_key = f"attestations/{file_id}/{file_name}"
        
        storage_key = storage_key.replace('//', '/')
        
        upload_url = self.client.get_presigned_url(
            "PUT",
            self.bucket,
            storage_key,
            expires=timedelta(seconds=expiration)
        )
        
        logger.info(f"MinIO presigned PUT: {upload_url}")
        return {
            'upload_url': upload_url,
            'storage_path': storage_key
        }
    
    def generate_download_url(self, storage_path, expiration=DEFAULT_EXPIRATION):
        """Generate presigned URL for MinIO download"""
        url = self.client.get_presigned_url(
            "GET",
            self.bucket,
            storage_path,
            expires=timedelta(seconds=expiration)
        )
        logger.info(f"MinIO presigned GET: {url}")
        return url
    
    def delete_file(self, storage_path):
        """Delete MinIO object"""
        try:
            self.client.remove_object(self.bucket, storage_path)
            logger.info(f"Successfully deleted MinIO object: {storage_path}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete MinIO object {storage_path}: {e}")
            return False


def create_test_file(content=None):
    """Create a temporary test file"""
    if content is None:
        content = TEST_CONTENT
    
    fd, path = tempfile.mkstemp(suffix=".json")
    try:
        with os.fdopen(fd, 'w') as f:
            f.write(content)
        return path
    except Exception as e:
        logger.error(f"Failed to create test file: {e}")
        os.unlink(path)
        sys.exit(1)


def calculate_file_hash(file_path):
    """Calculate SHA-256 hash of a file"""
    sha256 = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


def upload_to_presigned_url(file_path, upload_url, content_type="application/json"):
    """Upload a file using a presigned URL"""
    try:
        with open(file_path, "rb") as f:
            headers = {"Content-Type": content_type}
            response = requests.put(upload_url, data=f, headers=headers)
            response.raise_for_status()
            logger.info(f"Successfully uploaded file to {upload_url}, status: {response.status_code}")
            return True
    except Exception as e:
        logger.error(f"Failed to upload file using presigned URL: {e}")
        return False


def download_from_presigned_url(download_url, output_path):
    """Download a file using a presigned URL"""
    try:
        response = requests.get(download_url)
        response.raise_for_status()
        with open(output_path, "wb") as f:
            f.write(response.content)
        logger.info(f"Successfully downloaded file from {download_url} to {output_path}")
        return True
    except Exception as e:
        logger.error(f"Failed to download file using presigned URL: {e}")
        return False


def verify_files_match(original_path, downloaded_path):
    """Verify that the original and downloaded files match"""
    original_hash = calculate_file_hash(original_path)
    downloaded_hash = calculate_file_hash(downloaded_path)
    
    if original_hash == downloaded_hash:
        logger.info(f"Files match! Hash: {original_hash}")
        return True
    else:
        logger.error(f"Files do not match! Original hash: {original_hash}, Downloaded hash: {downloaded_hash}")
        return False


def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description="Test S3 operations with presigned URLs")
    parser.add_argument("--provider", choices=["aws", "minio"], default="aws",
                        help="Storage provider (aws or minio)")
    
    # AWS S3 options
    parser.add_argument("--aws-access-key", help="AWS Access Key ID")
    parser.add_argument("--aws-secret-key", help="AWS Secret Access Key")
    parser.add_argument("--aws-region", default=DEFAULT_REGION, help="AWS Region")
    parser.add_argument("--s3-bucket", help="S3 Bucket name")
    
    # MinIO options
    parser.add_argument("--minio-endpoint", help="MinIO endpoint (e.g., 'minio.example.com:9000')")
    parser.add_argument("--minio-access-key", help="MinIO Access Key")
    parser.add_argument("--minio-secret-key", help="MinIO Secret Key")
    parser.add_argument("--minio-bucket", help="MinIO Bucket name")
    parser.add_argument("--minio-secure", default=True, type=bool, help="Use secure connection (HTTPS)")
    
    # Common options
    parser.add_argument("--file", help="Path to file to upload (creates a test file if not provided)")
    parser.add_argument("--expiration", type=int, default=DEFAULT_EXPIRATION,
                        help="Presigned URL expiration time in seconds")
    parser.add_argument("--no-cleanup", action="store_true", help="Do not delete the uploaded file")
    
    return parser.parse_args()


def read_env_vars():
    """Read S3 configuration from environment variables"""
    env_vars = {
        "aws_access_key": os.environ.get("AWS_ACCESS_KEY_ID"),
        "aws_secret_key": os.environ.get("AWS_SECRET_ACCESS_KEY"),
        "aws_region": os.environ.get("AWS_DEFAULT_REGION", DEFAULT_REGION),
        "s3_bucket": os.environ.get("S3_BUCKET"),
        "minio_endpoint": os.environ.get("MINIO_ENDPOINT"),
        "minio_access_key": os.environ.get("MINIO_ACCESS_KEY"),
        "minio_secret_key": os.environ.get("MINIO_SECRET_KEY"),
        "minio_bucket": os.environ.get("MINIO_BUCKET"),
        "minio_secure": os.environ.get("MINIO_SECURE", "true").lower() == "true",
    }
    return env_vars


def extract_hostname_from_url(url):
    """Extract hostname from URL for logging purposes"""
    parsed = urlparse(url)
    return parsed.netloc


def main():
    """Main function"""
    args = parse_args()
    env_vars = read_env_vars()
    
    # Combine command line args with env vars (cmd line takes precedence)
    config = {k: getattr(args, k) if getattr(args, k) is not None else env_vars.get(k) for k in env_vars}
    
    # Create storage backend
    if args.provider == "aws":
        required = ["aws_access_key", "aws_secret_key", "aws_region", "s3_bucket"]
        if any(config.get(k) is None for k in required):
            logger.error("Missing required AWS S3 configuration. Please provide AWS credentials and bucket name.")
            sys.exit(1)
            
        storage = S3StorageBackend(
            access_key=config["aws_access_key"],
            secret_key=config["aws_secret_key"],
            region=config["aws_region"],
            bucket=config["s3_bucket"]
        )
    else:  # minio
        required = ["minio_endpoint", "minio_access_key", "minio_secret_key", "minio_bucket"]
        if any(config.get(k) is None for k in required):
            logger.error("Missing required MinIO configuration. Please provide MinIO credentials and bucket name.")
            sys.exit(1)
            
        storage = MinIOStorageBackend(
            endpoint=config["minio_endpoint"],
            access_key=config["minio_access_key"],
            secret_key=config["minio_secret_key"],
            bucket=config["minio_bucket"],
            secure=config["minio_secure"]
        )
    
    # Create a test file if not provided
    input_file = args.file
    temp_file_created = False
    if not input_file:
        input_file = create_test_file()
        temp_file_created = True
    
    logger.info(f"Using test file: {input_file}")
    file_name = os.path.basename(input_file)
    
    try:
        # Step 1: Generate presigned upload URL
        logger.info("Step 1: Generating presigned upload URL...")
        upload_result = storage.generate_upload_url(
            file_name=file_name,
            content_type="application/json",
            expiration=args.expiration
        )
        upload_url = upload_result["upload_url"]
        storage_path = upload_result["storage_path"]
        
        logger.info(f"Storage path: {storage_path}")
        logger.info(f"Upload URL hostname: {extract_hostname_from_url(upload_url)}")
        
        # Step 2: Upload file using presigned URL
        logger.info("\nStep 2: Uploading file using presigned URL...")
        if not upload_to_presigned_url(input_file, upload_url):
            logger.error("Upload failed. Exiting.")
            # sys.exit(1)
        
        # Step 3: Generate presigned download URL
        logger.info("\nStep 3: Generating presigned download URL...")
        download_url = storage.generate_download_url(storage_path, args.expiration)
        logger.info(f"Download URL hostname: {extract_hostname_from_url(download_url)}")
        
        # Step 4: Download file using presigned URL
        logger.info("\nStep 4: Downloading file using presigned URL...")
        download_file = tempfile.mktemp(suffix=".json")
        if not download_from_presigned_url(download_url, download_file):
            logger.error("Download failed. Exiting.")
            # sys.exit(1)
        
        # Step 5: Verify files match
        logger.info("\nStep 5: Verifying file integrity...")
        if verify_files_match(input_file, download_file):
            logger.info("✓ TEST PASSED: Files match perfectly!")
        else:
            logger.error("✗ TEST FAILED: Files do not match!")
            sys.exit(1)
            
        # Step 6: Clean up (delete the uploaded file)
        if not args.no_cleanup:
            logger.info("\nStep 6: Cleaning up...")
            if storage.delete_file(storage_path):
                logger.info("Successfully deleted the uploaded file.")
            else:
                logger.warning("Failed to delete the uploaded file.")
    
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        # sys.exit(1)
    finally:
        # Clean up temporary files
        if temp_file_created and os.path.exists(input_file):
            os.unlink(input_file)
            logger.debug(f"Deleted temporary input file: {input_file}")
        
        if 'download_file' in locals() and os.path.exists(download_file):
            os.unlink(download_file)
            logger.debug(f"Deleted temporary download file: {download_file}")


if __name__ == "__main__":
    main()