1  import hashlib
  2  import random
  3  import RNS
  4  import os
  5  import time
  6  import unittest
  7  
  8  class TestSHA256(unittest.TestCase):
  9      def setUp(self):
 10          self.f = RNS.Cryptography.sha256
 11  
 12      def test_empty(self):
 13          self.assertEqual(
 14              self.f(''.encode("utf-8")),
 15              bytes.fromhex("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"))
 16  
 17      def test_less_than_block_length(self):
 18          self.assertEqual(
 19              self.f('abc'.encode("utf-8")),
 20              bytes.fromhex("ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"))
 21  
 22      def test_block_length(self):
 23          self.assertEqual(
 24              self.f('a'.encode("utf-8")*64),
 25              bytes.fromhex("ffe054fe7ae0cb6dc65c3af9b61d5209f439851db43d0ba5997337df154668eb"))
 26  
 27      def test_several_blocks(self):
 28          self.assertEqual(
 29              self.f('a'.encode("utf-8")*1000000),
 30              bytes.fromhex("cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0"))
 31  
 32      def test_random_blocks(self):
 33          max_rounds = 5000
 34  
 35          b = 0
 36          i = 0
 37          ok = True
 38          start = time.time()
 39          print("")
 40          while ok and i < max_rounds:
 41              i += 1
 42              rlen = random.randint(0, 1024*16)
 43              rdat = os.urandom(rlen)
 44              b += rlen
 45              msg = rdat
 46              ok = RNS.Cryptography.sha256(msg) == hashlib.sha256(msg).digest()
 47              # t = RNS.Cryptography.sha256(msg)
 48              # t = hashlib.sha256(msg).digest()
 49              if (i%1000 == 0):
 50                  gbytes = round(b/1000000000,3)
 51                  mbps = round((b*8/1000000)/(time.time()-start), 2)
 52                  print(str(i)+" rounds OK, total data: "+str(gbytes)+"GB, "+str(mbps)+"mbps")
 53  
 54          if not ok:
 55              print("Failed at round "+str(i))
 56          else:
 57              print("SHA-256 test OK")
 58  
 59          self.assertEqual(ok, True)
 60  
 61  
 62  class TestSHA512(unittest.TestCase):
 63      def setUp(self):
 64          self.f = RNS.Cryptography.sha512
 65  
 66      def test_empty(self):
 67          self.assertEqual(
 68              self.f(''.encode("utf-8")),
 69              bytes.fromhex(
 70                  'cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce'+
 71                  '47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e'))
 72  
 73      def test_less_than_block_length(self):
 74          self.assertEqual(self.f('abc'.encode("utf-8")),
 75              bytes.fromhex(
 76                  'ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a'+
 77                  '2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f'))
 78  
 79      def test_block_length(self):
 80          self.assertEqual(self.f('a'.encode("utf-8")*128),
 81              bytes.fromhex(
 82                  'b73d1929aa615934e61a871596b3f3b33359f42b8175602e89f7e06e5f658a24'+
 83                  '3667807ed300314b95cacdd579f3e33abdfbe351909519a846d465c59582f321'))
 84  
 85      def test_several_blocks(self):
 86          self.assertEqual(self.f('a'.encode("utf-8")*1000000),
 87              bytes.fromhex(
 88                  'e718483d0ce769644e2e42c7bc15b4638e1f98b13b2044285632a803afa973eb'+
 89                  'de0ff244877ea60a4cb0432ce577c31beb009c5c2c49aa2e4eadb217ad8cc09b'))
 90  
 91      def test_random_blocks(self):
 92          max_rounds = 5000
 93  
 94          b = 0
 95          i = 0
 96          ok = True
 97          start = time.time()
 98          print("")
 99          while ok and i < max_rounds:
100              i += 1
101              rlen = random.randint(0, 1024*16)
102              rdat = os.urandom(rlen)
103              b += rlen
104              msg = rdat
105              ok = RNS.Cryptography.sha512(msg) == hashlib.sha512(msg).digest()
106              # t = RNS.Cryptography.sha512(msg)
107              # t = hashlib.sha512(msg).digest()
108              if (i%1000 == 0):
109                  gbytes = round(b/1000000000,3)
110                  mbps = round((b*8/1000000)/(time.time()-start), 2)
111                  print(str(i)+" rounds OK, total data: "+str(gbytes)+"GB, "+str(mbps)+"mbps")
112  
113          if not ok:
114              print("Failed at round "+str(i))
115          else:
116              print("SHA-512 test OK")
117  
118          self.assertEqual(ok, True)
119  
120  
121  if __name__ == '__main__':
122      unittest.main(verbosity=2)