1#!/usr/bin/env python
2#
3# Copyright (c) 2015, Linaro Limited
4# All rights reserved.
5#
6# Redistribution and use in source and binary forms, with or without
7# modification, are permitted provided that the following conditions are met:
8#
9# 1. Redistributions of source code must retain the above copyright notice,
10# this list of conditions and the following disclaimer.
11#
12# 2. Redistributions in binary form must reproduce the above copyright notice,
13# this list of conditions and the following disclaimer in the documentation
14# and/or other materials provided with the distribution.
15#
16# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
17# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
20# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
21# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
22# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
23# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
24# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
25# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
26# POSSIBILITY OF SUCH DAMAGE.
27#
28
29def get_args():
30	from argparse import ArgumentParser
31
32	parser = ArgumentParser()
33	parser.add_argument('--in', required=True, dest='inf', \
34			help='Name of input file (unsigned TA)')
35	return parser.parse_args()
36
37def assert_file_exists(fname):
38	import os.path
39
40	if(os.path.isfile(fname)):
41		return True
42	else:
43		raise FileNotFoundError('File ' + fname + ' was not found')
44
45def main():
46	from Crypto.Signature import PKCS1_v1_5
47	from Crypto.Hash import SHA256
48	from Crypto.PublicKey import RSA
49	import struct, base64, os.path, sys
50
51	args = get_args()
52
53	assert_file_exists(args.inf)
54
55	# Read input file (unsigned TA)
56	f = open(args.inf, 'rb')
57	img = f.read()
58	f.close()
59
60	h = SHA256.new()
61
62	digest_len = h.digest_size
63	#We plan to use RSA 2048 bit keys so signature is 256 bytes
64	sig_len = 256 #len(signer.sign(h))
65	img_size = len(img)
66
67	magic = 0x4f545348	# SHDR_MAGIC
68	img_type = 0		# SHDR_TA
69	algo = 0x70004830	# TEE_ALG_RSASSA_PKCS1_V1_5_SHA256
70	shdr = struct.pack('<IIIIHH', magic, img_type, img_size, algo, digest_len, sig_len)
71
72	h.update(shdr)
73	h.update(img)
74	dig = h.digest()
75
76	print("Image size is:", img_size)
77	print("Digest length is:", digest_len)
78	print("Digest is:", base64.b64encode(dig))
79
80	digest_filename = os.path.splitext(args.inf)[0].split(".")[0]+'.dig'
81
82	print('Digest Filename:', digest_filename)
83
84	# Write digest to file
85	dig_out = open(digest_filename, 'wb+')
86	dig_out.write(base64.b64encode(dig))
87	dig_out.close()
88
89if __name__ == "__main__":
90	main()
91