#!/usr/bin/env python

"""
    chacha8_test.py
     
    A Python demo/test program for ChaCha8.
     
    Larry Bugbee
    rev June 2010

"""

from chacha8 import ChaCha8
import sys
    
#---------------------------------------------------------------
# utility functions
    
def strnumlist(L):
    return ''.join([chr(x) for x in L])

def printlong(label, bs, bpl, dent, hex=True):
    tmpl = '%-' + str(dent) + 's'
    while bs:
        if hex:
            print (tmpl % label)[:dent], bs[:bpl].encode('hex')
        else:
            print (tmpl % label)[:dent], bs[:bpl]
        label = ''
        bs = bs[bpl:]

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
# Print usage

def print_usage():
    print """
        Usage:
            python %s [-sample] [-vectors] [-timing]
    """ % (sys.argv[0])

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
# Simple usage demo

def sample_usage_demo():
    print '-'*40
    print '  Simple usage demo:'
    key        = 'qwerty7890123456'     # 128-bit key (16 bytes)
    iv         = 'iv345678'
    plaintext  = 'Kilroy was here, there, and everywhere.'
    printlong('    key:',  key,  64, 16, hex=False)
    printlong('    iv:',   iv,   64, 16, hex=False)
    printlong('    plaintext: ', plaintext,  64, 16, hex=False)

    ciphertext = ChaCha8(key, iv).encrypt(plaintext)
    printlong('    ciphertext:', ciphertext, 64, 16)

    message = ChaCha8(key, iv).decrypt(ciphertext)
    printlong('    message:   ', message,    64, 16, hex=False)

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
# Published Vectors Test

def published_vectors_test():
    print '-'*40
    print '  Published Vectors Test:'

    # tuples of key, iv, plaintext, and expected ciphertext
    test_vectors = (  
      (
        # test_vector 1:  128-bit key (16 bytes) 
        "80000000000000000000000000000000", 
        "0000000000000000",            
        "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
        "beb1e81e0f747e43ee51922b3e87fb38d0163907b4ed49336032ab78b67c24579fe28f751bd3703e51d876c017faa43589e63593e03355a7d57b2366f30047c5",
      ),
      (
        # test_vector 2:  128-bit key (16 bytes) 
        "0053a6f94c9ff24598eb3e91e4378add",
        "0d74db42a91077de",
        "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
        "509b267e7266355fa2dc0a25c023fce47922d03dd9275423d7cb7118b2aedf220568854bf47920d6fc0fd10526cfe7f9de472835afc73c916b849e91eee1f529",
      ),
      (
        # test_vector 3:  256-bit key (32 bytes) 
        "0000200000000000000000000000000000000000000000000000000000000000",
        "0000000000000000",
        "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
        "653f4a18e3d27daf51f841a00b6c1a2bd2489852d4ae0711e1a4a32ad166fa6f881a2843238c7e17786ba5162bc019d573849c167668510ada2f62b4ff31ad04",
      ),
      (
        # test_vector 4:  256-bit key (32 bytes) 
        "0f62b5085bae0154a7fa4da0f34699ec3f92e5388bde3184d72a7dd02376c91c",
        "288ff65dc42b92f9",
        "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
        "db165814f66733b7a8e34d1ffc1234271256d3bf8d8da2166922e598acac70f412b3fe35a94190ad0ae2e8ec62134819ab61addcccfe99d867ca3d73183fa3fd",
      ),
    )
    
    for test_vector in test_vectors:
        key        = test_vector[0].decode('hex')
        iv         = test_vector[1].decode('hex')
        plaintext  = test_vector[2].decode('hex')
        expected   = test_vector[3].decode('hex')
        printlong('    key:',  key,  64, 14)
        printlong('    iv:',   iv,   64, 14)
        printlong('    plaintext: ', plaintext, 64, 14)
        printlong('    expected:  ', expected,  64, 14)
        
        result = ChaCha8(key, iv).encrypt(plaintext)
        printlong('    result:    ', result, 64, 14)
        
        if result != expected:
            print '      *** failed ***'
        else:
            print '          passed '

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
# Timing test

def timing_test():
    # caution!!!  this may take some time on slow processors
    ITER = 1000

    print '-'*40
    print '  Timing test (Python encrypting 512 bytes %d times):' % (ITER)
    import time
    
    key  = ('80000000000000000000000000000000' +
            '00000000000000000000000000000000'
           ).decode('hex')
    iv   = ('0000000000000000').decode('hex')
    data = ('00'*512).decode('hex')      # 512
    
    t0 = time.time()
    for i in xrange(ITER):
        cc8 = ChaCha8(key, iv)
        ciphertext = cc8.encrypt(data)
    t1 = time.time()
    print '    %5.3f secs' % (t1-t0)

#---------------------------------------------------------------

if __name__ == '__main__':

    if len(sys.argv) < 2:
        print_usage()
    else:
        for arg in sys.argv:
            if arg == '-sample' or arg == '-simple':
                sample_usage_demo()
            elif arg == '-vectors':
                published_vectors_test()
            elif arg == '-timing':
                timing_test()
        print '-'*40

#---------------------------------------------------------------
#---------------------------------------------------------------
#---------------------------------------------------------------