#!/usr/bin/env python

import mobius

g_count = 0
g_errors = 0
g_unittests = 0

class unittest (object):

  def __init__ (self, title):
    self.__title = title
    self.__count = 0
    self.__errors = 0
    global g_unittests
    g_unittests += 1

  def end (self):
    dots = 0 if len (self.__title) > 60 else 60 - len (self.__title)
    print "\033[1;39m%s\033[0m %s" % (self.__title, '.' * dots),

    if (self.__errors == 0):
      print "\033[1;32mOK",
    else:
      print "\033[1;31mERROR",
  
    print "\033[0m(%d/%d)" % (self.__count - self.__errors, self.__count)

  def assert_equal (self, v1, v2):
    global g_count, g_errors
    self.__count += 1
    g_count += 1
    
    if v1 != v2:
      print "\033[1;31mERROR:\033[0m %s == %s failed" % (v1, v2)
      self.__errors += 1
      g_errors += 1

  def assert_true (self, v1):
    global g_count, g_errors
    self.__count += 1
    g_count += 1
    
    if not v1:
      print "\033[1;31mERROR:\033[0m %s is false" % v1
      self.__errors += 1
      g_errors += 1

  @staticmethod
  def final_summary ():
    global g_unittests, g_count, g_errors
    print
    print "unittests:", g_unittests
    print "tests    :", g_count
    print "errors   :", g_errors

def string_to_hex (s):
  return ''.join ('%02x' % ord (c) for c in s)

def get_hex_digest (h):
  return string_to_hex (h.digest)

def unittest_core_application ():
  app = mobius.core.application ()

  test = unittest ("mobius.core.application")
  test.assert_equal (app.name, "Mobius Forensic Toolkit")
  test.assert_true (len (app.version) > 0)
  test.assert_true (len (app.title) > 0)
  test.assert_true (len (app.copyright) > 0)

  test.end ()

def unittest_crypt_cipher_aes ():

  test = unittest ("mobius.crypt.cipher_aes")
  
  plaintext = '\x6b\xc1\xbe\xe2\x2e\x40\x9f\x96\xe9\x3d\x7e\x11\x73\x93\x17\x2a\xae\x2d\x8a\x57\x1e\x03\xac\x9c\x9e\xb7\x6f\xac\x45\xaf\x8e\x51\x30\xc8\x1c\x46\xa3\x5c\xe4\x11\xe5\xfb\xc1\x19\x1a\x0a\x52\xef\xf6\x9f\x24\x45\xdf\x4f\x9b\x17\xad\x2b\x41\x7b\xe6\x6c\x37\x10'

  iv = '\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'

  # AES128 - mode ECB
  ciphertext = '\x3a\xd7\x7b\xb4\x0d\x7a\x36\x60\xa8\x9e\xca\xf3\x24\x66\xef\x97\xf5\xd3\xd5\x85\x03\xb9\x69\x9d\xe7\x85\x89\x5a\x96\xfd\xba\xaf\x43\xb1\xcd\x7f\x59\x8e\xce\x23\x88\x1b\x00\xe3\xed\x03\x06\x88\x7b\x0c\x78\x5e\x27\xe8\xad\x3f\x82\x23\x20\x71\x04\x72\x5d\xd4'
  
  aes1 = mobius.crypt.cipher_aes ('\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c')

  test.assert_equal (aes1.encrypt (plaintext), ciphertext)
  test.assert_equal (aes1.decrypt (ciphertext), plaintext)

  # AES192 - mode ECB
  ciphertext = '\xbd\x33\x4f\x1d\x6e\x45\xf2\x5f\xf7\x12\xa2\x14\x57\x1f\xa5\xcc\x97\x41\x04\x84\x6d\x0a\xd3\xad\x77\x34\xec\xb3\xec\xee\x4e\xef\xef\x7a\xfd\x22\x70\xe2\xe6\x0a\xdc\xe0\xba\x2f\xac\xe6\x44\x4e\x9a\x4b\x41\xba\x73\x8d\x6c\x72\xfb\x16\x69\x16\x03\xc1\x8e\x0e'

  aes2 = mobius.crypt.cipher_aes ('\x8e\x73\xb0\xf7\xda\x0e\x64\x52\xc8\x10\xf3\x2b\x80\x90\x79\xe5\x62\xf8\xea\xd2\x52\x2c\x6b\x7b')

  test.assert_equal (aes2.encrypt (plaintext), ciphertext)
  test.assert_equal (aes2.decrypt (ciphertext), plaintext)

  # AES256 - mode ECB
  ciphertext = '\xf3\xee\xd1\xbd\xb5\xd2\xa0\x3c\x06\x4b\x5a\x7e\x3d\xb1\x81\xf8\x59\x1c\xcb\x10\xd4\x10\xed\x26\xdc\x5b\xa7\x4a\x31\x36\x28\x70\xb6\xed\x21\xb9\x9c\xa6\xf4\xf9\xf1\x53\xe7\xb1\xbe\xaf\xed\x1d\x23\x30\x4b\x7a\x39\xf9\xf3\xff\x06\x7d\x8d\x8f\x9e\x24\xec\xc7'
  
  aes3 = mobius.crypt.cipher_aes ('\x60\x3d\xeb\x10\x15\xca\x71\xbe\x2b\x73\xae\xf0\x85\x7d\x77\x81\x1f\x35\x2c\x07\x3b\x61\x08\xd7\x2d\x98\x10\xa3\x09\x14\xdf\xf4')

  test.assert_equal (aes3.encrypt (plaintext), ciphertext)
  test.assert_equal (aes3.decrypt (ciphertext), plaintext)
  
  # AES128 - mode CBC
  ciphertext = '\x76\x49\xab\xac\x81\x19\xb2\x46\xce\xe9\x8e\x9b\x12\xe9\x19\x7d\x50\x86\xcb\x9b\x50\x72\x19\xee\x95\xdb\x11\x3a\x91\x76\x78\xb2\x73\xbe\xd6\xb8\xe3\xc1\x74\x3b\x71\x16\xe6\x9e\x22\x22\x95\x16\x3f\xf1\xca\xa1\x68\x1f\xac\x09\x12\x0e\xca\x30\x75\x86\xe1\xa7'

  aes4a = mobius.crypt.cipher_aes ('\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c', 'cbc', iv)

  test.assert_equal (aes4a.encrypt (plaintext), ciphertext)
  
  aes4b = mobius.crypt.cipher_aes ('\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c', 'cbc', iv)

  test.assert_equal (aes4b.decrypt (ciphertext), plaintext)
  
  test.end ()

