mirror of https://github.com/acidanthera/audk.git
505 lines
15 KiB
Python
505 lines
15 KiB
Python
#!/usr/bin/env python
|
|
# @ CommonUtility.py
|
|
# Common utility script
|
|
#
|
|
# Copyright (c) 2016 - 2021, Intel Corporation. All rights reserved.<BR>
|
|
# SPDX-License-Identifier: BSD-2-Clause-Patent
|
|
#
|
|
##
|
|
|
|
import os
|
|
import sys
|
|
import shutil
|
|
import subprocess
|
|
import string
|
|
from ctypes import ARRAY, c_char, c_uint16, c_uint32, \
|
|
c_uint8, Structure, sizeof
|
|
from importlib.machinery import SourceFileLoader
|
|
from SingleSign import single_sign_gen_pub_key
|
|
|
|
|
|
# Key types defined should match with cryptolib.h
|
|
PUB_KEY_TYPE = {
|
|
"RSA": 1,
|
|
"ECC": 2,
|
|
"DSA": 3,
|
|
}
|
|
|
|
# Signing type schemes defined should match with cryptolib.h
|
|
SIGN_TYPE_SCHEME = {
|
|
"RSA_PKCS1": 1,
|
|
"RSA_PSS": 2,
|
|
"ECC": 3,
|
|
"DSA": 4,
|
|
}
|
|
|
|
# Hash values defined should match with cryptolib.h
|
|
HASH_TYPE_VALUE = {
|
|
"SHA2_256": 1,
|
|
"SHA2_384": 2,
|
|
"SHA2_512": 3,
|
|
"SM3_256": 4,
|
|
}
|
|
|
|
# Hash values defined should match with cryptolib.h
|
|
HASH_VAL_STRING = dict(map(reversed, HASH_TYPE_VALUE.items()))
|
|
|
|
AUTH_TYPE_HASH_VALUE = {
|
|
"SHA2_256": 1,
|
|
"SHA2_384": 2,
|
|
"SHA2_512": 3,
|
|
"SM3_256": 4,
|
|
"RSA2048SHA256": 1,
|
|
"RSA3072SHA384": 2,
|
|
}
|
|
|
|
HASH_DIGEST_SIZE = {
|
|
"SHA2_256": 32,
|
|
"SHA2_384": 48,
|
|
"SHA2_512": 64,
|
|
"SM3_256": 32,
|
|
}
|
|
|
|
|
|
class PUB_KEY_HDR (Structure):
|
|
_pack_ = 1
|
|
_fields_ = [
|
|
('Identifier', ARRAY(c_char, 4)), # signature ('P', 'U', 'B', 'K')
|
|
('KeySize', c_uint16), # Length of Public Key
|
|
('KeyType', c_uint8), # RSA or ECC
|
|
('Reserved', ARRAY(c_uint8, 1)),
|
|
('KeyData', ARRAY(c_uint8, 0)),
|
|
]
|
|
|
|
def __init__(self):
|
|
self.Identifier = b'PUBK'
|
|
|
|
|
|
class SIGNATURE_HDR (Structure):
|
|
_pack_ = 1
|
|
_fields_ = [
|
|
('Identifier', ARRAY(c_char, 4)),
|
|
('SigSize', c_uint16),
|
|
('SigType', c_uint8),
|
|
('HashAlg', c_uint8),
|
|
('Signature', ARRAY(c_uint8, 0)),
|
|
]
|
|
|
|
def __init__(self):
|
|
self.Identifier = b'SIGN'
|
|
|
|
|
|
class LZ_HEADER(Structure):
|
|
_pack_ = 1
|
|
_fields_ = [
|
|
('signature', ARRAY(c_char, 4)),
|
|
('compressed_len', c_uint32),
|
|
('length', c_uint32),
|
|
('version', c_uint16),
|
|
('svn', c_uint8),
|
|
('attribute', c_uint8)
|
|
]
|
|
_compress_alg = {
|
|
b'LZDM': 'Dummy',
|
|
b'LZ4 ': 'Lz4',
|
|
b'LZMA': 'Lzma',
|
|
}
|
|
|
|
|
|
def print_bytes(data, indent=0, offset=0, show_ascii=False):
|
|
bytes_per_line = 16
|
|
printable = ' ' + string.ascii_letters + string.digits + string.punctuation
|
|
str_fmt = '{:s}{:04x}: {:%ds} {:s}' % (bytes_per_line * 3)
|
|
bytes_per_line
|
|
data_array = bytearray(data)
|
|
for idx in range(0, len(data_array), bytes_per_line):
|
|
hex_str = ' '.join(
|
|
'%02X' % val for val in data_array[idx:idx + bytes_per_line])
|
|
asc_str = ''.join('%c' % (val if (chr(val) in printable) else '.')
|
|
for val in data_array[idx:idx + bytes_per_line])
|
|
print(str_fmt.format(
|
|
indent * ' ',
|
|
offset + idx, hex_str,
|
|
' ' + asc_str if show_ascii else ''))
|
|
|
|
|
|
def get_bits_from_bytes(bytes, start, length):
|
|
if length == 0:
|
|
return 0
|
|
byte_start = (start) // 8
|
|
byte_end = (start + length - 1) // 8
|
|
bit_start = start & 7
|
|
mask = (1 << length) - 1
|
|
val = bytes_to_value(bytes[byte_start:byte_end + 1])
|
|
val = (val >> bit_start) & mask
|
|
return val
|
|
|
|
|
|
def set_bits_to_bytes(bytes, start, length, bvalue):
|
|
if length == 0:
|
|
return
|
|
byte_start = (start) // 8
|
|
byte_end = (start + length - 1) // 8
|
|
bit_start = start & 7
|
|
mask = (1 << length) - 1
|
|
val = bytes_to_value(bytes[byte_start:byte_end + 1])
|
|
val &= ~(mask << bit_start)
|
|
val |= ((bvalue & mask) << bit_start)
|
|
bytes[byte_start:byte_end+1] = value_to_bytearray(
|
|
val,
|
|
byte_end + 1 - byte_start)
|
|
|
|
|
|
def value_to_bytes(value, length):
|
|
return value.to_bytes(length, 'little')
|
|
|
|
|
|
def bytes_to_value(bytes):
|
|
return int.from_bytes(bytes, 'little')
|
|
|
|
|
|
def value_to_bytearray(value, length):
|
|
return bytearray(value_to_bytes(value, length))
|
|
|
|
# def value_to_bytearray (value, length):
|
|
return bytearray(value_to_bytes(value, length))
|
|
|
|
|
|
def get_aligned_value(value, alignment=4):
|
|
if alignment != (1 << (alignment.bit_length() - 1)):
|
|
raise Exception(
|
|
'Alignment (0x%x) should to be power of 2 !' % alignment)
|
|
value = (value + (alignment - 1)) & ~(alignment - 1)
|
|
return value
|
|
|
|
|
|
def get_padding_length(data_len, alignment=4):
|
|
new_data_len = get_aligned_value(data_len, alignment)
|
|
return new_data_len - data_len
|
|
|
|
|
|
def get_file_data(file, mode='rb'):
|
|
return open(file, mode).read()
|
|
|
|
|
|
def gen_file_from_object(file, object):
|
|
open(file, 'wb').write(object)
|
|
|
|
|
|
def gen_file_with_size(file, size):
|
|
open(file, 'wb').write(b'\xFF' * size)
|
|
|
|
|
|
def check_files_exist(base_name_list, dir='', ext=''):
|
|
for each in base_name_list:
|
|
if not os.path.exists(os.path.join(dir, each + ext)):
|
|
return False
|
|
return True
|
|
|
|
|
|
def load_source(name, filepath):
|
|
mod = SourceFileLoader(name, filepath).load_module()
|
|
return mod
|
|
|
|
|
|
def get_openssl_path():
|
|
if os.name == 'nt':
|
|
if 'OPENSSL_PATH' not in os.environ:
|
|
openssl_dir = "C:\\Openssl\\bin\\"
|
|
if os.path.exists(openssl_dir):
|
|
os.environ['OPENSSL_PATH'] = openssl_dir
|
|
else:
|
|
os.environ['OPENSSL_PATH'] = "C:\\Openssl\\"
|
|
if 'OPENSSL_CONF' not in os.environ:
|
|
openssl_cfg = "C:\\Openssl\\openssl.cfg"
|
|
if os.path.exists(openssl_cfg):
|
|
os.environ['OPENSSL_CONF'] = openssl_cfg
|
|
openssl = os.path.join(
|
|
os.environ.get('OPENSSL_PATH', ''),
|
|
'openssl.exe')
|
|
else:
|
|
# Get openssl path for Linux cases
|
|
openssl = shutil.which('openssl')
|
|
|
|
return openssl
|
|
|
|
|
|
def run_process(arg_list, print_cmd=False, capture_out=False):
|
|
sys.stdout.flush()
|
|
if os.name == 'nt' and os.path.splitext(arg_list[0])[1] == '' and \
|
|
os.path.exists(arg_list[0] + '.exe'):
|
|
arg_list[0] += '.exe'
|
|
if print_cmd:
|
|
print(' '.join(arg_list))
|
|
|
|
exc = None
|
|
result = 0
|
|
output = ''
|
|
try:
|
|
if capture_out:
|
|
output = subprocess.check_output(arg_list).decode()
|
|
else:
|
|
result = subprocess.call(arg_list)
|
|
except Exception as ex:
|
|
result = 1
|
|
exc = ex
|
|
|
|
if result:
|
|
if not print_cmd:
|
|
print('Error in running process:\n %s' % ' '.join(arg_list))
|
|
if exc is None:
|
|
sys.exit(1)
|
|
else:
|
|
raise exc
|
|
|
|
return output
|
|
|
|
|
|
# Adjust hash type algorithm based on Public key file
|
|
def adjust_hash_type(pub_key_file):
|
|
key_type = get_key_type(pub_key_file)
|
|
if key_type == 'RSA2048':
|
|
hash_type = 'SHA2_256'
|
|
elif key_type == 'RSA3072':
|
|
hash_type = 'SHA2_384'
|
|
else:
|
|
hash_type = None
|
|
|
|
return hash_type
|
|
|
|
|
|
def rsa_sign_file(
|
|
priv_key, pub_key, hash_type, sign_scheme,
|
|
in_file, out_file, inc_dat=False, inc_key=False):
|
|
|
|
bins = bytearray()
|
|
if inc_dat:
|
|
bins.extend(get_file_data(in_file))
|
|
|
|
|
|
# def single_sign_file(priv_key, hash_type, sign_scheme, in_file, out_file):
|
|
|
|
out_data = get_file_data(out_file)
|
|
|
|
sign = SIGNATURE_HDR()
|
|
sign.SigSize = len(out_data)
|
|
sign.SigType = SIGN_TYPE_SCHEME[sign_scheme]
|
|
sign.HashAlg = HASH_TYPE_VALUE[hash_type]
|
|
|
|
bins.extend(bytearray(sign) + out_data)
|
|
if inc_key:
|
|
key = gen_pub_key(priv_key, pub_key)
|
|
bins.extend(key)
|
|
|
|
if len(bins) != len(out_data):
|
|
gen_file_from_object(out_file, bins)
|
|
|
|
|
|
def get_key_type(in_key):
|
|
|
|
# Check in_key is file or key Id
|
|
if not os.path.exists(in_key):
|
|
key = bytearray(gen_pub_key(in_key))
|
|
else:
|
|
# Check for public key in binary format.
|
|
key = bytearray(get_file_data(in_key))
|
|
|
|
pub_key_hdr = PUB_KEY_HDR.from_buffer(key)
|
|
if pub_key_hdr.Identifier != b'PUBK':
|
|
pub_key = gen_pub_key(in_key)
|
|
pub_key_hdr = PUB_KEY_HDR.from_buffer(pub_key)
|
|
|
|
key_type = next(
|
|
(key for key,
|
|
value in PUB_KEY_TYPE.items() if value == pub_key_hdr.KeyType))
|
|
return '%s%d' % (key_type, (pub_key_hdr.KeySize - 4) * 8)
|
|
|
|
|
|
def get_auth_hash_type(key_type, sign_scheme):
|
|
if key_type == "RSA2048" and sign_scheme == "RSA_PKCS1":
|
|
hash_type = 'SHA2_256'
|
|
auth_type = 'RSA2048_PKCS1_SHA2_256'
|
|
elif key_type == "RSA3072" and sign_scheme == "RSA_PKCS1":
|
|
hash_type = 'SHA2_384'
|
|
auth_type = 'RSA3072_PKCS1_SHA2_384'
|
|
elif key_type == "RSA2048" and sign_scheme == "RSA_PSS":
|
|
hash_type = 'SHA2_256'
|
|
auth_type = 'RSA2048_PSS_SHA2_256'
|
|
elif key_type == "RSA3072" and sign_scheme == "RSA_PSS":
|
|
hash_type = 'SHA2_384'
|
|
auth_type = 'RSA3072_PSS_SHA2_384'
|
|
else:
|
|
hash_type = ''
|
|
auth_type = ''
|
|
return auth_type, hash_type
|
|
|
|
|
|
# def single_sign_gen_pub_key(in_key, pub_key_file=None):
|
|
|
|
|
|
def gen_pub_key(in_key, pub_key=None):
|
|
|
|
keydata = single_sign_gen_pub_key(in_key, pub_key)
|
|
|
|
publickey = PUB_KEY_HDR()
|
|
publickey.KeySize = len(keydata)
|
|
publickey.KeyType = PUB_KEY_TYPE['RSA']
|
|
|
|
key = bytearray(publickey) + keydata
|
|
|
|
if pub_key:
|
|
gen_file_from_object(pub_key, key)
|
|
|
|
return key
|
|
|
|
|
|
def decompress(in_file, out_file, tool_dir=''):
|
|
if not os.path.isfile(in_file):
|
|
raise Exception("Invalid input file '%s' !" % in_file)
|
|
|
|
# Remove the Lz Header
|
|
fi = open(in_file, 'rb')
|
|
di = bytearray(fi.read())
|
|
fi.close()
|
|
|
|
lz_hdr = LZ_HEADER.from_buffer(di)
|
|
offset = sizeof(lz_hdr)
|
|
if lz_hdr.signature == b"LZDM" or lz_hdr.compressed_len == 0:
|
|
fo = open(out_file, 'wb')
|
|
fo.write(di[offset:offset + lz_hdr.compressed_len])
|
|
fo.close()
|
|
return
|
|
|
|
temp = os.path.splitext(out_file)[0] + '.tmp'
|
|
if lz_hdr.signature == b"LZMA":
|
|
alg = "Lzma"
|
|
elif lz_hdr.signature == b"LZ4 ":
|
|
alg = "Lz4"
|
|
else:
|
|
raise Exception("Unsupported compression '%s' !" % lz_hdr.signature)
|
|
|
|
fo = open(temp, 'wb')
|
|
fo.write(di[offset:offset + lz_hdr.compressed_len])
|
|
fo.close()
|
|
|
|
compress_tool = "%sCompress" % alg
|
|
if alg == "Lz4":
|
|
try:
|
|
cmdline = [
|
|
os.path.join(tool_dir, compress_tool),
|
|
"-d",
|
|
"-o", out_file,
|
|
temp]
|
|
run_process(cmdline, False, True)
|
|
except Exception:
|
|
msg_string = "Could not find/use CompressLz4 tool, " \
|
|
"trying with python lz4..."
|
|
print(msg_string)
|
|
try:
|
|
import lz4.block
|
|
if lz4.VERSION != '3.1.1':
|
|
msg_string = "Recommended lz4 module version " \
|
|
"is '3.1.1'," + lz4.VERSION \
|
|
+ " is currently installed."
|
|
print(msg_string)
|
|
except ImportError:
|
|
msg_string = "Could not import lz4, use " \
|
|
"'python -m pip install lz4==3.1.1' " \
|
|
"to install it."
|
|
print(msg_string)
|
|
exit(1)
|
|
decompress_data = lz4.block.decompress(get_file_data(temp))
|
|
with open(out_file, "wb") as lz4bin:
|
|
lz4bin.write(decompress_data)
|
|
else:
|
|
cmdline = [
|
|
os.path.join(tool_dir, compress_tool),
|
|
"-d",
|
|
"-o", out_file,
|
|
temp]
|
|
run_process(cmdline, False, True)
|
|
os.remove(temp)
|
|
|
|
|
|
def compress(in_file, alg, svn=0, out_path='', tool_dir=''):
|
|
if not os.path.isfile(in_file):
|
|
raise Exception("Invalid input file '%s' !" % in_file)
|
|
|
|
basename, ext = os.path.splitext(os.path.basename(in_file))
|
|
if out_path:
|
|
if os.path.isdir(out_path):
|
|
out_file = os.path.join(out_path, basename + '.lz')
|
|
else:
|
|
out_file = os.path.join(out_path)
|
|
else:
|
|
out_file = os.path.splitext(in_file)[0] + '.lz'
|
|
|
|
if alg == "Lzma":
|
|
sig = "LZMA"
|
|
elif alg == "Tiano":
|
|
sig = "LZUF"
|
|
elif alg == "Lz4":
|
|
sig = "LZ4 "
|
|
elif alg == "Dummy":
|
|
sig = "LZDM"
|
|
else:
|
|
raise Exception("Unsupported compression '%s' !" % alg)
|
|
|
|
in_len = os.path.getsize(in_file)
|
|
if in_len > 0:
|
|
compress_tool = "%sCompress" % alg
|
|
if sig == "LZDM":
|
|
shutil.copy(in_file, out_file)
|
|
compress_data = get_file_data(out_file)
|
|
elif sig == "LZ4 ":
|
|
try:
|
|
cmdline = [
|
|
os.path.join(tool_dir, compress_tool),
|
|
"-e",
|
|
"-o", out_file,
|
|
in_file]
|
|
run_process(cmdline, False, True)
|
|
compress_data = get_file_data(out_file)
|
|
except Exception:
|
|
msg_string = "Could not find/use CompressLz4 tool, " \
|
|
"trying with python lz4..."
|
|
print(msg_string)
|
|
try:
|
|
import lz4.block
|
|
if lz4.VERSION != '3.1.1':
|
|
msg_string = "Recommended lz4 module version " \
|
|
"is '3.1.1', " + lz4.VERSION \
|
|
+ " is currently installed."
|
|
print(msg_string)
|
|
except ImportError:
|
|
msg_string = "Could not import lz4, use " \
|
|
"'python -m pip install lz4==3.1.1' " \
|
|
"to install it."
|
|
print(msg_string)
|
|
exit(1)
|
|
compress_data = lz4.block.compress(
|
|
get_file_data(in_file),
|
|
mode='high_compression')
|
|
elif sig == "LZMA":
|
|
cmdline = [
|
|
os.path.join(tool_dir, compress_tool),
|
|
"-e",
|
|
"-o", out_file,
|
|
in_file]
|
|
run_process(cmdline, False, True)
|
|
compress_data = get_file_data(out_file)
|
|
else:
|
|
compress_data = bytearray()
|
|
|
|
lz_hdr = LZ_HEADER()
|
|
lz_hdr.signature = sig.encode()
|
|
lz_hdr.svn = svn
|
|
lz_hdr.compressed_len = len(compress_data)
|
|
lz_hdr.length = os.path.getsize(in_file)
|
|
data = bytearray()
|
|
data.extend(lz_hdr)
|
|
data.extend(compress_data)
|
|
gen_file_from_object(out_file, data)
|
|
|
|
return out_file
|