Spaces:

dn-scribe
/

attstore

Runtime error

attstore / storage.py

fix: get things working with s3 (ux and valint)

8f47551 12 months ago

11.6 kB

import os
import uuid
import boto3
from botocore.config import Config as BotoConfig
from minio import Minio
from urllib.parse import urljoin
from flask import current_app
from logging_utils import get_logger
from datetime import timedelta


logger = get_logger(__name__)

class StorageBackend:
    """Base storage backend interface"""
    
    def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
        """Generate a presigned URL for upload"""
        raise NotImplementedError
    
    def generate_download_url(self, storage_path):
        """Generate a download URL"""
        raise NotImplementedError
    
    def delete_file(self, storage_path):
        """Delete a file"""
        raise NotImplementedError

class FileStorageBackend(StorageBackend):
    """Local file storage backend"""
    
    def __init__(self, storage_path, base_url):
        self.storage_path = storage_path
        self.base_url = base_url
        os.makedirs(storage_path, exist_ok=True)
    
    def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
        """For file storage, return a local path and upload endpoint
        
        Uses request_base_url if provided (for proxy/k8s scenarios), otherwise falls back to
        configured base_url or relative URL.
        """
        # Sanitize file_name to prevent directory traversal
        # Remove any path separators and normalize
        safe_file_name = os.path.basename(file_name.replace('\\', '/').split('/')[-1])
        if not safe_file_name or safe_file_name.startswith('.'):
            safe_file_name = 'file.bin'
        
        file_id = session_id or str(uuid.uuid4())
        storage_key = f"{file_id}/{safe_file_name}"
        local_path = os.path.join(self.storage_path, storage_key)
        
        # Verify the resolved path is within storage_path (prevent traversal)
        real_storage_path = os.path.realpath(self.storage_path)
        real_local_path = os.path.realpath(local_path)
        if not real_local_path.startswith(real_storage_path):
            raise ValueError(f"Invalid file path: directory traversal detected")

        # Create directory
        os.makedirs(os.path.dirname(local_path), exist_ok=True)

        # Build upload URL - prefer request_base_url (works with proxies/k8s)
        relative_upload_path = f"/evidence/upload/{file_id}"
        if request_base_url:
            # Use the actual request base URL (includes correct protocol/host/port)
            # Remove trailing slash from base_url if present
            base = request_base_url.rstrip('/')
            upload_url = f"{base}{relative_upload_path}"
        else:
            # Fallback to relative URL for same-origin requests
            upload_url = relative_upload_path

        return {
            'upload_url': upload_url,
            'storage_path': storage_key,
            'local_path': local_path
        }
    
    def generate_download_url(self, storage_path):
        """Generate download URL for local files"""
        # For API access, prepend 'evidence/download/' to the storage path
        api_path = f"evidence/download/{storage_path}"
        return urljoin(self.base_url, api_path)
    
    def delete_file(self, storage_path):
        """Delete local file"""
        local_path = os.path.join(self.storage_path, storage_path)
        
        # SECURITY: Verify the resolved path is within storage directory
        real_storage_path = os.path.realpath(self.storage_path)
        real_local_path = os.path.realpath(local_path)
        if not real_local_path.startswith(real_storage_path):
            logger.error(f"Directory traversal attempt detected in delete: {storage_path}")
            return False
        
        try:
            if os.path.exists(local_path):
                os.remove(local_path)
                # Also remove directory if empty
                dir_path = os.path.dirname(local_path)
                if os.path.exists(dir_path) and not os.listdir(dir_path):
                    os.rmdir(dir_path)
            return True
        except Exception as e:
            logger.error(f"Failed to delete file {storage_path}: {e}")
            return False

