#encryption modes for babyblock

#I did not include the decryption, as these
#are part of the homework exercises.

import babyblock

#The plaintext is a sequence of bytes.  If the length is odd, it will be padded with
#one 0 byte so that there is a whole number of 16-bit blocks. The key as usual
#is a 16-bit integer. The number of rounds for the block cipher is fixed at 6.
def encrypt_ecb(key,pt):
    ct=[]
    if len(pt)%2==1:
        pt=pt+[0]
    for j in range(0,len(pt),2):
        block=256*pt[j]+pt[j+1]
        ctblock=babyblock.encrypt(key,block)
        ct.extend([ctblock/256,ctblock%256])
    return ct

#CBC encryption.  The intialization vector
#is sent as two bytes prepended to the  ciphertext
def encrypt_cbc(key,pt,iv):
    ct=[iv/256,iv%256]
    if len(pt)%2==1:
        pt=pt+[0]
    for j in range(0,len(pt),2):
        block=256*pt[j]+pt[j+1]
        newblock=(256*ct[j]+ct[j+1])^block
        ctblock=babyblock.encrypt(key,newblock)
        ct.extend([ctblock/256,ctblock%256])
    return ct


#CTR encrytion. The counter is sent as two bytes prepended to
#the ciphertext.
def encrypt_ctr(key,pt,counter):
    ct=[counter/256,counter%256]
    if len(pt)%2==1:
        pt=pt+[0]
    for j in range(0,len(pt),2):
        mask=babyblock.encrypt(key,counter+j+1,6)
        ctblock=mask^(256*pt[j]+pt[j+1])
        ct.extend([ctblock/256,ctblock%256])
    return ct