#
#  SelfTest/Signature/test_dss.py: Self-test for DSS signatures
#
# ===================================================================
#
# Copyright (c) 2014, Legrandin <helderijs@gmail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ===================================================================

import re
import json
import unittest
from binascii import hexlify, unhexlify

from Crypto.Util.py3compat import tobytes, bord, bchr

from Crypto.Hash import SHA1, SHA224, SHA256, SHA384, SHA512
from Crypto.Signature import DSS
from Crypto.PublicKey import DSA, ECC
from Crypto.SelfTest.st_common import list_test_cases
from Crypto.SelfTest.loader import load_tests
from Crypto.Util.number import bytes_to_long, long_to_bytes

from Crypto.Util._file_system import pycryptodome_filename
from Crypto.Util.strxor import strxor


def t2b(hexstring):
    ws = hexstring.replace(" ", "").replace("\n", "")
    return unhexlify(tobytes(ws))


def t2l(hexstring):
    ws = hexstring.replace(" ", "").replace("\n", "")
    return int(ws, 16)


def load_hash_by_name(hash_name):
    return __import__("Crypto.Hash." + hash_name, globals(), locals(), ["new"])


class StrRNG:

    def __init__(self, randomness):
        length = len(randomness)
        self._idx = 0
        # Fix required to get the right K (see how randint() works!)
        self._randomness = long_to_bytes(bytes_to_long(randomness) - 1, length)

    def __call__(self, n):
        out = self._randomness[self._idx:self._idx + n]
        self._idx += n
        return out


class FIPS_DSA_Tests(unittest.TestCase):

    # 1st 1024 bit key from SigGen.txt
    P = 0xa8f9cd201e5e35d892f85f80e4db2599a5676a3b1d4f190330ed3256b26d0e80a0e49a8fffaaad2a24f472d2573241d4d6d6c7480c80b4c67bb4479c15ada7ea8424d2502fa01472e760241713dab025ae1b02e1703a1435f62ddf4ee4c1b664066eb22f2e3bf28bb70a2a76e4fd5ebe2d1229681b5b06439ac9c7e9d8bde283
    Q = 0xf85f0f83ac4df7ea0cdf8f469bfeeaea14156495
    G = 0x2b3152ff6c62f14622b8f48e59f8af46883b38e79b8c74deeae9df131f8b856e3ad6c8455dab87cc0da8ac973417ce4f7878557d6cdf40b35b4a0ca3eb310c6a95d68ce284ad4e25ea28591611ee08b8444bd64b25f3f7c572410ddfb39cc728b9c936f85f419129869929cdb909a6a3a99bbe089216368171bd0ba81de4fe33
    X = 0xc53eae6d45323164c7d07af5715703744a63fc3a
    Y = 0x313fd9ebca91574e1c2eebe1517c57e0c21b0209872140c5328761bbb2450b33f1b18b409ce9ab7c4cd8fda3391e8e34868357c199e16a6b2eba06d6749def791d79e95d3a4d09b24c392ad89dbf100995ae19c01062056bb14bce005e8731efde175f95b975089bdcdaea562b32786d96f5a31aedf75364008ad4fffebb970b

    key_pub = DSA.construct((Y, G, P, Q))
    key_priv = DSA.construct((Y, G, P, Q, X))

    def shortDescription(self):
        return "FIPS DSA Tests"

    def test_loopback(self):
        hashed_msg = SHA512.new(b"test")
        signer = DSS.new(self.key_priv, 'fips-186-3')
        signature = signer.sign(hashed_msg)

        verifier = DSS.new(self.key_pub, 'fips-186-3')
        verifier.verify(hashed_msg, signature)

    def test_negative_unapproved_hashes(self):
        """Verify that unapproved hashes are rejected"""

        from Crypto.Hash import RIPEMD160

        self.description = "Unapproved hash (RIPEMD160) test"
        hash_obj = RIPEMD160.new()
        signer = DSS.new(self.key_priv, 'fips-186-3')
        self.assertRaises(ValueError, signer.sign, hash_obj)
        self.assertRaises(ValueError, signer.verify, hash_obj, b"\x00" * 40)

    def test_negative_unknown_modes_encodings(self):
        """Verify that unknown modes/encodings are rejected"""

        self.description = "Unknown mode test"
        self.assertRaises(ValueError, DSS.new, self.key_priv, 'fips-186-0')

        self.description = "Unknown encoding test"
        self.assertRaises(ValueError, DSS.new, self.key_priv, 'fips-186-3', 'xml')

    def test_asn1_encoding(self):
        """Verify ASN.1 encoding"""

        self.description = "ASN.1 encoding test"
        hash_obj = SHA1.new()
        signer = DSS.new(self.key_priv, 'fips-186-3', 'der')
        signature = signer.sign(hash_obj)

        # Verify that output looks like a DER SEQUENCE
        self.assertEqual(bord(signature[0]), 48)
        signer.verify(hash_obj, signature)

        # Verify that ASN.1 parsing fails as expected
        signature = bchr(7) + signature[1:]
        self.assertRaises(ValueError, signer.verify, hash_obj, signature)

    def test_sign_verify(self):
        """Verify public/private method"""

        self.description = "can_sign() test"
        signer = DSS.new(self.key_priv, 'fips-186-3')
        self.failUnless(signer.can_sign())

        signer = DSS.new(self.key_pub, 'fips-186-3')
        self.failIf(signer.can_sign())