def unittest_crypt_cipher_des ():

  test = unittest ("mobius.crypt.cipher_des")
  c = mobius.crypt.cipher_des ("\x13\x34\x57\x79\x9b\xbc\xdf\xf1")

  text = "\x01\x23\x45\x67\x89\xab\xcd\xef"
  test.assert_equal (string_to_hex (text), "0123456789abcdef")
  
  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "85e813540f0ab405")
  
  text2 = c.decrypt (text2)
  test.assert_equal (string_to_hex (text2), "0123456789abcdef")
  
  # LM password "ABCD123"
  c = mobius.crypt.cipher_des ("ABCD123")

  text = "KGS!@#$%"
  test.assert_equal (string_to_hex (text), "4b47532140232425")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "6f87cd328120cc55")
  
  text2 = c.decrypt (text2)
  test.assert_equal (string_to_hex (text2), "4b47532140232425")
  
  # DES - mode CBC
  c = mobius.crypt.cipher_des ("ABCDE12", 'cbc', '\x00\x00\x00\x00\x00\x00\x00\x00')
  
  text = "KGS!@#$%"
  test.assert_equal (string_to_hex (text), "4b47532140232425")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "722ac01404a75156")
  
  c.reset ()
  text2 = c.decrypt (text2)
  test.assert_equal (string_to_hex (text2), "4b47532140232425")
  
  test.end ()

def unittest_crypt_cipher_rc4 ():

  # key="Key", Plaintext="Plaintext"
  test = unittest ("mobius.crypt.cipher_rc4")
  c = mobius.crypt.cipher_rc4 ("Key")

  text = "Plaintext"
  test.assert_equal (text, "Plaintext")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "bbf316e8d940af0ad3")

  c.reset ()
  text2 = c.decrypt (text2)
  test.assert_equal (text, text2)

  # key="Secret", Plaintext="Attack at dawn"
  c = mobius.crypt.cipher_rc4 ("Secret")

  text = "Attack at dawn"
  test.assert_equal (text, "Attack at dawn")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "45a01f645fc35b383552544b9bf5")

  c.reset ()
  text2 = c.decrypt (text2)
  test.assert_equal (text, text2)

  test.end ()

