| |
| """ |
| 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 |
|
|
| |
| logging.basicConfig( |
| level=logging.INFO, |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' |
| ) |
| logger = logging.getLogger("s3_test") |
|
|
| |
| DEFAULT_EXPIRATION = 3600 |
| 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 |
| |
| |
| cfg = BotoConfig(signature_version="s3v4", s3={"addressing_style": "virtual"}) |
| |
| |
| 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: |
| 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 |
| ) |
| |
| |
| 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)") |
| |
| |
| 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") |
| |
| |
| 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)") |
| |
| |
| 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() |
| |
| |
| config = {k: getattr(args, k) if getattr(args, k) is not None else env_vars.get(k) for k in env_vars} |
| |
| |
| 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: |
| 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"] |
| ) |
| |
| |
| 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: |
| |
| 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)}") |
| |
| |
| logger.info("\nStep 2: Uploading file using presigned URL...") |
| if not upload_to_presigned_url(input_file, upload_url): |
| logger.error("Upload failed. Exiting.") |
| |
| |
| |
| 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)}") |
| |
| |
| 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.") |
| |
| |
| |
| 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) |
| |
| |
| 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}") |
| |
| finally: |
| |
| 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() |
|
|
