#!/usr/bin/env python

import mobius.api

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_01 = unittest ("mobius.core.application")
  test_01.assert_equal (app.name, "Mobius Forensic Toolkit")
  test_01.assert_true (len (app.version) > 0)
  test_01.assert_true (len (app.title) > 0)
  test_01.assert_true (len (app.copyright) > 0)

  test_01.end ()

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

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

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

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

  test_01.end ()

def unittest_crypt_cipher_rc4 ():

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

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

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

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

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

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

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

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

  test_01.end ()

def unittest_crypt_cipher_rot13 ():

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

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

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

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

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

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

  test_01.end ()

def unittest_crypt_cipher_zip ():

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

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

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

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

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

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

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

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

  test_01.end ()

def unittest_crypt_cipher_des ():

  test_01 = 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_01.assert_equal (string_to_hex (text), "0123456789abcdef")
  
  text2 = c.encrypt (text)
  test_01.assert_equal (string_to_hex (text2), "85e813540f0ab405")
  
  text2 = c.decrypt (text2)
  test_01.assert_equal (string_to_hex (text2), "0123456789abcdef")
  
  # LM password "ABCD123"
  c = mobius.crypt.cipher_des ("ABCD123")

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

  text2 = c.encrypt (text)
  test_01.assert_equal (string_to_hex (text2), "6f87cd328120cc55")
  
  text2 = c.decrypt (text2)
  test_01.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_01.assert_equal (string_to_hex (text), "4b47532140232425")

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

def unittest_io_file ():

  test = unittest ("mobius.api.file")
  f = mobius.api.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_hash_zip ()
unittest_crypt_cipher_des ()
unittest_crypt_cipher_rc4 ()
unittest_crypt_cipher_rot13 ()
unittest_crypt_cipher_zip ()
unittest_io_file ()
unittest.final_summary ()
