# Test hashlib module
#
# $Id$
#
#  Copyright (C) 2005-2010   Gregory P. Smith (greg@krypto.org)
#  Licensed to PSF under a Contributor Agreement.
#

import array
from binascii import unhexlify
import functools
import hashlib
import importlib
import itertools
import os
import sys
import threading
import unittest
import warnings
from test import support
from test.support import _4G, bigmemtest, import_fresh_module
from test.support import requires_hashdigest
from http.client import HTTPException

# Were we compiled --with-pydebug or with #define Py_DEBUG?
COMPILED_WITH_PYDEBUG = hasattr(sys, 'gettotalrefcount')

c_hashlib = import_fresh_module('hashlib', fresh=['_hashlib'])
py_hashlib = import_fresh_module('hashlib', blocked=['_hashlib'])

try:
    from _hashlib import HASH
except ImportError:
    HASH = None

try:
    import _blake2
except ImportError:
    _blake2 = None

requires_blake2 = unittest.skipUnless(_blake2, 'requires _blake2')

try:
    import _sha3
except ImportError:
    _sha3 = None

requires_sha3 = unittest.skipUnless(_sha3, 'requires _sha3')


def hexstr(s):
    assert isinstance(s, bytes), repr(s)
    h = "0123456789abcdef"
    r = ''
    for i in s:
        r += h[(i >> 4) & 0xF] + h[i & 0xF]
    return r


URL = "http://www.pythontest.net/hashlib/{}.txt"

def read_vectors(hash_name):
    url = URL.format(hash_name)
    try:
        testdata = support.open_urlresource(url)
    except (OSError, HTTPException):
        raise unittest.SkipTest("Could not retrieve {}".format(url))
    with testdata:
        for line in testdata:
            line = line.strip()
            if line.startswith('#') or not line:
                continue
            parts = line.split(',')
            parts[0] = bytes.fromhex(parts[0])
            yield parts


