#!/usr/bin/env python3

"""
Created by: Wonder Gumise and Ahmad MK Nasser
Date created: 12/04/2019
Version: 1.0.1
Date last modified 15/04/2019
Python Version: 3.3
py-crypto Version: 2.6.1
"""

from Crypto.Cipher import AES
import drbg
import CMAC
import os
import sys
import hexfile
import binascii

def GetBlockListOfSampleAddresses(size, sample_size, seed, start_addr): #Block sampling for each address
    '''Generate of `size` randomly sampled address blocks using the given `seed`
    
    :param size: The size of the data image in bytes
    :param sample_size: The percentage of the data image being sampled
    :param seed: A 32 byte seed of type bytes for use by the AES CTR DRBG generator
    :param start_addr: The start address of the hex image
    '''
    cipher = drbg.CTRDRBG('aes128' , seed)  #AES CTR DRBG object for use in generating the addresses
    rnd_addr_list = []
    mask = (size >> 16) - 1  #create a mask to later limit the address range
    sample_size = int(size*(sample_size/100))  #convert sample size as a percentage to bytes
    block_size = int(sample_size/16)  #convert sample size to blocks
    for _ in range(block_size):
        buffer = cipher._generate(3)
        #create a range limited hex address from the generated bytes corresponding to the size of the hex image
        address = ((int(buffer[2]) & mask)<<16 | int(buffer[1])<<8 | (int(buffer[0]) & 0xF0))
        address = address + start_addr

        for j in range(16):
            rnd_addr_list.append(address+j)
        #print(rnd_addr_list)
    return rnd_addr_list #sorted(rnd_addr_list)

def GetOnePerBlockListOfSampleAddresses(size, seed, start_addr): #Byte sampling for each address
    '''Generate of one randomly sampled address for every block in the hex image using the given `seed`
    
    :param size: The size of the data image in bytes
    :param seed: A 32 byte seed of type bytes for use by the AES CTR DRBG generator
    :param start_addr: The start address of the hex file
    '''
    cipher = drbg.CTRDRBG('aes128' , seed)  #AES CTR DRBG object for use in generating the addresses
    rnd_addr_list = []
    size = (int(size/16))  #convert size to blocks
    for _ in range(int(size/2)):
        offset_b = cipher._generate(1)
        ##offset = int.from_bytes(offset_b, byteorder='big')
        #print(offset) #debug only
        address = start_addr + int(offset_b[0] & 0x0F)
        rnd_addr_list.append(address)
        start_addr += 16

        address = start_addr + ((int(offset_b[0]) & 0xF0) >> 4)
        rnd_addr_list.append(address)
        start_addr += 16
    return rnd_addr_list

def GetTwoPerBlockListOfSampleAddresses(size, seed, start_addr): #Byte sampling for each address
    '''Generate of two randomly sampled addresses for every block in the hex image using the given `seed`
    
    :param size: The size of the data image in bytes
    :param seed: A 32 byte seed of type bytes for use by the AES CTR DRBG generator
    :param start_addr: The start address of the hex file
    '''
    cipher = drbg.CTRDRBG('aes128' , seed)  #AES CTR DRBG object for use in generating the addresses
    rnd_addr_list = []
    size = (int(size/16))
    #fetch size from hexfile
    for _ in range(size):
        offset_b = cipher._generate(1)
        ##offset = int.from_bytes(offset_b, byteorder='big')  #debug only
        #print(offset) #debug only
        address = start_addr + int(offset_b[0] & 0x0F)
        rnd_addr_list.append(address)
        address = start_addr + ((int(offset_b[0]) & 0xF0) >> 4)
        rnd_addr_list.append(address)
        start_addr += 16
    return rnd_addr_list


if __name__ == "__main__":
    print('        Verification phase:\n')
    f_seed = open("Saved_Parameters\seed.txt", 'r')
    seed = f_seed.read()  #retrieve saved seed
    f_seed.close()
    seed = str.encode(seed)  #convert hex string to hex bytes 
    seed = binascii.unhexlify(seed)  #convert hex bytes to bytes

    f_key = open("Saved_Parameters\key.txt", 'r')
    key = f_key.read()  #retrieve saved key
    f_key.close()
    key = str.encode(key) #convert hex string to hex bytes
    key = binascii.unhexlify(key)   #convert hex bytes to bytes

    # fileDir = os.path.dirname(os.path.realpath('__file__'))  #get path of running script
    filename = input('Enter filename/relative path of hexfile to be verified: ')  #get name of hexfile to be verified
    f_hex = hexfile.load(filename)
    info = f_hex.segments
    info = str(info)  #convert hex segments info list to a string
    info = info[:-2]  #remove '>]' from the end of segments info string
    info = [x for x in info.split() if x.startswith('0x')]  #extract hex number from the info string
    start_addr = int(info[0], 16)  #convert the hex number to type int
    # print(start_addr)  #debug only

    f_scheme = open("Saved_Parameters\scheme.txt", 'r')
    scheme = f_scheme.read()  #retrieve saved scheme
    f_scheme.close()
    if scheme == 'rbs':
        f_sample_size = open("Saved_Parameters\sample_size.txt", 'r')
        sample_size = f_sample_size.read()  #retrieve saved key
        f_sample_size.close()
        sample_size = int(sample_size)
        s_addr = GetBlockListOfSampleAddresses(f_hex.size, sample_size, seed, start_addr)
    elif scheme == 'pbs1':
        s_addr = GetOnePerBlockListOfSampleAddresses(f_hex.size, seed, start_addr)
    elif scheme == 'pbs2':
        s_addr = GetTwoPerBlockListOfSampleAddresses(f_hex.size, seed, start_addr)
    else:
        print('Error: Invalid sample boot method')
        print('    Valid options are: rbs, pbs1 and pbs2')
        sys.exit(0)

    s_data = []
    for addr in s_addr:
        data = f_hex[addr]  #get data from sampled addresses
        s_data.append(data)
    # print(s_data)
    s_data = bytes(s_data)

    cmac = CMAC.new(key,ciphermod=AES)
    cmac.update(s_data)  #calculate CMAC for sampled data

    f_tag = open("Saved_Parameters\cmac.txt", 'r')
    tag = f_tag.read()  #retrieve saved tag to be used for verification
    f_tag.close()
    if (cmac.hexdigest() == tag):  #verify if tag is the same
        print('No modification detected')
    else:
        print('Modification detected')

sys.exit(0)