def unittest_crypt_cipher_rot13 ():

  # nowhere <-> abjurer
  test = unittest ("mobius.crypt.cipher_rot13")
  c = mobius.crypt.cipher_rot13 ()

  text = "nowhere"
  test.assert_equal (text, "nowhere")

  text2 = c.encrypt (text)
  test.assert_equal (text2, "abjurer")

  # ARES <-> NERF
  c = mobius.crypt.cipher_rot13 ()

  text = "ARES"
  test.assert_equal (text, "ARES")

  text2 = c.encrypt (text)
  test.assert_equal (text2, "NERF")

  test.end ()

def unittest_crypt_cipher_zip ():

  # key="Key", Plaintext="Plaintext"
  test = unittest ("mobius.crypt.cipher_zip")
  c = mobius.crypt.cipher_zip ("Key")

  text = "Plaintext"
  test.assert_equal (text, "Plaintext")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "fe1995e4fe54a8c6f3")

  c.reset ()
  text2 = c.decrypt (text2)
  test.assert_equal (text, text2)

  # key="Secret", Plaintext="Attack at dawn"
  c = mobius.crypt.cipher_zip ("Secret")

  text = "Attack at dawn"
  test.assert_equal (text, "Attack at dawn")

  text2 = c.encrypt (text)
  test.assert_equal (string_to_hex (text2), "7595da02f5ec5c2c78755fd4069e")

  c.reset ()
  text2 = c.decrypt (text2)
  test.assert_equal (text, text2)

  test.end ()

def unittest_crypt_hash_adler32 ():
  test = unittest ("mobius.crypt.hash_adler32")

  h1 = mobius.crypt.hash_adler32 ()
  test.assert_equal (get_hex_digest (h1), "00000001")

  h2 = mobius.crypt.hash_adler32 ()
  h2.update ('abc')
  test.assert_equal (get_hex_digest (h2), "024d0127")

  h2.update ('d')
  test.assert_equal (get_hex_digest (h2), "03d8018b")

  h3 = mobius.crypt.hash_adler32 ()
  h3.update ('abcd')
  test.assert_equal (get_hex_digest (h3), get_hex_digest (h2))

  test.end ()

def unittest_crypt_hash_crc32 ():
  test = unittest ("mobius.crypt.hash_crc32")

  h1 = mobius.crypt.hash_crc32 ()
  test.assert_equal (get_hex_digest (h1), "00000000")

  h2 = mobius.crypt.hash_crc32 ()
  h2.update ('abc')
  test.assert_equal (get_hex_digest (h2), "352441c2")

  h2.update ('d')
  test.assert_equal (get_hex_digest (h2), "ed82cd11")

  h3 = mobius.crypt.hash_crc32 ()
  h3.update ('abcd')
  test.assert_equal (get_hex_digest (h3), get_hex_digest (h2))

  test.end ()

def unittest_crypt_hash_md5 ():
  test = unittest ("mobius.crypt.hash_md5")

  # RFC 1321 - section A.5
  h1 = mobius.crypt.hash_md5 ();
  test.assert_equal (get_hex_digest (h1), "d41d8cd98f00b204e9800998ecf8427e");

  h2 = mobius.crypt.hash_md5 ();
  h2.update ("a");
  test.assert_equal (get_hex_digest (h2), "0cc175b9c0f1b6a831c399e269772661");

  h3 = mobius.crypt.hash_md5 ();
  h3.update ("abc");
  test.assert_equal (get_hex_digest (h3), "900150983cd24fb0d6963f7d28e17f72");

  h4 = mobius.crypt.hash_md5 ();
  h4.update ("message digest");
  test.assert_equal (get_hex_digest (h4), "f96b697d7cb7938d525a2f31aaf161d0");

  h5 = mobius.crypt.hash_md5 ();
  h5.update ("abcdefghijklmnopqrstuvwxyz");
  test.assert_equal (get_hex_digest (h5), "c3fcd3d76192e4007dfb496cca67e13b");
  
  h6 = mobius.crypt.hash_md5 ();
  h6.update ("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789");
  test.assert_equal (get_hex_digest (h6), "d174ab98d277d9f5a5611c2c9f419d9f");
  
  h7 = mobius.crypt.hash_md5 ();
  h7.update ("12345678901234567890123456789012345678901234567890123456789012345678901234567890");
  test.assert_equal (get_hex_digest (h7), "57edf4a22be3c955ac49da2e2107b67a");
  
  h8 = mobius.crypt.hash_md5 ();
  h8.update ('\0' * 1000);
  test.assert_equal (get_hex_digest (h8), "ede3d3b685b4e137ba4cb2521329a75e");

  test.end ()

