# encoding: utf8
"""

TODO
=====

- finalize on naming of extra/depends_on
    The only other thing I can thing of is include/depends_on

- consider adding loads_ex which returns (data, meta)
    Would need to add `extra` to `dumps` for this to make sense.
    
- consider adding `mode` parameter to dumps:
    - possible modes:
        None -> byte string
        'u'  -> unicode
        'r'  -> deterministic repr; parsable with ast.literal_eval
        'p'  -> pickle
    - would then only have dumps/loads
        
- document that while creation time is encoded, it is still public for signatures

- document the various signature/meta keys
    t -> creation time
    x -> expiration time
    z -> compressed (ignores value)
    i -> encryption iv
    n -> nonce
    s -> signature MAC
    m -> data type

"""

from __future__ import division

import base64
import hashlib
import hmac
import json
import os
import re
import string
import struct
import time
import urllib
import urlparse
import zlib
import pickle

from . import repr as deterministic_repr
import tomcrypt.cipher


CREATION_TIME_KEY = 't' # legacy
EXPIRATION_TIME_KEY = 'x' # legacy
COMPRESSION_FLAG_KEY = 'z'
ENCRYPTION_IV_KEY = 'i'
NONCE_KEY = 'n' # legacy, but not checked
SIGNATURE_MAC_KEY = 's' # legacy
DATA_FORMAT_KEY = 'f'

meta_keys = set([
    CREATION_TIME_KEY,
    EXPIRATION_TIME_KEY,
    COMPRESSION_FLAG_KEY,
    ENCRYPTION_IV_KEY,
    NONCE_KEY,
    SIGNATURE_MAC_KEY,
    DATA_FORMAT_KEY,
])

# Build up the map we will use to determine which hash was used by the size of
# its digest.
hash_by_size = {}
algorithms = getattr(hashlib, 'algorithms', 'md5 sha1 sha256'.split())
for name in algorithms:
    hash_by_size[getattr(hashlib, name)().digest_size] = getattr(hashlib, name)


# Build up the map we will use to encode transport-unsafe characters.
encode_seal_map = {}
for i in xrange(256):
    encode_seal_map[chr(i)] = '$%02x' % i
for c in (string.letters + string.digits + '_-\'()'):
    encode_seal_map[c] = c


def dollar_encode(input, safe=''):
    """Encode all characters that are not URL/cookie safe.
    
    This is similar to percent-encoding, but we use a dollar sign to avoid
    collisions with it.
    
    >>> dollar_encode('hello')
    'hello'
    
    >>> dollar_encode('before.after')
    'before$2eafter'
    
    >>> dollar_encode('before.after', '.')
    'before.after'
    
    >>> dollar_encode(''.join(chr(i) for i in range(10)))
    '$00$01$02$03$04$05$06$07$08$09'
    
    """
    if safe:
        char_map = encode_seal_map.copy()
        for c in safe:
            char_map[c] = c
    else:
        char_map = encode_seal_map
    return ''.join(char_map[c] for c in input)


def dollar_decode(input):
    """Decode a dollar-encoded string.
    
    This is more accepting than dollar_encode; it will decode any characters that
    have been encoded, not just unsafe ones.
    
    >>> dollar_decode('hello')
    'hello'
    
    >>> dollar_decode('$00$01$02$03$04$05$06$07$08$09')
    '\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07\\x08\\t'
    
    >>> all_bytes = ''.join(chr(i) for i in range(256))
    >>> dollar_decode(dollar_encode(all_bytes)) == all_bytes
    True
    
    """
    return re.sub(r'\$([a-fA-F0-9]{2})', lambda m: chr(int(m.group(1), 16)), input)


