#encrypt bitmapped images
#This uses the crypto function modules
#from the pycrypto package.  You need software
#for viewing the images to see the result.

from Crypto.Cipher import DES
from Crypto.Cipher import AES

#The argument is the name of a .bmp file (with extension)
#in the course home directory. This is the AES version.

def bm_encrypt_ecb(filename):
    #modify this if you try it at home!'
    path = '/Users/straubin/teaching/crypto2017/'
    fullpathname=path+filename
    image_file=open(fullpathname,'r')
    filecontents=image_file.read()
    #for standard bmp files, the first 54 bytes contain
    #header information.  Let's strip this and save it.
    header=filecontents[:54]
    bitmap=filecontents[54:]
    #pad so that the bitmap length is divisible by 16.
    #We will add 16 full bytes of padding in the case that
    #it is divisible by 16.
    defect=16-len(bitmap)%16
    bitmap += ('0'*defect)
    #now encrypt
    obj=AES.new('sixteen byte key',AES.MODE_ECB)
    enc_bitmap=obj.encrypt(bitmap)
    enc_bitmap=enc_bitmap[:-defect]
    #now rejoin the header to the encrypted bitmap,
    #and write out to a new file
    enc_contents=header+enc_bitmap
    outfile=open(path+'enc_image.bmp','w')
    outfile.write(enc_contents)
    outfile.close()

#The argument is the name of a .bmp file (with extension)
#in the course home directory. This is the DES version
    
def bm_encrypt_ecb_des(filename):
    #modify this if you try it at home!'
    path = '/Users/straubin/teaching/crypto2017/'
    fullpathname=path+filename
    image_file=open(fullpathname,'r')
    filecontents=image_file.read()
    #for standard bmp files, the first 54 bytes contain
    #header information.  Let's strip this and save it.
    header=filecontents[:54]
    bitmap=filecontents[54:]
    #pad so that the bitmap length is divisible by 8.
    #We will add 8 full bytes of padding in the case that
    #it is divisible by 8.
    defect=8-len(bitmap)%8
    bitmap += ('0'*defect)
    #now encrypt
    obj=DES.new('8bytekey',DES.MODE_ECB)
    enc_bitmap=obj.encrypt(bitmap)
    enc_bitmap=enc_bitmap[:-defect]
    #now rejoin the header to the encrypted bitmap,
    #and write out to a new file
    enc_contents=header+enc_bitmap
    outfile=open(path+'enc_image.bmp','w')
    outfile.write(enc_contents)
    outfile.close()

#The CBC version is the same, but we have to supply an IV argument.
#Here the IV is all zeros
def bm_encrypt_cbc(filename):
    #modify this if you try it at home!'
    path = '/Users/straubin/teaching/crypto2017/'
    fullpathname=path+filename
    image_file=open(fullpathname,'r')
    filecontents=image_file.read()
    #for standard bmp files, the first 54 bytes contain
    #header information.  Let's strip this and save it.
    header=filecontents[:54]
    bitmap=filecontents[54:]
    #pad so that the bitmap length is divisible by 16.
    #We will add 16 full bytes of padding in the case that
    #it is divisible by 16.
    defect=16-len(bitmap)%16
    bitmap += ('0'*defect)
    #now encrypt
    obj=AES.new('sixteen byte key',AES.MODE_CBC,'0'*16)
    enc_bitmap=obj.encrypt(bitmap)
    enc_bitmap=enc_bitmap[:-defect]
    #now rejoin the header to the encrypted bitmap,
    #and write out to a new file
    enc_contents=header+enc_bitmap
    outfile=open(path+'enc_image.bmp','w')
    outfile.write(enc_contents)
    outfile.close()

#The DES version of the CBC demo.
def bm_encrypt_cbc_des(filename):
    #modify this if you try it at home!'
    path = '/Users/straubin/teaching/crypto2017/'
    fullpathname=path+filename
    image_file=open(fullpathname,'r')
    filecontents=image_file.read()
    #for standard bmp files, the first 54 bytes contain
    #header information.  Let's strip this and save it.
    header=filecontents[:54]
    bitmap=filecontents[54:]
    #pad so that the bitmap length is divisible by 8.
    #We will add 8 full bytes of padding in the case that
    #it is divisible by 8.
    defect=8-len(bitmap)%8
    bitmap += ('0'*defect)
    #now encrypt
    obj=DES.new('8bytekey',DES.MODE_CBC,'0'*8)
    enc_bitmap=obj.encrypt(bitmap)
    enc_bitmap=enc_bitmap[:-defect]
    #now rejoin the header to the encrypted bitmap,
    #and write out to a new file
    enc_contents=header+enc_bitmap
    outfile=open(path+'enc_image.bmp','w')
    outfile.write(enc_contents)
    outfile.close()