def unittest_crypt_hash_sha1 ():
  test = unittest ("mobius.crypt.hash_sha1")

  # RFC 3174 - section 7.3
  h1 = mobius.crypt.hash_sha1 ()
  test.assert_equal (get_hex_digest (h1), "da39a3ee5e6b4b0d3255bfef95601890afd80709")

  h2 = mobius.crypt.hash_sha1 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "a9993e364706816aba3e25717850c26c9cd0d89d")

  h3 = mobius.crypt.hash_sha1 ()
  h3.update ("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq")
  test.assert_equal (get_hex_digest (h3), "84983e441c3bd26ebaae4aa1f95129e5e54670f1")

  h4 = mobius.crypt.hash_sha1 ()
  
  for i in xrange (1000000):
    h4.update ('a')

  test.assert_equal (get_hex_digest (h4), "34aa973cd4c4daa4f61eeb2bdbad27316534016f");

  h5 = mobius.crypt.hash_sha1 ();
  
  for i in xrange (20):
    h5.update ("01234567012345670123456701234567")

  test.assert_equal (get_hex_digest (h5), "dea356a2cddd90c7a7ecedc5ebb563934f460452")
  
  test.end ()

def unittest_crypt_hash_sha2_224 ():
  test = unittest ("mobius.crypt.hash_sha2_224")

  h1 = mobius.crypt.hash_sha2_224 ()
  test.assert_equal (get_hex_digest (h1), "d14a028c2a3a2bc9476102bb288234c415a2b01f828ea62ac5b3e42f")

  h2 = mobius.crypt.hash_sha2_224 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "23097d223405d8228642a477bda255b32aadbce4bda0b3f7e36c9da7")

  h3 = mobius.crypt.hash_sha2_224 ()
  h3.update ("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq")
  test.assert_equal (get_hex_digest (h3), "75388b16512776cc5dba5da1fd890150b0c6455cb4f58b1952522525")
  
  test.end ()

def unittest_crypt_hash_sha2_256 ():
  test = unittest ("mobius.crypt.hash_sha2_256")

  h1 = mobius.crypt.hash_sha2_256 ()
  test.assert_equal (get_hex_digest (h1), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")

  h2 = mobius.crypt.hash_sha2_256 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")

  h3 = mobius.crypt.hash_sha2_256 ()
  h3.update ("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq")
  test.assert_equal (get_hex_digest (h3), "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1")
  
  test.end ()

def unittest_crypt_hash_sha2_384 ():
  test = unittest ("mobius.crypt.hash_sha2_384")

  h1 = mobius.crypt.hash_sha2_384 ()
  test.assert_equal (get_hex_digest (h1), "38b060a751ac96384cd9327eb1b1e36a21fdb71114be07434c0cc7bf63f6e1da274edebfe76f65fbd51ad2f14898b95b")

  h2 = mobius.crypt.hash_sha2_384 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "cb00753f45a35e8bb5a03d699ac65007272c32ab0eded1631a8b605a43ff5bed8086072ba1e7cc2358baeca134c825a7")

  h3 = mobius.crypt.hash_sha2_384 ()
  h3.update ("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu")
  test.assert_equal (get_hex_digest (h3), "09330c33f71147e83d192fc782cd1b4753111b173b3b05d22fa08086e3b0f712fcc7c71a557e2db966c3e9fa91746039")
  
  test.end ()

def unittest_crypt_hash_sha2_512 ():
  test = unittest ("mobius.crypt.hash_sha2_512")

  h1 = mobius.crypt.hash_sha2_512 ()
  test.assert_equal (get_hex_digest (h1), "cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e")

  h2 = mobius.crypt.hash_sha2_512 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f")

  h3 = mobius.crypt.hash_sha2_512 ()
  h3.update ("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu")
  test.assert_equal (get_hex_digest (h3), "8e959b75dae313da8cf4f72814fc143f8f7779c6eb9f7fa17299aeadb6889018501d289e4900f7e4331b99dec4b5433ac7d329eeb6dd26545e96e55b874be909")
  
  test.end ()

