Ticket #6390: s3.py

from mimetypes import guess_type
from StringIO import StringIO
import urlparse
import os

from django.core.exceptions import ImproperlyConfigured
from django.core.filestorage.base import Backend, RemoteFile
from django.utils.functional import curry

ACCESS_KEY_NAME = 'AWS_ACCESS_KEY_ID'
SECRET_KEY_NAME = 'AWS_SECRET_ACCESS_KEY'

try:
    from boto.s3.connection import S3Connection
except ImportError:
    raise ImproperlyConfigured, "Could not load boto's S3 bindings."

class S3Backend(Backend):
    """Amazon Simple Storage Service"""

    def __init__(self, bucket, access_key=None, secret_key=None, acl='public-read'):
        self.bucket = bucket

        if not access_key and not secret_key:
             access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(access_key, secret_key)
        self.generator = S3.QueryStringAuthGenerator(access_key, secret_key)

    def _get_access_keys(self):
        access_key = getattr(settings, ACCESS_KEY_NAME, None)
        secret_key = getattr(settings, SECRET_KEY_NAME, None)
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_connection(self):
        return S3Connection(*self._get_access_keys())

    def _put_file(self, filename, raw_contents):
        content_type = guess_type(filename)[0] or "application/x-octet-stream"
        headers = {'x-amz-acl':  self.acl, 'Content-Type': content_type}
        self.connection.put(self.bucket, self.filename, raw_contents, headers)

    def get_filename(self, filename):
        return urlparse.urljoin('/%s/' % self.bucket, filename)

    def get_absolute_url(self, filename):
        return self.generator.make_bare_url(self.bucket, filename)

    def get_filesize(self, filename):
        response = self.connection.make_request('HEAD', self.bucket, filename)
        return int(response.getheader('Content-Length'))

    def open_file(self, filename, mode='rb'):
        response = self.connection.get(self.bucket, filename)
        writer = curry(self._put_file, filename)
        return RemoteFile(self, response.object.data, mode, writer)

    def file_exists(self, filename):
        response = self.connection.make_request('HEAD', self.bucket, filename)
        return response.status == 200

    def save_file(self, filename, raw_contents):
        filename = self.get_available_filename(filename)
        self._put_file(filename, raw_contents)
        return filename

    def delete_file(self, filename):
        self.connection.delete(self.bucket, self.key)