ims-gen

#! /usr/bin/env python

#
# Copyright (c) 2015 Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from this
# software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# Generate <count> legal IMS values (i.e. 280-bit random numbers having a
# Hamming weight of 128 over the LSB 32 bytes)

from __future__ import print_function
from struct import pack_into, unpack_from
import sys
import os
import binascii
import hashlib
import random
import argparse
from util import error, print_to_error

# Program return values
PROGRAM_SUCCESS = 0
PROGRAM_WARNINGS = 1
PROGRAM_ERROR = 2

# output formats
FORMAT_HUMAN = 1
FORMAT_TOSHIBA_EFUSE = 2
FORMAT_binascii = 3

# IMS is 35 bytes long, but boot ROM only cares about the first 32 bytes
IMS_BITS = 280
IMS_BYTES = (IMS_BITS + 7)//8
IMS_DWORDS = (IMS_BYTES + 3)//4
IMS_MEANINGFUL_LENGTH = 32
IMS_MEANINGFUL_BITS = IMS_MEANINGFUL_LENGTH * 8

CMS_BITS = 192
CMS_BYTES = (CMS_BITS + 7)//8
CMS_DWORDS = (CMS_BYTES + 3)//4

HASH_DIGEST_SIZE = 32


def binlify_lsb_msb(binarray):
    """ Convert a binary array into an ASCII binary string, LSb...MSb """
    asciibin = ""
    for byte in binarray:
        for bit_pos in range(0, 8, 1):
            if (byte & (1 << bit_pos)):
                asciibin += "1"
            else:
                asciibin += "0"
    return asciibin


def binlify_msb_lsb(binarray):
    """ Convert a binary array into an ASCII binary string, MSb...LSb """
    asciibin = ""
    for index in range(len(binarray) - 1, -1, -1):
        byte = binarray[index]
        for bit_pos in range(7, -1, -1):
            if (byte & (1 << bit_pos)):
                asciibin += "1"
            else:
                asciibin += "0"
    return asciibin


def permute_cms(cms_buf):
    """ Permute the CMS and return whether or not it has a valid hamming weight

    Returns true if valid, false otherwise
    """
    temp = bytearray(4)

    # Permute the CMS, one 32-bit word at a time
    while True:
        for i in range(0, CMS_DWORDS):
            # Replace 1 32-bit word at a time with randomness
            offset = i * 4
            length = min(4, IMS_BYTES - offset)
            pack_into("<L", temp, 0, random.getrandbits(32))

            cms_buf[offset:(offset+length)] = temp[:length]

            # Calculate the Hamming weight
            hamming_wt = 0
            for j in range(0, CMS_BYTES):
                x = cms_buf[j]
                hamming_wt += bin(x).count('1')

            # If the hamming weight is correct, use this IMS
            if hamming_wt == (CMS_BITS / 2):
                return


def print_cms(cms_buf, output_format, outfile):
    """ Print the CMS """
    if output_format == FORMAT_HUMAN:
        # Print the CMS in human-friendly form
        print("CMS {0:s}.{1:s}.{2:s}.{3:s}.{4:s}.{5:s}".
              format(binascii.hexlify(cms_buf[0:4]),
                     binascii.hexlify(cms_buf[4:8]),
                     binascii.hexlify(cms_buf[8:12]),
                     binascii.hexlify(cms_buf[12:16]),
                     binascii.hexlify(cms_buf[16:20]),
                     binascii.hexlify(cms_buf[20:24])), file=outfile)
    elif output_format == FORMAT_TOSHIBA_EFUSE:
        # Print the CMS in Toshiba .efz form
        cms = unpack_from("<LLLLLL", cms_buf)
        print("CMS[191:0] = {0:08x}_{1:08x}_{2:08x}_{3:08x}_{4:08x}_{5:08x}".
              format(cms[5], cms[4], cms[3], cms[2], cms[1], cms[0]),
              file=outfile)
    elif output_format == FORMAT_binascii:
        # Print the CMS in ascii binary MSB..LSB
        print(binlify_msb_lsb(cms_buf[0:CMS_BYTES]), file=outfile)


def permute_ims(ims_buf):
    """ Permute the IMS until it reaches a Hamming weight of 128

    Complicating matters is the fact that while the IMS is 280 bits (37 bytes),
    the bootloader only cares about the LSB 256 bits (32 bytes), so the hamming
    weight for those 32 bytes must equal 128. The high-order bits are random,
    but don't contribute to the hamming weight.

    ims_buf: a bytearray(IMS_BYTES)
    """
    temp = bytearray(4)

    # Permute the IMS, one 32-bit word at a time
    while True:
        for i in range(0, IMS_DWORDS):
            # Replace 1 32-bit word at a time with randomness
            offset = i * 4
            length = min(4, IMS_BYTES - offset)
            pack_into("<L", temp, 0, random.getrandbits(32))

            ims_buf[offset:(offset+length)] = temp[:length]

            # Calculate the Hamming weight
            hamming_wt = 0
            for j in range(0, IMS_MEANINGFUL_LENGTH):
                x = ims_buf[j]
                hamming_wt += bin(x).count('1')

            # If the hamming weight is correct, use this IMS
            if hamming_wt == (IMS_MEANINGFUL_BITS / 2):
                return