def unittest_crypt_hash_sha2_512_224 ():
  test = unittest ("mobius.crypt.hash_sha2_512_224")

  h1 = mobius.crypt.hash_sha2_512_224 ()
  test.assert_equal (get_hex_digest (h1), "6ed0dd02806fa89e25de060c19d3ac86cabb87d6a0ddd05c333b84f4")

  h2 = mobius.crypt.hash_sha2_512_224 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "4634270f707b6a54daae7530460842e20e37ed265ceee9a43e8924aa")

  h3 = mobius.crypt.hash_sha2_512_224 ()
  h3.update ("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu")
  test.assert_equal (get_hex_digest (h3), "23fec5bb94d60b23308192640b0c453335d664734fe40e7268674af9")
  
  test.end ()

def unittest_crypt_hash_sha2_512_256 ():
  test = unittest ("mobius.crypt.hash_sha2_512_256")

  h1 = mobius.crypt.hash_sha2_512_256 ()
  test.assert_equal (get_hex_digest (h1), "c672b8d1ef56ed28ab87c3622c5114069bdd3ad7b8f9737498d0c01ecef0967a")

  h2 = mobius.crypt.hash_sha2_512_256 ()
  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "53048e2681941ef99b2e29b76b4c7dabe4c2d0c634fc6d46e0e2f13107e7af23")

  h3 = mobius.crypt.hash_sha2_512_256 ()
  h3.update ("abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu")
  test.assert_equal (get_hex_digest (h3), "3928e184fb8690f840da3988121d31be65cb9d3ef83ee6146feac861e19b563a")
  
  test.end ()

def unittest_crypt_hash_zip ():
  h1 = mobius.crypt.hash_zip ()
  h2 = mobius.crypt.hash_zip ()
  h3 = mobius.crypt.hash_zip ()

  test = unittest ("mobius.crypt.hash_zip")
  test.assert_equal (get_hex_digest (h1), "123456782345678934567890")
  test.assert_equal (get_hex_digest (h2), "123456782345678934567890")
  test.assert_equal (get_hex_digest (h3), "123456782345678934567890")

  h2.update ("abc")
  test.assert_equal (get_hex_digest (h2), "5dd2af4d589d03b43cf5ffa4")

  h2.update ('d')
  test.assert_equal (get_hex_digest (h2), "42ef4ac38d167254428e6d93")
  
  h3.update ("abcd")
  test.assert_equal (get_hex_digest (h3), get_hex_digest (h2))

  test.end ()

def unittest_io_file ():

  test = unittest ("mobius.io.file")
  f = mobius.io.file ("file:///dev/zero")
  reader = f.new_reader ()

  # initial values
  test.assert_equal (reader.size, 0)
  test.assert_equal (reader.tell (), 0)
  
  # read bytes
  data = reader.read (16)
  test.assert_equal (len (data), 16)
  test.assert_equal (data[0], '\0')
  test.assert_equal (data[15], '\0')
  
  # values
  test.assert_equal (reader.tell (), 16)

  # seek
  reader.seek (100)
  test.assert_equal (reader.tell (), 100)

  # read more bytes
  data = reader.read (16)
  test.assert_equal (len (data), 16)
  test.assert_equal (data[0], '\0')
  test.assert_equal (data[15], '\0')

  # values
  test.assert_equal (reader.tell (), 116)

  test.end ()

# main code
unittest_core_application ()
unittest_crypt_cipher_aes ()
unittest_crypt_cipher_des ()
unittest_crypt_cipher_rc4 ()
unittest_crypt_cipher_rot13 ()
unittest_crypt_cipher_zip ()
unittest_crypt_hash_adler32 ()
unittest_crypt_hash_crc32 ()
unittest_crypt_hash_md5 ()
unittest_crypt_hash_sha1 ()
unittest_crypt_hash_sha2_224 ()
unittest_crypt_hash_sha2_256 ()
unittest_crypt_hash_sha2_384 ()
unittest_crypt_hash_sha2_512 ()
unittest_crypt_hash_sha2_512_224 ()
unittest_crypt_hash_sha2_512_256 ()
unittest_crypt_hash_zip ()
unittest_io_file ()
unittest.final_summary ()
