#!/usr/bin/env python
"""
Firmware Bundle-and-Protect Tool
"""
import argparse
import shutil
import struct
import json
import sys
import zlib
import pdb
import binascii
from hashlib import sha256
from cStringIO import StringIO
from intelhex import IntelHex
import random, os, struct
#from Crypto.Cipher import AES
from simon import SimonCipher

# Swaps bytes in a list, see StackOverflow
def swap_order(d, wsz=4, gsz=2 ):
        return "".join(["".join([m[i:i+gsz] for i in range(wsz-gsz,-gsz,-gsz)]) for m in [d[i:i+wsz] for i in range(0,len(d),wsz)]])

def complement(input_int):
    input_str = (bin(input_int)[2:]).zfill(8)
    result = ''
    for i in input_str:
        if i == '1':
            result = result + '0'
        else:
            result = result + '1'
    return int(result, 2) + 1

def compute_checksum(intelhex_str):
	try:
            checksum_hex = bytearray(intelhex_str.decode('hex'))
	except:
	    import pdb; pdb.set_trace()
        checksum_val = sum(checksum_hex) & 0xFF
	two_complement  = complement(checksum_val)
	checksum_bytes = hex(two_complement)
	checksum_str = checksum_bytes[2:].zfill(2)[-2:]
	return checksum_str

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Firmware Update Tool')

    parser.add_argument("--infile",
                        help="Path to the firmware image to protect.",
                        required=True)
    parser.add_argument("--outfile", help="Filename for the output firmware.",
                        required=True)
    parser.add_argument("--version", help="Version number of this firmware.",
                        required=True)
    parser.add_argument("--message", help="Release message for this firmware.",
                        required=True)
    args = parser.parse_args()

    # Parse Intel hex file.
    firmware = IntelHex(args.infile)


    # Get version and size.
    firmware_size = firmware.maxaddr() + 1
    version = int(args.version)

    # Add release message to end of hex (null-terminated).
    sio = StringIO()
    firmware.putsz(firmware_size, (args.message + '\0'))
    firmware.write_hex_file(sio)
    hex_data = sio.getvalue()

    # Hex data is in IntelHex format, see wikipedia article on formatting

    # Parse the key from secret_configure_output.txt
    f = open('secret_configure_output.txt', 'r') 
    secret_config_file_txt = f.read()
    secret_config_json = json.loads(secret_config_file_txt)
    key_string = secret_config_json['SIMONKEY']
    key = int(key_string, 16)

    #create simon cipher, key goes in here in 0xhex, or int('hex-string-here',16)
    my_simon = SimonCipher(key,key_size=128, block_size=64)

    # split each line
    byline = hex_data.splitlines()

    # remove first two lines from encryption because there is no data there
    encrypted_byline = byline[0:2]

    # only use data lines
    byline_data = byline[2:-1]
    import pdb; pdb.set_trace()

    # for hash output
    hash_input_temp = ''
    hash_input = []
    i = 0
    total_flash_data = ''
    first_address = byline_data[0][3:7]
    # for each line with data
    for line in byline_data:
        # read numBytes in line
        numBytes = int(line[1:3],16);
        print numBytes
        # get data hex values
        data = line[9:9+numBytes*2]
        total_flash_data = total_flash_data + data

    if len(total_flash_data) % 32 != 0:
        total_flash_data.ljust((len(total_flash_data)/32 * 32) + 32, '0')
    total_flash_list = [total_flash_data[i:i + 32] for i in range(0, len(total_flash_data), 32)]
    import pdb; pdb.set_trace()

    for j, data in enumerate(total_flash_list):    
        if len(data) < 32:
            data = data.ljust(32, '0')
        elif len(data) > 32:
            print "Not a valid data length" 
            sys.exit()
        # encrypt the 16B using simon
        data1 = swap_order(data[0:16],wsz=16,gsz=2)
        data2 = swap_order(data[16:],wsz=16,gsz=2)
        encrypted1 = hex(my_simon.encrypt(int(data1,16)))[2:-1]
        encrypted2 = hex(my_simon.encrypt(int(data2,16)))[2:-1]
        swapped1 = swap_order(encrypted1.zfill(16),wsz = 16, gsz=2)
        swapped2 = swap_order(encrypted2.zfill(16), wsz = 16, gsz = 2)
        encryptedData = swapped1 + swapped2
        current_address = hex(int(first_address, 16) + 16 * j)[2:]

        # create new line with encrypted data, convert to hex string
        encryptedline = ':10' + current_address + '00' +  encryptedData
        checksum_str = compute_checksum(encryptedline[1:])
        encryptedline = encryptedline + checksum_str

        hash_input_temp = hash_input_temp + encryptedData
        if len(hash_input_temp) == 512:
            hash_input.append(hash_input_temp)
            hash_input_temp = ''

        encrypted_byline.append(encryptedline)

    hash_input.append(hash_input_temp)

    # add the last line back
    encrypted_byline.extend(byline[-1:])

    # go back to intel_hex SIO() format for hex_data
    encrypted_hex_data = "\n".join(encrypted_byline)
    # #HASH for pages
    tags = []
    hash_locals = []

    for input_data in hash_input:
         hash_local = sha256(input_data.decode('hex')).hexdigest().zfill(64)
         hash_locals.append(hash_local)
         first_q = swap_order(hash_local[0:16], wsz=16, gsz=2)
         second_q = swap_order(hash_local[16:32], wsz=16, gsz=2)
         third_q = swap_order(hash_local[32:48], wsz=16, gsz=2)
         fourth_q = swap_order(hash_local[48:], wsz=16, gsz=2)
         fq = hex(my_simon.encrypt(int(first_q, 16)))[2:-1].zfill(16)
         sq = hex(my_simon.encrypt(int(second_q, 16)))[2:-1].zfill(16)
         tq = hex(my_simon.encrypt(int(third_q, 16)))[2:-1].zfill(16)
         foq = hex(my_simon.encrypt(int(fourth_q, 16)))[2:-1].zfill(16)
         fs = swap_order(fq, wsz=16, gsz=2)
         ss = swap_order(sq, wsz=16, gsz=2)
         ts = swap_order(tq, wsz=16, gsz=2)
         fos = swap_order(foq, wsz=16, gsz=2)
         tags.append(fs + ss + ts + fos)

    # Sign Result
    # Save as Version-Bytes
    version_hash_input = (hex(version)[2:]).zfill(4) + (hex(key)[2:-1]).zfill(32)
    import pdb; pdb.set_trace()
    version_hash = sha256(version_hash_input.decode('hex')).hexdigest().zfill(64)
   
    # Encode the data as json and write to outfile.
    data = {
        'firmware_size' : firmware_size,
        'version_hash' : version_hash,
	'version' : version,
        'hex_data' : encrypted_hex_data,
        'tags' : tags
    }

    with open(args.outfile, 'wb+') as outfile:
        data = json.dumps(data)
        data = zlib.compress(data)
        outfile.write(data)

