demo_otp_with_image.py

import argparse
import os

BMP_HEADER_LENGTH = 54
BMP_HEADER = bytes.fromhex('424d36000c000000000036000000280000000002000000020000010018000000000000000c0074120000741200000000000000000000')

first_encrypted_path = 'first_encrypted_image.bmp'
second_encrypted_path = 'second_encrypted_image.bmp'
mix_encrypted_path = 'mix_encrypted_image.bmp'

def demo(first_image_path, second_image_path):
    with open(first_image_path, 'rb') as rf1:
        first_image_binary = rf1.read()

    with open(second_image_path, 'rb') as rf2:
        second_image_binary = rf2.read()

    if len(first_image_binary) != len(second_image_binary):
        raise ValueError('The length of second image must equals the length of first image.')
    
    pad = os.urandom(len(first_image_binary) - BMP_HEADER_LENGTH)

    first_encrypted_binary = otp_encrypt(first_image_binary, pad)
    with (open(first_encrypted_path, 'wb')) as wf1:
        wf1.write(first_encrypted_binary)
    print(f'Encrypted version of "{first_image_path}" is written to: {first_encrypted_path}')

    second_encrypted_binary = otp_encrypt(second_image_binary, pad)
    with (open(second_encrypted_path, 'wb')) as wf2:
        wf2.write(second_encrypted_binary)
    print(f'Encrypted version of "{second_image_path}" is written to: {second_encrypted_path}')

    mix_encrypted_data = xor_data(first_encrypted_binary[BMP_HEADER_LENGTH:], \
        second_encrypted_binary[BMP_HEADER_LENGTH:])
    mixed_encrypted_binary = BMP_HEADER + mix_encrypted_data
    with (open(mix_encrypted_path, 'wb')) as wf3:
        wf3.write(mixed_encrypted_binary)
    print(f'Mixed of both encrypted files is written to: {mix_encrypted_path}')

def otp_encrypt(image_binary, pad):
    image_data = image_binary[BMP_HEADER_LENGTH:]
    encrypted_image_data = xor_data(image_data, pad)
    encrypted_image_binary = BMP_HEADER + encrypted_image_data

    return encrypted_image_binary

def xor_data(data, pad):
    return bytes(a ^ b for (a, b) in zip(data, pad))

def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--first_image', help='Path to the first test image', required=True)
    parser.add_argument('--second_image', help='Path to the second test image', required=True)
    args=parser.parse_args()

    try:
        demo(args.first_image, args.second_image)
    except Exception as e:
        print('An error occurred.')
        print(str(e))


if __name__ == '__main__':
    main()