#!/usr/bin/env python3
#
# Copyright (C) 2021 Amlogic, Inc. All rights reserved.
#
# This source code is subject to the terms and conditions defined in the
# file 'LICENSE' which is part of this source code package.

ALGO = 'gen-prot-aes128'
MRK_NAME = 'DGPK2'
EK3 = 'a1f9288d5e382b34a2897e9194ee2911'
EK2 = '754ac2fdb002e3fa13407d3367ce293d'
#EK1 = '8b8f5a227bb1c1487d21a0fd6d842833'

TARGET_NAME = 'dac'
MAX_IDX = 100000000
DERIVED_FILE_NUM = 0

def get_args():
	from argparse import ArgumentParser

	parser = ArgumentParser()
	parser.add_argument('--soc', type = str, required = True, help = 'soc')
	parser.add_argument('--key_tool', type = str, required = True, help = 'vendor key tool path')
	parser.add_argument('--pfpk_hmac_generator', type = str, required = True, help = 'pfpk hmac generator path')
	parser.add_argument('--dgpk2', type = str, default = '', help = 'dgpk2 file or directory')
	parser.add_argument('--pfid', type = str, default = '', help = 'pfid file or directory')
	parser.add_argument('--pfpk_hmac', type = str, default = '', help = 'pfpk_hmac file or directory')
	parser.add_argument('--nonce', type = str, default = '', help = 'nonce file or nonce hex string')
	parser.add_argument('--segid', type = str, default = '', help = 'segid file or segid hex string')
	parser.add_argument('--enhance', nargs = '?', default = None, const = True, help = 'enable security enhanced mode')
	parser.add_argument('--stage', type = str, default = '2', help = 'boot stage string')
	parser.add_argument('--out_dir', type = str, required = True, help = 'output directory')
	return parser.parse_args()

def derive_one_dac(file_dgpk2, file_pfid, out_dir, idx, key_tool, soc, pfpk_hmac_generator, nonce, segid, stage, enhance):
	import os
	import sys
	import binascii
	import hmac
	import hashlib

	if int(idx) > MAX_IDX:
		print('the number of files exceeds the limit')
		sys.exit(1)

	f = open(file_dgpk2, 'rb')
	mrk = f.read()
	f.close()

	f = open(file_pfid, 'rb')
	ek1 = f.read()
	f.close()

	if os.path.isfile(segid):
		f = open(segid, 'rb')
		segid = f.read()
		f.close()
	else:
		segid = binascii.unhexlify(segid)
	if len(segid) == 4:
		ek1 = ek1[:4] + segid + ek1[8:]
	elif len(segid):
		print('segid length error')
		sys.exit(1)

	if os.path.isfile(nonce):
		f = open(nonce, 'rb')
		nonce = f.read()
		f.close()
	else:
		nonce = binascii.unhexlify(nonce)

	# due to the other para aren't supported, can only use "SC2"
	cmd = key_tool + ' ' + ALGO + ' --chipset=SC2' \
	      + ' --mrk=' + bytes.decode(binascii.b2a_hex(mrk)) \
	      + ' --mrk-name=' + MRK_NAME + ' --boot-stage=' + stage \
	      + ' --ek3=' + EK3 + ' --ek2=' + EK2 \
	      + ' --ek1=' + bytes.decode(binascii.b2a_hex(ek1))
	pfpk = os.popen(cmd).read()[:32]

	if enhance is True:
		cmd = pfpk_hmac_generator + ' -p ' + pfpk + ' -e '
	else:
		cmd = pfpk_hmac_generator + ' -p ' + pfpk

	pfpk_hmac = binascii.unhexlify((os.popen(cmd).read())[:64])

	dac = hmac.HMAC(pfpk_hmac, ek1 + nonce, hashlib.sha256).digest()

	file_dac = out_dir + '/' + soc + '_' + TARGET_NAME + '_' + idx + '.bin'
	f = open(file_dac, 'wb')
	f.write(dac)
	f.close()

	global DERIVED_FILE_NUM
	DERIVED_FILE_NUM = DERIVED_FILE_NUM + 1
	print(os.path.abspath(file_dac) + ' derived successfully')