class FIPS_DSA_Tests_KAT(unittest.TestCase):
    pass

test_vectors_verify = load_tests(("Crypto", "SelfTest", "Signature", "test_vectors", "DSA"),
                                 "FIPS_186_3_SigVer.rsp",
                                 "Signature Verification 186-3",
                                 {'result' : lambda x: x})

for idx, tv in enumerate(test_vectors_verify):

    if isinstance(tv, str):
        res = re.match(r"\[mod = L=([0-9]+), N=([0-9]+), ([a-zA-Z0-9-]+)\]", tv)
        assert(res)
        hash_name = res.group(3).replace("-", "")
        hash_module = load_hash_by_name(hash_name)
        continue

    if hasattr(tv, "p"):
        modulus = tv.p
        generator = tv.g
        suborder = tv.q
        continue

    hash_obj = hash_module.new(tv.msg)
   
    comps = [bytes_to_long(x) for x in (tv.y, generator, modulus, suborder)]
    key = DSA.construct(comps, False) # type: ignore
    verifier = DSS.new(key, 'fips-186-3')

    def positive_test(self, verifier=verifier, hash_obj=hash_obj, signature=tv.r+tv.s):
        verifier.verify(hash_obj, signature)

    def negative_test(self, verifier=verifier, hash_obj=hash_obj, signature=tv.r+tv.s):
        self.assertRaises(ValueError, verifier.verify, hash_obj, signature)

    if tv.result == 'p':
        setattr(FIPS_DSA_Tests_KAT, "test_verify_positive_%d" % idx, positive_test)
    else:
        setattr(FIPS_DSA_Tests_KAT, "test_verify_negative_%d" % idx, negative_test)


test_vectors_sign = load_tests(("Crypto", "SelfTest", "Signature", "test_vectors", "DSA"),
                               "FIPS_186_3_SigGen.txt",
                               "Signature Creation 186-3",
                               {})

for idx, tv in enumerate(test_vectors_sign):

    if isinstance(tv, str):
        res = re.match(r"\[mod = L=([0-9]+), N=([0-9]+), ([a-zA-Z0-9-]+)\]", tv)
        assert(res)
        hash_name = res.group(3).replace("-", "")
        hash_module = load_hash_by_name(hash_name)
        continue

    if hasattr(tv, "p"):
        modulus = tv.p
        generator = tv.g
        suborder = tv.q
        continue

    hash_obj = hash_module.new(tv.msg)
    comps_dsa = [bytes_to_long(x) for x in (tv.y, generator, modulus, suborder, tv.x)]
    key = DSA.construct(comps_dsa, False) # type: ignore
    signer = DSS.new(key, 'fips-186-3', randfunc=StrRNG(tv.k))

    def new_test(self, signer=signer, hash_obj=hash_obj, signature=tv.r+tv.s):
        self.assertEqual(signer.sign(hash_obj), signature)
    setattr(FIPS_DSA_Tests_KAT, "test_sign_%d" % idx, new_test)


