import sha3
import unittest


class TestSha3(unittest.TestCase):

    def test_sha3_224(self):
        print("Testing SHA3-224 for NIST vectors...")
        msg = []
        rslt_ex = []
        path = 'test_vectors/SHA3_224ShortMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        path = 'test_vectors/SHA3_224LongMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        for i in range(len(msg)):
            rslt_own = sha3.sha3_224(msg[i])
            self.assertEqual(rslt_own, rslt_ex[i], "Wrong hash! is: " + str(rslt_own
                ) + ", should be: " + rslt_ex[i] + ", input is: " + msg[i])

    def test_sha3_256(self):
        print("Testing SHA3-256 for NIST vectors...")
        msg = []
        rslt_ex = []
        path = 'test_vectors/SHA3_256ShortMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        path = 'test_vectors/SHA3_256LongMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        for i in range(len(msg)):
            rslt_own = sha3.sha3_256(msg[i])
            self.assertEqual(rslt_own, rslt_ex[i], "Wrong hash! is: " + str(rslt_own
                ) + ", should be: " + rslt_ex[i] + ", input is: " + msg[i])

    def test_sha3_384(self):
        print("Testing SHA3-384 for NIST vectors...")
        msg = []
        rslt_ex = []
        path = 'test_vectors/SHA3_384ShortMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        path = 'test_vectors/SHA3_384LongMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        for i in range(len(msg)):
            rslt_own = sha3.sha3_384(msg[i])
            self.assertEqual(rslt_own, rslt_ex[i], "Wrong hash! is: " + str(rslt_own
                ) + ", should be: " + rslt_ex[i] + ", input is: " + msg[i])

    def test_sha3_512(self):
        print("Testing SHA3-512 for NIST vectors...")
        msg = []
        rslt_ex = []
        path = 'test_vectors/SHA3_512ShortMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        path = 'test_vectors/SHA3_512LongMsg.txt'
        with open(path, 'r') as infile:
            for line in infile:
                if line.startswith("Msg"):
                    if line[6:-1] == "00":
                        msg.append("")
                    else:
                        msg.append(line[6:-1])  # comma on the end prevents the double spacing from printing a file line
                elif line.startswith("MD"):
                    rslt_ex.append(line[5:-1])
            infile.close()
        for i in range(len(msg)):
            rslt_own = sha3.sha3_512(msg[i])
            self.assertEqual(rslt_own, rslt_ex[i], "Wrong hash! is: " + str(rslt_own
                ) + ", should be: " + rslt_ex[i] + ", input is: " + msg[i])


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