def encode_seal(data, meta):
    """Serialized a data string with associated meta data.
    
    The output data is formatted like:
    
        data.key1:value1,key2:value2 [...]
    
    All characters that are not letters, digits, underscore or hyphen will be
    dollar-encoded. The characters ".:,$" are used since they are all safe for
    use in URL queries, URL paths, and cookie values (with quoting), and in this
    usage they look roughly appropriate. Ergo, the results of this function are
    safe to include in URLs and cookies without escaping.
    
    The data, and all keys and values of the meta dict MUST be byte strings.
    
    >>> encode_seal('data', dict(key='value', key2='value2'))
    'data.key:value,key2:value2'
    
    >>> encode_seal(u'data', {})
    Traceback (most recent call last):
    ...
    TypeError: data must be str; not <type 'unicode'>
    
    >>> encode_seal('a.b', {'!@#': '$%^'})
    'a.b.$21$40$23:$24$25$5e'
    
    >>> encode_seal(u'data', {'key': 'value'})
    Traceback (most recent call last):
    ...
    TypeError: data must be str; not <type 'unicode'>
    
    >>> encode_seal('data', {u'key': 'value'})
    Traceback (most recent call last):
    ...
    TypeError: meta key u'key' must be str; not <type 'unicode'>
    
    >>> encode_seal('data', {'key': u'value'})
    Traceback (most recent call last):
    ...
    TypeError: meta 'key' value u'value' must be str; not <type 'unicode'>
    
    
    """
    if not isinstance(data, str):
        raise TypeError('data must be str; not %r' % type(data))
    meta_parts = []
    for k, v in sorted(meta.iteritems()):
        if not isinstance(k, str):
            raise TypeError('meta key %r must be str; not %r' % (k, type(k)))
        if not isinstance(v, str):
            raise TypeError('meta %r value %r must be str; not %r' % (k, v, type(v)))
        meta_parts.append((dollar_encode(k), dollar_encode(v)))
    return '%s.%s' % (dollar_encode(data, '.,:'), ','.join('%s:%s' % x for x in meta_parts))


def decode_seal(data):
    """Unserialize data from encode_seal; returns a tuple of original data and meta.
    
    >>> decode_seal('data.key:value')
    ('data', {'key': 'value'})
    
    >>
> decode_seal('a.b.$21$40$23:$24$25$5e')
    ('a.b', {'!@#': '$%^'})
    
    >>> all_bytes 
= ''.join(chr(i) for i in range(256))
    >>> meta = {}
    >>> for i in range(0, 256, 16):
    ...     key = ''.join(chr(c) for c in range(i, i + 8))
    ...     val = ''.join(chr(c) for c in range(i + 8, i + 16))
    ...     meta[key] = val
    >>> out_data, out_meta = decode_seal(encode_seal(all_bytes, meta))
    >>> out_data == all_bytes
    True
    >>> out_meta == meta
    True
    
    """
    data, meta = data.rsplit('.', 1)
    meta = dict(map(dollar_decode, x.split(':')) for x in meta.split(','))
    return dollar_decode(data), meta


def encode_query(data):
    """Encode a dictionary of strings for use as a query string.
    
    Key-value pairs will be serialized in sorted order in order for the results
    from this function to be deterministic.
    
    >>> encode_query(dict(a="one", b="two"))
    'a=one&b=two'
    
    """
    if isinstance(data, dict):
        data = data.items()
    return urllib.urlencode(sorted(data))


def decode_query(data):
    """Decode a URL query into a dictionary.
    
    >>> decode_query('a=one&b=two')
    {'a': 'one', 'b': 'two'}
    
    """
    return dict(urlparse.parse_qsl(data))


def encode_binary(data):
    """Encode binary data into a URL/cookie safe form (base64).
    
    We remove trailing padding characters to reduce space. Since this is no
    longer valid base64 we must use `decode_binary` to recover the original
    data.
    
    >>> encode_binary('hello')
    'aGVsbG8'
    
    >>> encode_binary('<=>?+@')
    'PD0-PytA'
    
    >>> encode_binary(''.join(chr(i) for i in range(10)))
    'AAECAwQFBgcICQ'
    
    """
    return base64.urlsafe_b64encode(data).rstrip('=')


def decode_binary(data):
    """Decode binary data from base64.
    
    This function will automatically restore stripped base64 padding.
    
    >>> decode_binary('aGVsbG8')
    'hello'
    
    >>> decode_binary('PD0-PytA')
    '<=>?+@'
    
    >>> decode_binary('AAECAwQFBgcICQ')
    '\\x00\\x01\\x02\\x03\\x04\\x05\\x06\\x07\\x08\\t'
    
    """
    # Fix base64 padding.
    pad_size = len(data) % 4
    if pad_size:
        data += ((4 - pad_size) * '=')
    # Must cast to string for decode.
    return base64.urlsafe_b64decode(str(data))


def encode_legacy_int(value):
    """Encode a positive integer into a base64 string.
    
    This implementation is limited to 64 bit integers, and is not as space
    efficient as it could be, but we are keeping it this way for backwards
    compatibility (for query signing).
    
    >>> [encode_legacy_int(x) for x in range(10)]
    ['', 'Q', 'g', 'w', 'BA', 'BQ', 'Bg'