class S3StorageBackend(StorageBackend):
    """
    Works with real AWS S3.
    - Uses SigV4.
    - Supports temporary creds (session token).
    - Avoids forcing endpoint_url (let boto choose correct regional endpoint).
    - Uses virtual-hosted–style addressing if bucket name is DNS-safe.
    """

    def __init__(self, access_key, secret_key, region, bucket, session_token=None):
        self.bucket = bucket

        # IMPORTANT: do NOT force endpoint_url for AWS; let boto3 pick the regional S3 endpoint.
        # Only set endpoint_url when talking to S3-compatible services like MinIO.
        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)


    def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
        file_id = session_id or str(uuid.uuid4())
        storage_key = f"attestations/{file_id}/{file_name}"
        
        # Normalize storage path (replace double slashes with single)
        storage_key = storage_key.replace('//', '/')
        
        # For S3 compatibility with valint (CLI client that always uses multipart/form-data)
        # Use multipart/form-data as content type to match what valint sends
        # This is a raw upload with the multipart/form-data header, not actual multipart encoding
        base_ct = "multipart/form-data"
        
        logger.info(f"S3 Upload URL generation - Original CT: '{content_type}', Using: '{base_ct}' for S3 compatibility, File: {file_name}")

        params = {
            "Bucket": self.bucket,
            "Key": storage_key,
            "ContentType": base_ct,
        }

        url = self.s3_client.generate_presigned_url(
            "put_object",
            Params=params,
            ExpiresIn=current_app.config["STORAGE_CONFIG"].PRESIGNED_URL_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):
        try:
            # Extract filename from storage path for Content-Disposition header
            filename = storage_path.split('/')[-1] if '/' in storage_path else storage_path
            
            url = self.s3_client.generate_presigned_url(
                "get_object",
                Params={
                    "Bucket": self.bucket,
                    "Key": storage_path,
                    "ResponseContentDisposition": f'attachment; filename="{filename}"'
                },
                ExpiresIn=current_app.config["STORAGE_CONFIG"].PRESIGNED_URL_EXPIRATION,
            )
            logger.info(f"S3 presigned GET: {url}")
            return url
        except Exception as e:
            logger.error(f"Failed to generate S3 download URL for {storage_path}: {e}")
            return None

    def delete_file(self, storage_path):
        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(StorageBackend):
    """MinIO storage backend"""
    
    def __init__(self, endpoint, access_key, secret_key, bucket, secure=True):
        self.bucket = bucket
        
        # Strip http:// or https:// from endpoint if present
        if endpoint.startswith("http://"):
            endpoint = endpoint[7:]
            secure = False
        elif endpoint.startswith("https://"):
            endpoint = endpoint[8:]
            secure = True
            
        self.client = Minio(
            endpoint,
            access_key=access_key,
            secret_key=secret_key,
            secure=secure
        )
        
        # Create bucket if it doesn't exist
        try:
            if not self.client.bucket_exists(bucket):
                self.client.make_bucket(bucket)
        except Exception as e:
            logger.error(f"Failed to create/check bucket {bucket}: {e}")
    
    def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
        """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=current_app.config['STORAGE_CONFIG'].PRESIGNED_URL_EXPIRATION)
        )

        
        return {
            'upload_url': upload_url,
            'storage_path': storage_key
        }
    
    def generate_download_url(self, storage_path):
        """Generate presigned URL for MinIO download"""
        return self.client.get_presigned_url(
            "GET",
            self.bucket,
            storage_path,
            expires=timedelta(seconds=current_app.config['STORAGE_CONFIG'].PRESIGNED_URL_EXPIRATION)
        )
    
    def delete_file(self, storage_path):
        """Delete MinIO object"""
        try:
            self.client.remove_object(self.bucket, storage_path)
            return True
        except Exception as e:
            logger.error(f"Failed to delete MinIO object {storage_path}: {e}")
            return False

def get_storage_backend():
    """Get configured storage backend"""
    config = current_app.config['STORAGE_CONFIG']
    
    if config.STORAGE_TYPE == 'FILE_MOUNT':
        return FileStorageBackend(config.FILE_STORAGE_PATH, config.FILE_STORAGE_BASE_URL)
    
    elif config.STORAGE_TYPE == 'CLOUD_STORAGE':
        # Check for MinIO first
        if config.CLOUD_STORAGE_PROVIDER == 'MINIO':
            return MinIOStorageBackend(
                config.MINIO_ENDPOINT,
                config.MINIO_ACCESS_KEY,
                config.MINIO_SECRET_KEY,
                config.MINIO_BUCKET,
                config.MINIO_SECURE
            )
        # Otherwise use AWS S3
        elif config.CLOUD_STORAGE_PROVIDER == 'AWS':
            return S3StorageBackend(
                config.AWS_ACCESS_KEY_ID,
                config.AWS_SECRET_ACCESS_KEY,
                config.AWS_REGION,
                config.S3_BUCKET,
                session_token=config.AWS_SESSION_TOKEN if hasattr(config, 'AWS_SESSION_TOKEN') else None
            )
        else:
            raise ValueError("Cloud storage configured but no valid backend settings found")
    
    else:
        raise ValueError(f"Unknown storage type: {config.STORAGE_TYPE}")

	import os
	import uuid
	import boto3
	from botocore.config import Config as BotoConfig
	from minio import Minio
	from urllib.parse import urljoin
	from flask import current_app
	from logging_utils import get_logger
	from datetime import timedelta


	logger = get_logger(__name__)

	class StorageBackend:
	"""Base storage backend interface"""

	def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
	"""Generate a presigned URL for upload"""
	raise NotImplementedError

	def generate_download_url(self, storage_path):
	"""Generate a download URL"""
	raise NotImplementedError

	def delete_file(self, storage_path):
	"""Delete a file"""
	raise NotImplementedError

	class FileStorageBackend(StorageBackend):
	"""Local file storage backend"""

	def __init__(self, storage_path, base_url):
	self.storage_path = storage_path
	self.base_url = base_url
	os.makedirs(storage_path, exist_ok=True)

	def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
	"""For file storage, return a local path and upload endpoint

	Uses request_base_url if provided (for proxy/k8s scenarios), otherwise falls back to
	configured base_url or relative URL.
	"""
	# Sanitize file_name to prevent directory traversal
	# Remove any path separators and normalize
	safe_file_name = os.path.basename(file_name.replace('\\', '/').split('/')[-1])
	if not safe_file_name or safe_file_name.startswith('.'):
	safe_file_name = 'file.bin'

	file_id = session_id or str(uuid.uuid4())
	storage_key = f"{file_id}/{safe_file_name}"
	local_path = os.path.join(self.storage_path, storage_key)

	# Verify the resolved path is within storage_path (prevent traversal)
	real_storage_path = os.path.realpath(self.storage_path)
	real_local_path = os.path.realpath(local_path)
	if not real_local_path.startswith(real_storage_path):
	raise ValueError(f"Invalid file path: directory traversal detected")

	# Create directory
	os.makedirs(os.path.dirname(local_path), exist_ok=True)

	# Build upload URL - prefer request_base_url (works with proxies/k8s)
	relative_upload_path = f"/evidence/upload/{file_id}"
	if request_base_url:
	# Use the actual request base URL (includes correct protocol/host/port)
	# Remove trailing slash from base_url if present
	base = request_base_url.rstrip('/')
	upload_url = f"{base}{relative_upload_path}"
	else:
	# Fallback to relative URL for same-origin requests
	upload_url = relative_upload_path

	return {
	'upload_url': upload_url,
	'storage_path': storage_key,
	'local_path': local_path
	}

	def generate_download_url(self, storage_path):
	"""Generate download URL for local files"""
	# For API access, prepend 'evidence/download/' to the storage path
	api_path = f"evidence/download/{storage_path}"
	return urljoin(self.base_url, api_path)

	def delete_file(self, storage_path):
	"""Delete local file"""
	local_path = os.path.join(self.storage_path, storage_path)

	# SECURITY: Verify the resolved path is within storage directory
	real_storage_path = os.path.realpath(self.storage_path)
	real_local_path = os.path.realpath(local_path)
	if not real_local_path.startswith(real_storage_path):
	logger.error(f"Directory traversal attempt detected in delete: {storage_path}")
	return False

	try:
	if os.path.exists(local_path):
	os.remove(local_path)
	# Also remove directory if empty
	dir_path = os.path.dirname(local_path)
	if os.path.exists(dir_path) and not os.listdir(dir_path):
	os.rmdir(dir_path)
	return True
	except Exception as e:
	logger.error(f"Failed to delete file {storage_path}: {e}")
	return False

	class S3StorageBackend(StorageBackend):
	"""
	Works with real AWS S3.
	- Uses SigV4.
	- Supports temporary creds (session token).
	- Avoids forcing endpoint_url (let boto choose correct regional endpoint).
	- Uses virtual-hosted–style addressing if bucket name is DNS-safe.
	"""

	def __init__(self, access_key, secret_key, region, bucket, session_token=None):
	self.bucket = bucket

	# IMPORTANT: do NOT force endpoint_url for AWS; let boto3 pick the regional S3 endpoint.
	# Only set endpoint_url when talking to S3-compatible services like MinIO.
	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)


	def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
	file_id = session_id or str(uuid.uuid4())
	storage_key = f"attestations/{file_id}/{file_name}"

	# Normalize storage path (replace double slashes with single)
	storage_key = storage_key.replace('//', '/')

	# For S3 compatibility with valint (CLI client that always uses multipart/form-data)
	# Use multipart/form-data as content type to match what valint sends
	# This is a raw upload with the multipart/form-data header, not actual multipart encoding
	base_ct = "multipart/form-data"

	logger.info(f"S3 Upload URL generation - Original CT: '{content_type}', Using: '{base_ct}' for S3 compatibility, File: {file_name}")

	params = {
	"Bucket": self.bucket,
	"Key": storage_key,
	"ContentType": base_ct,
	}

	url = self.s3_client.generate_presigned_url(
	"put_object",
	Params=params,
	ExpiresIn=current_app.config["STORAGE_CONFIG"].PRESIGNED_URL_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):
	try:
	# Extract filename from storage path for Content-Disposition header
	filename = storage_path.split('/')[-1] if '/' in storage_path else storage_path

	url = self.s3_client.generate_presigned_url(
	"get_object",
	Params={
	"Bucket": self.bucket,
	"Key": storage_path,
	"ResponseContentDisposition": f'attachment; filename="{filename}"'
	},
	ExpiresIn=current_app.config["STORAGE_CONFIG"].PRESIGNED_URL_EXPIRATION,
	)
	logger.info(f"S3 presigned GET: {url}")
	return url
	except Exception as e:
	logger.error(f"Failed to generate S3 download URL for {storage_path}: {e}")
	return None

	def delete_file(self, storage_path):
	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(StorageBackend):
	"""MinIO storage backend"""

	def __init__(self, endpoint, access_key, secret_key, bucket, secure=True):
	self.bucket = bucket

	# Strip http:// or https:// from endpoint if present
	if endpoint.startswith("http://"):
	endpoint = endpoint[7:]
	secure = False
	elif endpoint.startswith("https://"):
	endpoint = endpoint[8:]
	secure = True

	self.client = Minio(
	endpoint,
	access_key=access_key,
	secret_key=secret_key,
	secure=secure
	)

	# Create bucket if it doesn't exist
	try:
	if not self.client.bucket_exists(bucket):
	self.client.make_bucket(bucket)
	except Exception as e:
	logger.error(f"Failed to create/check bucket {bucket}: {e}")

	def generate_upload_url(self, file_name, content_type=None, session_id=None, request_base_url=None):
	"""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=current_app.config['STORAGE_CONFIG'].PRESIGNED_URL_EXPIRATION)
	)


	return {
	'upload_url': upload_url,
	'storage_path': storage_key
	}

	def generate_download_url(self, storage_path):
	"""Generate presigned URL for MinIO download"""
	return self.client.get_presigned_url(
	"GET",
	self.bucket,
	storage_path,
	expires=timedelta(seconds=current_app.config['STORAGE_CONFIG'].PRESIGNED_URL_EXPIRATION)
	)

	def delete_file(self, storage_path):
	"""Delete MinIO object"""
	try:
	self.client.remove_object(self.bucket, storage_path)
	return True
	except Exception as e:
	logger.error(f"Failed to delete MinIO object {storage_path}: {e}")
	return False

	def get_storage_backend():
	"""Get configured storage backend"""
	config = current_app.config['STORAGE_CONFIG']

	if config.STORAGE_TYPE == 'FILE_MOUNT':
	return FileStorageBackend(config.FILE_STORAGE_PATH, config.FILE_STORAGE_BASE_URL)

	elif config.STORAGE_TYPE == 'CLOUD_STORAGE':
	# Check for MinIO first
	if config.CLOUD_STORAGE_PROVIDER == 'MINIO':
	return MinIOStorageBackend(
	config.MINIO_ENDPOINT,
	config.MINIO_ACCESS_KEY,
	config.MINIO_SECRET_KEY,
	config.MINIO_BUCKET,
	config.MINIO_SECURE
	)
	# Otherwise use AWS S3
	elif config.CLOUD_STORAGE_PROVIDER == 'AWS':
	return S3StorageBackend(
	config.AWS_ACCESS_KEY_ID,
	config.AWS_SECRET_ACCESS_KEY,
	config.AWS_REGION,
	config.S3_BUCKET,
	session_token=config.AWS_SESSION_TOKEN if hasattr(config, 'AWS_SESSION_TOKEN') else None
	)
	else:
	raise ValueError("Cloud storage configured but no valid backend settings found")

	else:
	raise ValueError(f"Unknown storage type: {config.STORAGE_TYPE}")