140681896Sopenharmony_ci#!/usr/bin/env python3
240681896Sopenharmony_ci# -*- coding: utf-8 -*-
340681896Sopenharmony_ci
440681896Sopenharmony_ci# Copyright (c) 2023 Huawei Device Co., Ltd.
540681896Sopenharmony_ci# Licensed under the Apache License, Version 2.0 (the "License");
640681896Sopenharmony_ci# you may not use this file except in compliance with the License.
740681896Sopenharmony_ci# You may obtain a copy of the License at
840681896Sopenharmony_ci#
940681896Sopenharmony_ci# http://www.apache.org/licenses/LICENSE-2.0
1040681896Sopenharmony_ci#
1140681896Sopenharmony_ci# Unless required by applicable law or agreed to in writing, software
1240681896Sopenharmony_ci# distributed under the License is distributed on an "AS IS" BASIS,
1340681896Sopenharmony_ci# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1440681896Sopenharmony_ci# See the License for the specific language governing permissions and
1540681896Sopenharmony_ci# limitations under the License.
1640681896Sopenharmony_ci
1740681896Sopenharmony_ciimport hashlib
1840681896Sopenharmony_cifrom base64 import b64encode
1940681896Sopenharmony_cifrom build_pkcs7 import BLOCK_SIZE, sign_digest
2040681896Sopenharmony_cifrom log_exception import UPDATE_LOGGER
2140681896Sopenharmony_cifrom utils import OPTIONS_MANAGER
2240681896Sopenharmony_ci
2340681896Sopenharmony_ci
2440681896Sopenharmony_cidef sign_func_sha256(sign_file, private_key_file):
2540681896Sopenharmony_ci    """
2640681896Sopenharmony_ci    sign one file with private key
2740681896Sopenharmony_ci    :param sign_file: path of file ready to be signed
2840681896Sopenharmony_ci    :param private_key_file: private key path, ex. rsa_private_key2048.pem
2940681896Sopenharmony_ci    :return: base64 code of the signature
3040681896Sopenharmony_ci    """
3140681896Sopenharmony_ci    hash_sha256 = hashlib.sha256()
3240681896Sopenharmony_ci    with open(sign_file, 'rb') as file:
3340681896Sopenharmony_ci        chunk = file.read(BLOCK_SIZE)
3440681896Sopenharmony_ci        while chunk:
3540681896Sopenharmony_ci            hash_sha256.update(chunk)
3640681896Sopenharmony_ci            chunk = file.read(BLOCK_SIZE)
3740681896Sopenharmony_ci    signature = sign_digest(hash_sha256.digest(), private_key_file)
3840681896Sopenharmony_ci    if signature == False:
3940681896Sopenharmony_ci        UPDATE_LOGGER.print_log("sign digest failed", log_type=UPDATE_LOGGER.ERROR_LOG)
4040681896Sopenharmony_ci        return ""
4140681896Sopenharmony_ci    return str(b64encode(signature).decode("ascii"))
4240681896Sopenharmony_ci
4340681896Sopenharmony_ci
4440681896Sopenharmony_cidef generate_signed_data(file_lists, sign_func, private_key_file):
4540681896Sopenharmony_ci    """
4640681896Sopenharmony_ci    get hash signed data of file lists, hash signed data format:
4740681896Sopenharmony_ci    Name: build_tools/updater_binary
4840681896Sopenharmony_ci    signed-data: xxxxxxx
4940681896Sopenharmony_ci
5040681896Sopenharmony_ci    Name: build_tools/updater_binary
5140681896Sopenharmony_ci    signed-data: xxxxxxx
5240681896Sopenharmony_ci
5340681896Sopenharmony_ci    ....
5440681896Sopenharmony_ci    :param file_lists: path list of file ready to be signed, list item contains file_path and name_in_signed_data
5540681896Sopenharmony_ci    :param sign_func: signature function, ex. sign_func_sha256
5640681896Sopenharmony_ci    :param private_key_file: private key path, ex. rsa_private_key2048.pem
5740681896Sopenharmony_ci    :return: hash signed data of the file_lists
5840681896Sopenharmony_ci    """
5940681896Sopenharmony_ci    max_file_num = 32
6040681896Sopenharmony_ci    if not sign_func:
6140681896Sopenharmony_ci        UPDATE_LOGGER.print_log("please provide a sign function", log_type=UPDATE_LOGGER.ERROR_LOG)
6240681896Sopenharmony_ci        return ""
6340681896Sopenharmony_ci
6440681896Sopenharmony_ci    if len(file_lists) > max_file_num:
6540681896Sopenharmony_ci        UPDATE_LOGGER.print_log("signed file can't be more than %d" % max_file_num,
6640681896Sopenharmony_ci            log_type=UPDATE_LOGGER.ERROR_LOG)
6740681896Sopenharmony_ci        return ""
6840681896Sopenharmony_ci    sign_res_list = []
6940681896Sopenharmony_ci    for file, name in file_lists:
7040681896Sopenharmony_ci        sign_res = sign_func(file, private_key_file)
7140681896Sopenharmony_ci        if sign_res == "":
7240681896Sopenharmony_ci            UPDATE_LOGGER.print_log("sign file {} failed".format(name), log_type=UPDATE_LOGGER.ERROR_LOG)
7340681896Sopenharmony_ci            return ""
7440681896Sopenharmony_ci        sign_res_list += ["Name: {}\nsigned-data: {}\n".format(name, sign_res)]
7540681896Sopenharmony_ci    return "\n".join(sign_res_list)
7640681896Sopenharmony_ci
7740681896Sopenharmony_ci
7840681896Sopenharmony_cidef generate_signed_data_default(file_lists):
7940681896Sopenharmony_ci    return generate_signed_data(file_lists, sign_func_sha256, OPTIONS_MANAGER.private_key)