def print_ims(ims_buf, ep_uid, ep_uid_es3, output_format, outfile):
    """ Print the IMS """
    if output_format == FORMAT_HUMAN:
        # Print the IMS in human-friendly form
        print("IMS        {0:s}.{1:s}.{2:s}.{3:s}.{4:s}.{5:s}."
              "{6:s}.{7:s}.{8:s}".
              format(binascii.hexlify(ims_buf[0:4]),
                     binascii.hexlify(ims_buf[4:8]),
                     binascii.hexlify(ims_buf[8:12]),
                     binascii.hexlify(ims_buf[12:16]),
                     binascii.hexlify(ims_buf[16:20]),
                     binascii.hexlify(ims_buf[20:24]),
                     binascii.hexlify(ims_buf[24:28]),
                     binascii.hexlify(ims_buf[28:32]),
                     binascii.hexlify(ims_buf[32:35])), file=outfile)
        print("EP_UID     {0:s}.{1:s}".
              format(binascii.hexlify(ep_uid[0:4]),
                     binascii.hexlify(ep_uid[4:8])), file=outfile)
        print("EP_UID_ES3 {0:s}.{1:s}".
              format(binascii.hexlify(ep_uid_es3[0:4]),
                     binascii.hexlify(ep_uid_es3[4:8])), file=outfile)
        m = hashlib.sha256()
        m.update(ims_buf[0:IMS_BYTES])
        sha256 = m.digest()
        print("IMS SHA256 {0:s}.{1:s}.{2:s}.{3:s}.{4:s}.{5:s}."
              "{6:s}.{7:s}".
              format(binascii.hexlify(sha256[0:4]),
                     binascii.hexlify(sha256[4:8]),
                     binascii.hexlify(sha256[8:12]),
                     binascii.hexlify(sha256[12:16]),
                     binascii.hexlify(sha256[16:20]),
                     binascii.hexlify(sha256[20:24]),
                     binascii.hexlify(sha256[24:28]),
                     binascii.hexlify(sha256[28:32])), file=outfile)
    elif output_format == FORMAT_TOSHIBA_EFUSE:
        # Print the IMS in toshiba .efz form
        ims = unpack_from("<LLLLLLLLL", ims_buf)
        print("IMS[279:0] = {0:06x}_{1:08x}_{2:08x}_{3:08x}_{4:08x}_"
              "{5:08x}_{6:08x}_{7:08x}_{8:08x}".
              format(ims[8], ims[7], ims[6], ims[5], ims[4],
                     ims[3], ims[2], ims[1], ims[0]), file=outfile)
    elif output_format == FORMAT_binascii:
        # Print the IMS in ascii binary MSB..LSB
        print(binlify_msb_lsb(ims_buf[0:IMS_BYTES]), file=outfile)


def endpoint_unique_id(ims_buf):
    """ Calculate the Endpoint Unique ID from the IMS value

    ims_buf is the packed binary buffer containing the 280-bit IMS value
    """
    # Establish the default (i.e., no endpoint ID)
    endpoint_id = 0L

    # Compute Endpoint Unique ID from the IMS
    #
    # The algorithm used to calculate Endpoint Unique ID is:
    # Y1 = sha256(IMS[0:15] xor copy(0x3d, 16))
    # Z0 = sha256(Y1 || copy(0x01, 32))
    # EP_UID[0:7] = sha256(Z0)[0:7]
    #
    temp = bytearray(4)

    # grab IMS 4 bytes at a time and feed that to hash_update
    m = hashlib.sha256()
    for i in range(0, 4):
        up32 = unpack_from("<L", ims_buf, i * 4)
        pack_into("<L", temp, 0, up32[0] ^ 0x3d3d3d3d)
        m.update(temp)
    y1 = m.digest()

    m2 = hashlib.sha256()
    m2.update(y1)
    pack_into("<L", temp, 0, 0x01010101)
    for i in range(0, 8):
        m2.update(temp)
    z0 = m2.digest()

    m3 = hashlib.sha256()
    m3.update(z0)
    ep_uid = m3.digest()

    endpoint_id = bytearray(8)
    endpoint_id = ep_uid[0:8]
    return endpoint_id