def derive_one_dac_by_pfpk_hmac(file_pfpk_hmac, file_pfid, out_dir, idx, soc, nonce, segid):
	import os
	import sys
	import binascii
	import hmac
	import hashlib

	if int(idx) > MAX_IDX:
		print('the number of files exceeds the limit')
		sys.exit(1)

	f = open(file_pfpk_hmac, 'rb')
	pfpk_hmac = f.read()
	f.close()

	f = open(file_pfid, 'rb')
	ek1 = f.read()
	f.close()

	if os.path.isfile(segid):
		f = open(segid, 'rb')
		segid = f.read()
		f.close()
	else:
		segid = binascii.unhexlify(segid)
	if len(segid) == 4:
		ek1 = ek1[:4] + segid + ek1[8:]
	elif len(segid):
		print('segid length error')
		sys.exit(1)

	if os.path.isfile(nonce):
		f = open(nonce, 'rb')
		nonce = f.read()
		f.close()
	else:
		nonce = binascii.unhexlify(nonce)

	dac = hmac.HMAC(pfpk_hmac, ek1 + nonce, hashlib.sha256).digest()

	file_dac = out_dir + '/' + soc + '_' + TARGET_NAME + '_' + idx + '.bin'
	f = open(file_dac, 'wb')
	f.write(dac)
	f.close()

	global DERIVED_FILE_NUM
	DERIVED_FILE_NUM = DERIVED_FILE_NUM + 1
	print(os.path.abspath(file_dac) + ' derived successfully')

def main():
	import os
	import sys

	args = get_args()

	if not os.path.exists(args.key_tool):
		print(args.key_tool + ' not exist')
		sys.exit(1)

	if not os.path.exists(args.pfpk_hmac):
		if not os.path.exists(args.pfpk_hmac_generator):
			print(args.pfpk_hmac_generator + ' not exist')
			sys.exit(1)
		if not os.path.exists(args.dgpk2):
			print('dgpk2 not exist')
			sys.exit(1)
		if not os.path.exists(args.pfid):
			print('pfid not exist')
			sys.exit(1)
	else:
		if not os.path.exists(args.pfid):
			print('pfid not exist')
			sys.exit(1)
		if os.path.exists(args.dgpk2):
			print(args.dgpk2 + ' is not necessary')
			sys.exit(1)

	if not os.path.exists(args.out_dir):
		os.mkdir(args.out_dir)

	if os.path.isfile(args.dgpk2) and os.path.isfile(args.pfid):
		derive_one_dac(args.dgpk2, args.pfid, args.out_dir, '000000001', args.key_tool, args.soc, args.pfpk_hmac_generator, args.nonce, args.segid, args.stage, args.enhance)
	elif os.path.isdir(args.dgpk2) and os.path.isdir(args.pfid):
		for parent, dnames, fnames in os.walk(args.dgpk2):
			for fname in fnames:
				idx = fname[-13 : -4]
				file_pfid = args.pfid + '/' + args.soc + '_pfid_' + idx + '.bin'
				file_dgpk2 = os.path.join(parent, fname)
				if os.path.exists(file_pfid):
					derive_one_dac(file_dgpk2, file_pfid, args.out_dir, idx, args.key_tool, args.soc, args.pfpk_hmac_generator, args.nonce, args.segid, args.stage, args.enhance)
	elif os.path.isfile(args.pfpk_hmac) and os.path.isfile(args.pfid):
		derive_one_dac_by_pfpk_hmac(args.pfpk_hmac, args.pfid, args.out_dir, '000000001', args.soc, args.nonce, args.segid)
	elif os.path.isdir(args.pfpk_hmac) and os.path.isdir(args.pfid):
		for parent, dnames, fnames in os.walk(args.pfpk_hmac):
			for fname in fnames:
				idx = fname[-13 : -4]
				file_pfid = args.pfid + '/' + args.soc + '_pfid_' + idx + '.bin'
				file_pfpk_hmac = os.path.join(parent, fname)
				if os.path.exists(file_pfid):
					derive_one_dac_by_pfpk_hmac(file_pfpk_hmac, file_pfid, args.out_dir, idx, args.soc, args.nonce, args.segid)
	else:
		print('input parameter error')
		sys.exit(1)

	print(str(DERIVED_FILE_NUM) + ' ' + TARGET_NAME + ' files derived successfully')

if __name__ == '__main__':
	main()