# -*- coding: utf-8 -*-
import unittest
import random
import string

from flask import Flask
from flask.ext.passlib import Passlib

from werkzeug.security import generate_password_hash, check_password_hash

WERKZEUG_METHODS = ['md5', 'sha1']
PASSLIB_METHODS = ['werkzeug_salted_md5', 'werkzeug_salted_sha1']

class FlaskTestCase(unittest.TestCase):
    TESTING = True

    def setUp(self):
        self.app = Flask(__name__)
        self.app.config.from_object(self)

        self.assertTrue(self.app.testing)

        self.passlib = Passlib(self.app)

        self.ctx = self.app.test_request_context()
        self.ctx.push()

    def tearDown(self):
        self.ctx.pop()

    def generate_password(self, length=6):
        return ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(length))

class TestCase(object):
    ident = None
    werkzeug_method = None
    passlib_method = None
    salt_length = 8
    checksum_length = 0
    seperator = '$'

    def test_identify(self):
        password = self.generate_password()

        # Identify hashes generated by werkzeug
        passlib_hash = generate_password_hash(password, method=self.werkzeug_method)
        self.assertEqual(self.passlib.identify(passlib_hash), self.passlib_method)

        # Identify hashes generated by other werkzeug methods
        diff_werkzeug_method = random.choice([method for method in WERKZEUG_METHODS if method not in [self.werkzeug_method]])
        passlib_hash = generate_password_hash(password, method=diff_werkzeug_method)
        self.assertNotEqual(self.passlib.identify(passlib_hash), self.passlib_method)

        # Identify hashes generated by werkzeug with varying salt length
        passlib_hash = generate_password_hash(password, method=self.werkzeug_method, salt_length=self.salt_length + 1)
        self.assertEqual(self.passlib.identify(passlib_hash), self.passlib_method)

        # Identify hashes generated by itself
        passlib_hash = self.passlib.encrypt(password, method=self.passlib_method)
        self.assertEqual(self.passlib.identify(passlib_hash), self.passlib_method)

        # Identify hashes generated by other methods
        diff_passlib_method = random.choice([method for method in PASSLIB_METHODS if method not in [self.passlib_method]])
        passlib_hash = self.passlib.encrypt(password, method=diff_passlib_method)
        self.assertNotEqual(self.passlib.identify(passlib_hash), self.passlib_method)

        # Identify hashes generated by itself with varying salt length
        passlib_hash = self.passlib.encrypt(password, method=self.passlib_method, salt_length=self.salt_length + 1)
        self.assertEqual(self.passlib.identify(passlib_hash), self.passlib_method)

    def test_verify(self):
        password = self.generate_password()
        werkzeug_hash = generate_password_hash(password, method=self.werkzeug_method)
        passlib_hash = self.passlib.encrypt(password, scheme=self.passlib_method)

        self.assertTrue(self.passlib.verify(password, werkzeug_hash))
        self.assertTrue(self.passlib.verify(password, passlib_hash))


        password2 = self.generate_password()
        werkzeug_hash2 = generate_password_hash(password2, method=self.werkzeug_method)
        passlib_hash2 = self.passlib.encrypt(password2, scheme=self.passlib_method)

        self.assertTrue(self.passlib.verify(password2, werkzeug_hash2))
        self.assertTrue(self.passlib.verify(password2, passlib_hash2))

        self.assertFalse(self.passlib.verify(password, werkzeug_hash2))
        self.assertFalse(self.passlib.verify(password, passlib_hash2))
        self.assertFalse(self.passlib.verify(password2, werkzeug_hash))
        self.assertFalse(self.passlib.verify(password2, passlib_hash))

    def test_encrypt(self):
        password = self.generate_password()

        passlib_hash = self.passlib.encrypt(password, scheme=self.passlib_method, salt_length=self.salt_length)
        self.assertTrue(check_password_hash(passlib_hash, password))
        self.assertIn(self.ident, passlib_hash)
        self.assertTrue(len(passlib_hash.split(self.seperator)), 3)
        self.assertTrue(passlib_hash.startswith(self.ident))
        self.assertTrue(len(passlib_hash.split(self.seperator)[1]), self.salt_length)
        self.assertTrue(len(passlib_hash.split(self.seperator)[2]), self.checksum_length)


class MD5TestCase(TestCase, FlaskTestCase):
    ident = 'md5$'
    werkzeug_method = 'md5'
    passlib_method = 'werkzeug_salted_md5'
    checksum_length = 32

class SHA1TestCase(TestCase, FlaskTestCase):
    ident = 'sha1$'
    werkzeug_method = 'sha1'
    passlib_method = 'werkzeug_salted_sha1'
    checksum_length = 40

class SHA224TestCase(TestCase, FlaskTestCase):
    ident = 'sha224$'
    werkzeug_method = 'sha224'
    passlib_method = 'werkzeug_salted_sha224'
    checksum_length = 56

class SHA256TestCase(TestCase, FlaskTestCase):
    ident = 'sha256$'
    werkzeug_method = 'sha256'
    passlib_method = 'werkzeug_salted_sha256'
    checksum_length = 64

class SHA384TestCase(TestCase, FlaskTestCase):
    ident = 'sha384$'
    werkzeug_method = 'sha384'
    passlib_method = 'werkzeug_salted_sha384'
    checksum_length = 96

class SHA512TestCase(TestCase, FlaskTestCase):
    ident = 'sha512$'
    werkzeug_method = 'sha512'
    passlib_method = 'werkzeug_salted_sha512'
    checksum_length = 128


            
if __name__ == '__main__':
    unittest.main()