class FIPS_ECDSA_Tests(unittest.TestCase):

    key_priv = ECC.generate(curve="P-256")
    key_pub = key_priv.public_key()

    def shortDescription(self):
        return "FIPS ECDSA Tests"

    def test_loopback(self):
        hashed_msg = SHA512.new(b"test")
        signer = DSS.new(self.key_priv, 'fips-186-3')
        signature = signer.sign(hashed_msg)

        verifier = DSS.new(self.key_pub, 'fips-186-3')
        verifier.verify(hashed_msg, signature)

    def test_negative_unapproved_hashes(self):
        """Verify that unapproved hashes are rejected"""

        from Crypto.Hash import SHA1

        self.description = "Unapproved hash (SHA-1) test"
        hash_obj = SHA1.new()
        signer = DSS.new(self.key_priv, 'fips-186-3')
        self.assertRaises(ValueError, signer.sign, hash_obj)
        self.assertRaises(ValueError, signer.verify, hash_obj, b"\x00" * 40)

    def test_sign_verify(self):
        """Verify public/private method"""

        self.description = "can_sign() test"
        signer = DSS.new(self.key_priv, 'fips-186-3')
        self.failUnless(signer.can_sign())

        signer = DSS.new(self.key_pub, 'fips-186-3')
        self.failIf(signer.can_sign())

    def test_negative_unknown_modes_encodings(self):
        """Verify that unknown modes/encodings are rejected"""

        self.description = "Unknown mode test"
        self.assertRaises(ValueError, DSS.new, self.key_priv, 'fips-186-0')

        self.description = "Unknown encoding test"
        self.assertRaises(ValueError, DSS.new, self.key_priv, 'fips-186-3', 'xml')

    def test_asn1_encoding(self):
        """Verify ASN.1 encoding"""

        self.description = "ASN.1 encoding test"
        hash_obj = SHA256.new()
        signer = DSS.new(self.key_priv, 'fips-186-3', 'der')
        signature = signer.sign(hash_obj)

        # Verify that output looks like a DER SEQUENCE
        self.assertEqual(bord(signature[0]), 48)
        signer.verify(hash_obj, signature)

        # Verify that ASN.1 parsing fails as expected
        signature = bchr(7) + signature[1:]
        self.assertRaises(ValueError, signer.verify, hash_obj, signature)


class FIPS_ECDSA_Tests_KAT(unittest.TestCase):
    pass


test_vectors_verify = load_tests(("Crypto", "SelfTest", "Signature", "test_vectors", "ECDSA"),
                                 "SigVer.rsp",
                                 "ECDSA Signature Verification 186-3",
                                 {'result': lambda x: x,
                                  'qx': lambda x: int(x, 16),
                                  'qy': lambda x: int(x, 16),
                                  })

for idx, tv in enumerate(test_vectors_verify):

    if isinstance(tv, str):
        res = re.match(r"\[P-256,(SHA-[0-9]+)\]", tv)
        assert res
        hash_name = res.group(1).replace("-", "")
        hash_module = load_hash_by_name(hash_name)
        continue

    hash_obj = hash_module.new(tv.msg)
    ecc_key = ECC.construct(curve="P-256", point_x=tv.qx, point_y=tv.qy)
    verifier = DSS.new(ecc_key, 'fips-186-3')

    def positive_test(self, verifier=verifier, hash_obj=hash_obj, signature=tv.r+tv.s):
        verifier.verify(hash_obj, signature)

    def negative_test(self, verifier=verifier, hash_obj=hash_obj, signature=tv.r+tv.s):
        self.assertRaises(ValueError, verifier.verify, hash_obj, signature)

    if tv.result.startswith('p'):
        setattr(FIPS_ECDSA_Tests_KAT, "test_verify_positive_%d" % idx, positive_test)
    else:
        setattr(FIPS_ECDSA_Tests_KAT, "test_verify_negative_%d" % idx, negative_test)


test_vectors_sign = load_tests(("Crypto", "SelfTest", "Signature", "test_vectors", "ECDSA"),
                               "SigGen.txt",
                               "ECDSA Signature Verification 186-3",
                               {'d': lambda x: int(x, 16)})

for idx, tv in enumerate(test_vectors_sign):

    if isinstance(tv, str):
        res = re.match(r"\[P-256,(SHA-[0-9]+)\]", tv)
        assert res
        hash_name = res.group(1).replace("-", "")
        hash_module = load_hash_by_name(hash_name)
        continue

    hash_obj = hash_module.new(tv.msg)
    ecc_key = ECC.construct(curve="P-256", d=tv.d)
    signer = DSS.new(ecc_key, 'fips-186-3', randfunc=StrRNG(tv.k))

    def new_test(self, signer=signer, hash_obj=hash_obj, signature=tv.r+tv.s):
        self.assertEqual(signer.sign(hash_obj), signature)
    setattr(FIPS_ECDSA_Tests_KAT, "test_sign_%d" % idx, new_test)


class Det_DSA_Tests(unittest.TestCase):
    """Tests from rfc6979"""