def endpoint_unique_id_es3(ims_buf):
    """ Calculate the ES3-compatible Endpoint Unique ID from the IMS value

    ims_buf is the packed binary buffer containing the 280-bit IMS value
    """
    # Establish the default (i.e., no endpoint ID)
    endpoint_id = 0L

    # Generating EP_UID for ES3 compatible
    #
    # There is a bug in ES3 boot ROM, so the EP_UID is calculated by using
    # fewer IMS bytes:
    # 	Y1 = sha256((IMS[0] xor 0x3d) || (IMS[4] xor 0x3d) ||
    #       (IMS[8] xor 0x3d) || (IMS[12] xor 0x3d))
    # Z0 = sha256(Y1 || copy(0x01, 8))
    # EP_UID[0:7] = sha256(Z0)[0:7]
    #
    # Only 4 bytes of the IMS are used for this derivation to ensure, in an
    # information-theoretic sense, that no information about the rest of the
    # IMS is exposed through the publicly-exposed EP_UID value or its
    # derivation process.
    #
    temp = bytearray(1)

    # grab IMS 4 bytes at a time and feed that to hash_update
    m = hashlib.sha256()
    for i in range(0, 4):
        temp[0] = ims_buf[i * 4] ^ 0x3d
        m.update(temp)
    y1 = m.digest()

    m2 = hashlib.sha256()
    m2.update(y1)
    temp[0] = 0x01
    for i in range(0, 8):
        m2.update(temp)
    z0 = m2.digest()

    m3 = hashlib.sha256()
    m3.update(z0)
    ep_uid = m3.digest()

    endpoint_id = bytearray(8)
    endpoint_id = ep_uid[0:8]
    return endpoint_id


def main():
    """ Application for creating plausible stand-in IMS values """
    global ims, cms

    parser = argparse.ArgumentParser()

    parser.add_argument("--ims", "-i",
                        type=int,
                        help="How many IMS values to generate")

    parser.add_argument("--cms", "-c",
                        type=int,
                        help="How many CMS values to generate")

    parser.add_argument("--toshiba-efuse", "--efuse", "-e",
                        action='store_true',
                        help="Generates the IMS/CMS values in Toshiba "
                             "'.efz' format")

    parser.add_argument("--binascii", "-b",
                        action='store_true',
                        help="Generates the IMS/CMS values in ascii binary"
                             "format (MSB...LSB)")

    args = parser.parse_args()
    # Sanity-check the args
    if (args.ims == 0) and (args.cms == 0):
        error("You must specify at least one of --ims or --cms")
        sys.exit(PROGRAM_ERROR)

    if args.ims:
        num_generated = 0
        with open("ims_efuse.txt", "w") as ims_efuse:
            with open("ims_binascii.txt", "w") as ims_binascii:
                ims_buf = bytearray(IMS_DWORDS * 4)
                print_to_error("Generate", args.ims, "IMS values:")
                while num_generated < args.ims:
                    permute_ims(ims_buf)
                    num_generated += 1
                    print_ims(ims_buf,
                              endpoint_unique_id(ims_buf),
                              endpoint_unique_id_es3(ims_buf),
                              FORMAT_HUMAN, sys.stdout)
                    print_ims(ims_buf,
                              endpoint_unique_id(ims_buf),
                              endpoint_unique_id_es3(ims_buf),
                              FORMAT_TOSHIBA_EFUSE, ims_efuse)
                    print_ims(ims_buf,
                              endpoint_unique_id(ims_buf),
                              endpoint_unique_id_es3(ims_buf),
                              FORMAT_binascii, ims_binascii)
        # Remove unwanted formats (can't selectively choose "with")
        if not args.toshiba_efuse:
            os.remove("ims_efuse.txt")
        if not args.binascii:
            os.remove("ims_binascii.txt")

    if args.cms:
        num_generated = 0
        with open("cms_efuse.txt", "w") as cms_efuse:
            with open("cms_binascii.txt", "w") as cms_binascii:
                cms_buf = bytearray(CMS_DWORDS * 4)
                print_to_error("Generate", args.cms, "CMS values:")
                while num_generated < args.cms:
                    permute_cms(cms_buf)
                    num_generated += 1
                    print_cms(cms_buf, FORMAT_HUMAN, sys.stdout)
                    print_cms(cms_buf, FORMAT_TOSHIBA_EFUSE, cms_efuse)
                    print_cms(cms_buf, FORMAT_binascii, cms_binascii)
        # Remove unwanted formats (can't selectively choose "with")
        if not args.toshiba_efuse:
            os.remove("cms_efuse.txt")
        if not args.binascii:
            os.remove("cms_binascii.txt")


## Launch main
#
if __name__ == '__main__':
    main()