class HashLibTestCase(unittest.TestCase):
    supported_hash_names = ( 'md5', 'MD5', 'sha1', 'SHA1',
                             'sha224', 'SHA224', 'sha256', 'SHA256',
                             'sha384', 'SHA384', 'sha512', 'SHA512',
                             'blake2b', 'blake2s',
                             'sha3_224', 'sha3_256', 'sha3_384', 'sha3_512',
                             'shake_128', 'shake_256')

    shakes = {'shake_128', 'shake_256'}

    # Issue #14693: fallback modules are always compiled under POSIX
    _warn_on_extension_import = os.name == 'posix' or COMPILED_WITH_PYDEBUG

    def _conditional_import_module(self, module_name):
        """Import a module and return a reference to it or None on failure."""
        try:
            return importlib.import_module(module_name)
        except ModuleNotFoundError as error:
            if self._warn_on_extension_import:
                warnings.warn('Did a C extension fail to compile? %s' % error)
        return None

    def __init__(self, *args, **kwargs):
        algorithms = set()
        for algorithm in self.supported_hash_names:
            algorithms.add(algorithm.lower())

        _blake2 = self._conditional_import_module('_blake2')
        if _blake2:
            algorithms.update({'blake2b', 'blake2s'})

        self.constructors_to_test = {}
        for algorithm in algorithms:
            self.constructors_to_test[algorithm] = set()

        # For each algorithm, test the direct constructor and the use
        # of hashlib.new given the algorithm name.
        for algorithm, constructors in self.constructors_to_test.items():
            constructors.add(getattr(hashlib, algorithm))
            def _test_algorithm_via_hashlib_new(data=None, _alg=algorithm, **kwargs):
                if data is None:
                    return hashlib.new(_alg, **kwargs)
                return hashlib.new(_alg, data, **kwargs)
            constructors.add(_test_algorithm_via_hashlib_new)

        _hashlib = self._conditional_import_module('_hashlib')
        self._hashlib = _hashlib
        if _hashlib:
            # These two algorithms should always be present when this module
            # is compiled.  If not, something was compiled wrong.
            self.assertTrue(hasattr(_hashlib, 'openssl_md5'))
            self.assertTrue(hasattr(_hashlib, 'openssl_sha1'))
            for algorithm, constructors in self.constructors_to_test.items():
                constructor = getattr(_hashlib, 'openssl_'+algorithm, None)
                if constructor:
                    try:
                        constructor()
                    except ValueError:
                        # default constructor blocked by crypto policy
                        pass
                    else:
                        constructors.add(constructor)

        def add_builtin_constructor(name):
            constructor = getattr(hashlib, "__get_builtin_constructor")(name)
            self.constructors_to_test[name].add(constructor)

        _md5 = self._conditional_import_module('_md5')
        if _md5:
            add_builtin_constructor('md5')
        _sha1 = self._conditional_import_module('_sha1')
        if _sha1:
            add_builtin_constructor('sha1')
        _sha256 = self._conditional_import_module('_sha256')
        if _sha256:
            add_builtin_constructor('sha224')
            add_builtin_constructor('sha256')
        _sha512 = self._conditional_import_module('_sha512')
        if _sha512:
            add_builtin_constructor('sha384')
            add_builtin_constructor('sha512')
        if _blake2:
            add_builtin_constructor('blake2s')
            add_builtin_constructor('blake2b')

        _sha3 = self._conditional_import_module('_sha3')
        if _sha3:
            add_builtin_constructor('sha3_224')
            add_builtin_constructor('sha3_256')
            add_builtin_constructor('sha3_384')
            add_builtin_constructor('sha3_512')
            add_builtin_constructor('shake_128')
            add_builtin_constructor('shake_256')

        super(HashLibTestCase, self).__init__(*args, **kwargs)

    @property
    def hash_constructors(self):
        constructors = self.constructors_to_test.values()
        return itertools.chain.from_iterable(constructors)

    @support.refcount_test
    @unittest.skipIf(c_hashlib is None, 'Require _hashlib module')
    def test_refleaks_in_hash___init__(self):
        gettotalrefcount = support.get_attribute(sys, 'gettotalrefcount')
        sha1_hash = c_hashlib.new('sha1')
        refs_before = gettotalrefcount()
        for i in range(100):
            sha1_hash.__init__('sha1')
        self.assertAlmostEqual(gettotalrefcount() - refs_before, 0, delta=10)

    def test_hash_array(self):
        a = array.array("b", range(10))
        for cons in self.hash_constructors:
            c = cons(a)
            if c.name in self.shakes:
                c.hexdigest(16)
            else:
                c.hexdigest()

    def test_algorithms_guaranteed(self):
        self.assertEqual(hashlib.algorithms_guaranteed,
            set(_algo for _algo in self.supported_hash_names
                  if _algo.islower()))

    def test_algorithms_available(self):
        self.assertTrue(set(hashlib.algorithms_guaranteed).
                            issubset(hashlib.algorithms_available))

    def test_unknown_hash(self):
        self.assertRaises(ValueError, hashlib.new, 'spam spam spam spam spam')
        self.assertRaises(TypeError, hashlib.new, 1)

    def test_new_upper_to_lower(self):
        self.assertEqual(hashlib.new("SHA256").name, "sha256")

    def test_get_builtin_constructor(self):
        get_builtin_constructor = getattr(hashlib,
                                          '__get_builtin_constructor')
        builtin_constructor_cache = getattr(hashlib,
                                            '__builtin_constructor_cache')
        self.assertRaises(ValueError, get_builtin_constructor, 'test')
        try:
            import _md5
        except ImportError:
            self.skipTest("_md5 module not available")
        # This forces an ImportError for "import _md5" statements
        sys.modules['_md5'] = None
        # clear the cache
        builtin_constructor_cache.clear()
        try:
            self.assertRaises(ValueError, get_builtin_constructor, 'md5')
        finally:
            if '_md5' in locals():
                sys.modules['_md5'] = _md5
            else:
                del sys.modules['_md5']
        self.assertRaises(TypeError, get_builtin_constructor, 3)
        constructor = get_builtin_constructor('md5')
        self.assertIs(constructor, _md5.md5)
        self.assertEqual(sorted(builtin_constructor_cache), ['MD5', 'md5'])

    def test_hexdigest(self):
        for cons in self.hash_constructors:
            h = cons()
            if h.name in self.shakes:
                self.assertIsInstance(h.digest(16), bytes)
                self.assertEqual(hexstr(h.digest(16)), h.hexdigest(16))
            else:
                self.assertIsInstance(h.digest(), bytes)
                self.assertEqual(hexstr(h.digest()), h.hexdigest())

    def test_digest_length_overflow(self):
        # See issue #34922
        large_sizes = (2**29, 2**32-10, 2**32+10, 2**61, 2**64-10, 2**64+10)
        for cons in self.hash_constructors:
            h = cons()
            if h.name not in self.shakes:
                continue
            for digest in h.digest, h.hexdigest:
                self.assertRaises(ValueError, digest, -10)
                for length in large_sizes:
                    with self.assertRaises((ValueError, OverflowError)):
                        digest(length)

    def test_name_attribute(self):
        for cons in self.hash_constructors:
            h = cons()
            self.assertIsInstance(h.name, str)
            if h.name in self.supported_hash_names:
                self.assertIn(h.name, self.supported_hash_names)
            else:
                self.assertNotIn(h.name, self.supported_hash_names)
            self.assertEqual(h.name, hashlib.new(h.name).name)

    def test_large_update(self):
        aas = b'a' * 128
        bees = b'b' * 127
        cees = b'c' * 126
        dees = b'd' * 2048 #  HASHLIB_GIL_MINSIZE

        for cons in self.hash_constructors:
            m1 = cons()
            m1.update(aas)
            m1.update(bees)
            m1.update(cees)
            m1.update(dees)
            if m1.name in self.shakes:
                args = (16,)
            else:
                args = ()

            m2 = cons()
            m2.update(aas + bees + cees + dees)
            self.assertEqual(m1.digest(*args), m2.digest(*args))

            m3 = cons(aas + bees + cees + dees)
            self.assertEqual(m1.digest(*args), m3.digest(*args))

            # verify copy() doesn't touch original
            m4 = cons(aas + bees + cees)
            m4_digest = m4.digest(*args)
            m4_copy = m4.copy()
            m4_copy.update(dees)
            self.assertEqual(m1.digest(*args), m4_copy.digest(*args))
            self.assertEqual(m4.digest(*args), m4_digest)

    def check(self, name, data, hexdigest, shake=False, **kwargs):
        length = len(hexdigest)//2
        hexdigest = hexdigest.lower()
        constructors = self.constructors_to_test[name]
        # 2 is for hashlib.name(...) and hashlib.new(name, ...)
        self.assertGreaterEqual(len(constructors), 2)
        for hash_object_constructor in constructors:
            m = hash_object_constructor(data, **kwargs)
            computed = m.hexdigest() if not shake else m.hexdigest(length)
            self.assertEqual(
                    computed, hexdigest,
                    "Hash algorithm %s constructed using %s returned hexdigest"
                    " %r for %d byte input data that should have hashed to %r."
                    % (name, hash_object_constructor,
                       computed, len(data), hexdigest))
            computed = m.digest() if not shake else m.digest(length)
            digest = bytes.fromhex(hexdigest)
            self.assertEqual(computed, digest)
            if not shake:
                self.assertEqual(len(digest), m.digest_size)

    def check_no_unicode(self, algorithm_name):
        # Unicode objects are not allowed as input.
        constructors = self.constructors_to_test[algorithm_name]
        for hash_object_constructor in constructors:
            self.assertRaises(TypeError, hash_object_constructor, 'spam')

    def test_no_unicode(self):
        self.check_no_unicode('md5')
        self.check_no_unicode('sha1')
        self.check_no_unicode('sha224')
        self.check_no_unicode('sha256')
        self.check_no_unicode('sha384')
        self.check_no_unicode('sha512')

    @requires_blake2
    def test_no_unicode_blake2(self):
        self.check_no_unicode('blake2b')
        self.check_no_unicode('blake2s')

    @requires_sha3