a7e396811d043e4807a18c2ac3d2923b8d5088ab
#!/usr/bin/env python
"""
Firmware Bundle-and-Protect Tool
"""
import argparse
import shutil
import struct
import json
import sys
import zlib
import pdb
import binascii
from hashlib import sha256
from cStringIO import StringIO
from intelhex import IntelHex
import random, os, struct
#from Crypto.Cipher import AES
from simon import SimonCipher

# Swaps bytes in a list, see StackOverflow
def swap_order(d, wsz=4, gsz=2 ):
        return "".join(["".join([m[i:i+gsz] for i in range(wsz-gsz,-gsz,-gsz)]) for m in [d[i:i+wsz] for i in range(0,len(d),wsz)]])

def complement(input_int):
    input_str = (bin(input_int)[2:]).zfill(8)
    result = ''
    for i in input_str:
        if i == '1':
            result = result + '0'
        else:
            result = result + '1'
    return int(result, 2) + 1

def compute_checksum(intelhex_str):
	try:
            checksum_hex = bytearray(intelhex_str.decode('hex'))
	except:
	    import pdb; pdb.set_trace()
        checksum_val = sum(checksum_hex) & 0xFF
	two_complement  = complement(checksum_val)
	checksum_bytes = hex(two_complement)
	checksum_str = checksum_bytes[2:].zfill(2)[-2:]
	return checksum_str

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Firmware Update Tool')

    parser.add_argument("--infile",
                        help="Path to the firmware image to protect.",
                        required=True)
    parser.add_argument("--outfile", help="Filename for the output firmware.",
                        required=True)
    parser.add_argument("--version", help="Version number of this firmware.",
                        required=True)
    parser.add_argument("--message", help="Release message for this firmware.",
                        required=True)
    args = parser.parse_args()

    # Parse Intel hex file.
    firmware = IntelHex(args.infile)


    # Get version and size.
    firmware_size = firmware.maxaddr() + 1
    version = int(args.version)

    # Add release message to end of hex (null-terminated).
    sio = StringIO()
    firmware.putsz(firmware_size, (args.message + '\0'))
    firmware.write_hex_file(sio)
    hex_data = sio.getvalue()

    # Hex data is in IntelHex format, see wikipedia article on formatting

    # Parse the key from secret_configure_output.txt
    f = open('secret_configure_output.txt', 'r') 
    secret_config_file_txt = f.read()
    secret_config_json = json.loads(secret_config_file_txt)
    key_string = secret_config_json['SIMONKEY']
    key = int(key_string, 16)

    #create simon cipher, key goes in here in 0xhex, or int('hex-string-here',16)
    my_simon = SimonCipher(key,key_size=128, block_size=64)

    # split each line
    byline = hex_data.splitlines()

    # remove first two lines from encryption because there is no data there
    encrypted_byline = byline[0:2]

    # only use data lines
    byline_data = byline[2:-1]
    import pdb; pdb.set_trace()

    # for hash output
    hash_input_temp = ''
    hash_input = []
    i = 0
    total_flash_data = ''
    first_address = byline_data[0][3:7]
    # for each line with data
    for line in byline_data:
        # read numBytes in line
        numBytes = int(line[1:3],16);
        print numBytes
        # get data hex values
        data = line[9:9+numBytes*2]
        total_flash_data = total_flash_data + data

    if len(total_flash_data) % 32 != 0:
        total_flash_data.ljust((len(total_flash_data)/32 * 32) + 32, '0')
    total_flash_list = [total_flash_data[i:i + 32] for i in range(0, len(total_flash_data), 32)]
    import pdb; pdb.set_trace()

    for j, data in enumerate(total_flash_list):    
        if len(data) < 32:
            data = data.ljust(32, '0')
        elif len(data) > 32:
            print "Not a valid data length" 
            sys.exit()
        # encrypt the 16B using simon
        data1 = swap_order(data[0:16],wsz=16,gsz=2)
        data2 = swap_order(data[16:],wsz=16,gsz=2)
        encrypted1 = hex(my_simon.encrypt(int(data1,16)))[2:-1]
        encrypted2 = hex(my_simon.encrypt(int(data2,16)))[2:-1]
        swapped1 = swap_order(encrypted1.zfill(16),wsz = 16, gsz=2)
        swapped2 = swap_order(encrypted2.zfill(16), wsz = 16, gsz = 2)
        encryptedData = swapped1 + swapped2
        current_address = hex(int(first_address, 16) + 16 * j)[2:]

        # create new line with encrypted data, convert to hex string
        encryptedline = ':10' + current_address + '00' +  encryptedData
        checksum_str = compute_checksum(encryptedline[1:])
        encryptedline = encryptedline + checksum_str

        hash_input_temp = hash_input_temp + encryptedData
        if len(hash_input_temp) == 512:
            hash_input.append(hash_input_temp)
            hash_input_temp = ''

        encrypted_byline.append(encryptedline)

    hash_input.append(hash_input_temp)

    # add the last line back
    encrypted_byline.extend(byline[-1:])

    # go back to intel_hex SIO() format for hex_data
    encrypted_hex_data = "\n".join(encrypted_byline)
    # #HASH for pages
    tags = []
    hash_locals = []

    for input_data in hash_input:
         hash_local = sha256(input_data.decode('hex')).hexdigest().zfill(64)
         hash_locals.append(hash_local)
         first_q = swap_order(hash_local[0:16], wsz=16, gsz=2)
         second_q = swap_order(hash_local[16:32], wsz=16, gsz=2)
         third_q = swap_order(hash_local[32:48], wsz=16, gsz=2)
         fourth_q = swap_order(hash_local[48:], wsz=16, gsz=2)
         fq = hex(my_simon.encrypt(int(first_q, 16)))[2:-1].zfill(16)
         sq = hex(my_simon.encrypt(int(second_q, 16)))[2:-1].zfill(16)
         tq = hex(my_simon.encrypt(int(third_q, 16)))[2:-1].zfill(16)
         foq = hex(my_simon.encrypt(int(fourth_q, 16)))[2:-1].zfill(16)
         fs = swap_order(fq, wsz=16, gsz=2)
         ss = swap_order(sq, wsz=16, gsz=2)
         ts = swap_order(tq, wsz=16, gsz=2)
         fos = swap_order(foq, wsz=16, gsz=2)
         tags.append(fs + ss + ts + fos)

    # Sign Result
    # Save as Version-Bytes
    version_hash_input = (hex(version)[2:]).zfill(4) + (hex(key)[2:-1]).zfill(32)
    import pdb; pdb.set_trace()
    version_hash = sha256(version_hash_input.decode('hex')).hexdigest().zfill(64)
   
    # Encode the data as json and write to outfile.
    data = {
        'firmware_size' : firmware_size,
        'version_hash' : version_hash,
	'version' : version,
        'hex_data' : encrypted_hex_data,
        'tags' : tags
    }

    with open(args.outfile, 'wb+') as outfile:
        data = json.dumps(data)
        data = zlib.compress(data)
        outfile.write(data)
