"""
This module was (shamelessly) copied directly from tinys3.auth

Changes from tinys3
-------------------
Despite the AWS documentation which seems to imply that there may exist 
request query parameters which should not be signed - this module now 
signs every query parameter. (AWS_QUERY_PARAMS was removed from the 
implementation)
"""

from requests.auth import AuthBase
from requests.structures import CaseInsensitiveDict

import base64
import datetime
import hashlib
import hmac
import re

try:
    from urlparse import urlparse
except ImportError:
    from urllib.parse import urlparse

# A regexp used for detecting aws bucket names
BUCKET_VHOST_MATCH = re.compile(
        r'^([a-z0-9\-]+\.)?s3([a-z0-9\-]+)?\.amazonaws\.com$', 
        flags=re.IGNORECASE)

class S3Auth(AuthBase):
    """
    S3 Custom Authenticator class for requests
    This code was copied from tinys3.auth.

    This authenticator will sign your requests based on the RESTAuthentication 
    specs by Amazon 
    http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html

    You can read more about custom authenticators here:
    http://docs.python-requests.org/en/latest/user/advanced.html#custom-authentication

    Usage:

    >>> from s3 import S3Auth
    >>> requests.put(
                url, 
                auth=S3Auth(access_key_id, secret_access_key),
                ...
                )
    """

    def __init__(self, access_key_id, secret_access_key):
        self.access_key_id = access_key_id
        self.secret_access_key = secret_access_key

    def __call__(self, r):
        """
        The entry point of the custom authenticator.

        When used as an auth class, requests will call this method just before 
        sending the request.

        Params:
        r           The request object

        Returns:    The request object, after we've updated some headers
        """

        # Generate the string to sign
        msg = self._string_to_sign(r)
        # Sign the string and add the authorization header
        r.headers['Authorization'] = "AWS %s:%s" % (
                self.access_key_id, 
                self._sign(msg))

        # Fix an issue with 0 length requests
        self._fix_content_length(r)

        # return the request
        return r

    def _fix_content_length(self, request):
        """
        Amazon requires to have content-length header when using the put 
        request;  however, requests won't add this header, so we need to add 
        it ourselves.

        Params:
        request         The request object
        """

        if request.method == 'PUT' and not 'Content-Length' in request.headers:
            request.headers['Content-Length'] = '0'

    def _get_canonicalized_amz_headers(self, headers):
        """
        Collect the special Amazon headers, prepare them for signing

        Params:
        headers -   CaseInsensitiveDict with the header requests

        Returns:    String with the canonicalized headers

        More information about this process here:
        http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#RESTAuthenticationConstructingCanonicalizedAmzHeaders
        """ 

        # New dict for the amazon headers
        amz_dict = {}

        # Go over the existing headers
        for k, v in headers.items():
            # Decode the keys if they are encoded
            if isinstance(k, bytes):
                k = k.decode('ascii')

            # to lower case
            k = k.lower()

            # If it starts with 'x-amz' add it to our dict
            if k.startswith('x-amz'):
                amz_dict[k] = v

        result = ""
        # Sort the keys and iterate through them
        for k in sorted(amz_dict.keys()):
            # add stripped key and value to the result string
            result += "%s:%s\n" % (
                    k.strip(), 
                    amz_dict[k].strip().replace('\n', ' '))

        # Return the result string
        return result

    def _get_canonicalized_resource(self, request):
        """
        Generates the canonicalized resource string form a request

        You can read more about the process here:
        http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#ConstructingTheCanonicalizedResourceElement

        Params:
        request -       The request object

        Returns:        String the canoncicalized resource string.
        """

        r = ""

        # parse our url
        parts = urlparse(request.url)

        # get the host, remove any port identifiers
        host = parts.netloc.split(':')[0]

        if host:
            # try to match our host to <hostname>.s3.amazonaws.com/s3.amazonaws.com
            m = BUCKET_VHOST_MATCH.match(host)
            if m:
                bucket = (m.groups()[0] or '').rstrip('.')

                if bucket:
                    r += ('/' + bucket)
            else:
                # It's a virtual host, add it to the result
                r += ('/' + host)

        # Add the path string
        r += parts.path or '/'

        # add the special query strings
        r += self._get_subresource(parts.query)

        return r

    def _get_date(self):
        """
        Returns a string for the current date
        """
        return datetime.datetime.utcnow().strftime('%a, %d %b %Y %H:%M:%S GMT')

    def _get_subresource(self, qs):
        """
        Handle subresources in the query string

        More information about subresources:
        http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#ConstructingTheCanonicalizedResourceElement
        """
        r = []

        # Split the querystring
        keys = qs.split('&')
        # for each item
        for i in keys:
            # get the key
            item = i.split('=')
            k = item[0].lower()
            if k:
                r.append(i)
        # If we have result, convert them to query string
        if r:
            return '?' + '&'.join(r)

        return ''

    def _sign(self, string_to_sign):
        """
        Generates a signature for the given string

        Params:
        string_to_sign -    The string we want to sign
        
        Returns:            Signature in bytes
        """
        digest = hmac.new(
                self.secret_access_key.encode('utf8'),
                msg=string_to_sign.encode('utf8'),
                digestmod=hashlib.sha1).digest()

        return base64.b64encode(digest).strip().decode('ascii')

    def _string_to_sign(self, request):
        """
        Generates the string we need to sign on.

        Params:
        request -   The request object

        Returns:    String ready to be signed on
        """

        # We'll use case insensitive dict to store the headers
        h = CaseInsensitiveDict()
        # Add the hearders
        h.update(request.headers)

        # If we have an 'x-amz-date' header, we'll try to use it instead of 
        # the date
        if b'x-amz-date' in h or 'x-amz-date' in h:
            date = ''
        else:
            # No x-amz-header, we'll generate a date
            date = h.get('Date') or self._get_date()

        # Set the date header
        request.headers['Date'] = date

        # A fix for the content type header extraction in python 3
        # This have to be done because requests will try to set 
        # application/www-url-encoded header if we pass bytes as the content, 
        # and the content-type is set with a key that is b'Content-Type' and 
        #not 'Content-Type'
        content_type = ''
        if b'Content-Type' in request.headers:
            # Fix content type
            content_type = h.get(b'Content-Type')
            del request.headers[b'Content-Type']
            request.headers['Content-Type'] = content_type

        # The string we're about to generate
        # There's more information about it here:
        # http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html#ConstructingTheAuthenticationHeader
        msg = [
                # HTTP Method
                request.method,
                # MD5 If provided
                h.get(b'Content-MD5', '') or h.get('Content-MD5', ''),
                # Content type if provided
                content_type or h.get('Content-Type', ''),
                # Date
                date,
                # Canonicalized special amazon headers and resource uri
                self._get_canonicalized_amz_headers(h) 
                        + self._get_canonicalized_resource(request),
                ]
        return '\n'.join(msg)

