script_generator.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2021 Huawei Device Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Description : Create script file for updater
"""

import os
import re
import tempfile
from decimal import getcontext
from decimal import Decimal

from log_exception import VendorExpandError
from log_exception import UPDATE_LOGGER
from utils import OPTIONS_MANAGER
from utils import PARTITION_FILE
from utils import TOTAL_SCRIPT_FILE_NAME
from utils import SCRIPT_FILE_NAME
from utils import SCRIPT_KEY_LIST
from utils import UPDATE_BIN_FILE_NAME


class Script:

    def __init__(self):
        self.script = []
        self.version = 0
        self.info = {}

    def add_command(self, cmd=None):
        """
        Add command content to the script.
        :param cmd: command content
        :return:
        """
        self.script.append(cmd)

    def get_script(self):
        """
        Get the script list.
        :return: script list
        """
        return self.script

    def sha_check(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'sha_check')

    def image_sha_check(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'image_sha_check')

    def first_block_check(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'first_block_check')

    def image_patch(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'image_patch')

    def abort(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'abort')

    def show_progress(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'show_progress')

    def block_update(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'block_update')

    def raw_image_write(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'raw_image_write')

    def get_status(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'get_status')

    def set_status(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'set_status')

    def reboot_now(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'reboot_now')

    def updater_partitions(self, *args, **kwargs):
        raise VendorExpandError(type(self), 'updater_partitions')


class PreludeScript(Script):
    def __init__(self):
        super().__init__()


class VerseScript(Script):
    def __init__(self):
        super().__init__()

    def sha_check(self, ranges_str, expected_sha, partition):
        """
        Get the sha_check command.
        :param ranges_str: ranges string
        :param expected_sha: hash value
        :param partition: image name
        :return:
        """
        cmd = ('sha_check("/{partition}", "{ranges_str}", '
               '"{expected_sha}")').format(
            ranges_str=ranges_str,
            expected_sha=expected_sha, partition=partition)
        return cmd

    def image_sha_check(self, partition, src_size, src_hash,
        dst_size, dst_hash):
        """
        Get the image_sha_check command.
        :param ranges_str: ranges string
        :param expected_sha: hash value
        :param partition: image name
        :return:
        """
        cmd = ('image_sha_check("/{partition}", '
                '"{src_size}", "{src_hash}", '
                '"{dst_size}", "{dst_hash}")').format(
                partition=partition, src_size=src_size, src_hash=src_hash,
                dst_size=dst_size, dst_hash=dst_hash)
        return cmd

    def first_block_check(self, partition):
        """
        Get the first_block_check command.
        :param partition: image name
        :return:
        """
        cmd = 'first_block_check("/{partition}")'.format(
            partition=partition)
        return cmd

    def abort(self, partition):
        """
        Get the abort command.
        :param partition: image name
        :return:
        """
        cmd = 'abort("ERROR: {partition} partition ' \
              'fails to incremental check!");\n'.format(
                partition=partition)
        return cmd

    def show_progress(self, start_progress, dur):
        """
        Get the show_progress command.
        'dur' may be zero to advance the progress via SetProgress
        :param start_progress: start progress
        :param dur: seconds
        :return:
        """
        cmd = 'show_progress({start_progress}, {dur});\n'.format(
            start_progress=float(start_progress), dur=float(dur))
        return cmd

    def image_patch(self, partition, src_size, src_hash, target_size, target_hash):
        """
        Get the image_patch command.
        :param partition:  image name
        :return:
        """
        cmd = 'image_patch("/{partition}", "{src_size}", ' \
              '"{src_hash}", "{target_size}", "{target_hash}", ' \
              '"{partition}.patch.dat");\n'.format(partition=partition, src_size=src_size,
              src_hash=src_hash, target_size=target_size, target_hash=target_hash)
        return cmd

    def block_update(self, partition):
        """
        Get the block_update command.
        :param partition:  image name
        :return:
        """
        cmd = 'block_update("/{partition}", ' \
              '"{partition}.transfer.list", "{partition}.new.dat", ' \
              '"{partition}.patch.dat");\n'.format(partition=partition)
        return cmd

    def image_write(self, partition, image_name, image_path):
        return self.raw_image_write(partition, image_name)

    def raw_image_write(self, partition, image_name):
        """
        Get the raw_image_write command.
        :param partition:  image name
        :return:
        """
        cmd = 'raw_image_write("/%s", "/%s");\n' % (partition, image_name)
        return cmd

    def get_status(self):
        """
        Get the get_status command.
        :return:
        """
        cmd = 'get_status("/misc")'
        return cmd

    def set_status(self, status_value):
        """
        Get the set_status command.
        :param status_value: status value to be set
        :return:
        """
        cmd = 'set_status("/misc", %s);\n' % status_value
        return cmd

    def reboot_now(self):
        """
        Get the reboot_now command.
        :return:
        """
        cmd = 'reboot_now();\n'
        return cmd

    def updater_partitions(self):
        """
        Get the updater_partitions command.
        :return:
        """
        cmd = 'update_partitions("/%s");\n' % PARTITION_FILE
        return cmd

    def full_image_update(self, update_file_name):
        cmd = 'update_from_bin("%s");\n' % update_file_name
        return cmd

    def pkg_extract(self, pkg_file_name, dest_path):
        cmd = 'pkg_extract("%s", "%s");\n' % (pkg_file_name, dest_path)
        return cmd


class RefrainScript(Script):
    def __init__(self):
        super().__init__()


class EndingScript(Script):
    def __init__(self):
        super().__init__()


def write_script(script_content, opera_name):
    """
    Generate the {opera}script.
    :param script_content: script content
    :param opera_name: Opera phase names corresponding to the script content
                    'prelude', 'verse', 'refrain', and 'ending'.
    :return:
    """
    script_file = tempfile.NamedTemporaryFile(mode='w+')
    script_file.write(script_content)
    script_file.seek(0)
    script_file_name = ''.join([opera_name.title(), SCRIPT_FILE_NAME])
    OPTIONS_MANAGER.opera_script_file_name_dict[opera_name].\
        append((script_file_name, script_file))
    UPDATE_LOGGER.print_log("%s generation complete!" % script_file_name)


def generate_total_script():
    """
    Generate the overall script.
    """
    content_list = []
    for each_key, each_value in \
            OPTIONS_MANAGER.opera_script_file_name_dict.items():
        for each in each_value:
            each_content = "LoadScript(\"%s\", %s);" % \
                           (each[0], SCRIPT_KEY_LIST.index(each_key))
            content_list.append(each_content)
    script_total = tempfile.NamedTemporaryFile(mode='w+')
    script_total.write('\n'.join(content_list))
    script_total.seek(0)
    OPTIONS_MANAGER.total_script_file_obj = script_total
    UPDATE_LOGGER.print_log("%s generation complete!" % TOTAL_SCRIPT_FILE_NAME)


def get_progress_value(distributable_value=100):
    """
    Allocate a progress value to each image update.
    :param distributable_value: distributable value
    :return:
    """
    progress_value_dict = {}
    full_img_list = OPTIONS_MANAGER.full_img_list
    incremental_img_list = OPTIONS_MANAGER.incremental_img_list
    file_size_list = []
    each_img_size = 0
    if len(full_img_list) == 0 and len(incremental_img_list) == 0:
        UPDATE_LOGGER.print_log(
            "get progress value failed! > getting progress value failed!",
            UPDATE_LOGGER.ERROR_LOG)
        return False
    for partition in incremental_img_list:
        # Obtain the size of the incremental image file.
        if partition in OPTIONS_MANAGER.incremental_image_file_obj_dict:
            file_obj = OPTIONS_MANAGER.incremental_image_file_obj_dict[partition]
            each_img_size = os.path.getsize(file_obj.name)
        elif partition in OPTIONS_MANAGER.incremental_block_file_obj_dict:
            new_dat_file_obj, patch_dat_file_obj, transfer_list_file_obj =\
                OPTIONS_MANAGER.incremental_block_file_obj_dict[partition].get_file_obj()
            each_img_size = os.path.getsize(new_dat_file_obj.name) + os.path.getsize(patch_dat_file_obj.name)
        file_size_list.append(each_img_size)

    total_full_img_size = 0
    for idx, _ in enumerate(full_img_list):
        # Obtain the size of the full image file.
        file_obj = OPTIONS_MANAGER.full_image_file_obj_list[idx]
        total_full_img_size += os.path.getsize(file_obj.name)
    file_size_list.append(total_full_img_size)

    proportion_value_list = get_proportion_value_list(
        file_size_list, distributable_value=distributable_value)

    adjusted_proportion_value_list = adjust_proportion_value_list(
        proportion_value_list, distributable_value)

    all_img_list = incremental_img_list + [UPDATE_BIN_FILE_NAME]
    current_progress = 40
    for idx, each_img in enumerate(all_img_list):
        temp_progress = current_progress + adjusted_proportion_value_list[idx]
        progress_value_dict[each_img] = (current_progress, temp_progress)
        current_progress = temp_progress
    return progress_value_dict


def get_proportion_value_list(file_size_list, distributable_value=100):
    """
    Obtain the calculated progress proportion value list
    (proportion_value_list).
    :param file_size_list: file size list
    :param distributable_value: distributable value
    :return proportion_value_list: progress proportion value list
    """
    sum_size = sum(file_size_list)
    getcontext().prec = 2
    proportion_value_list = []
    for each_size_value in file_size_list:
        proportion = Decimal(str(float(each_size_value))) / Decimal(
            str(float(sum_size)))
        proportion_value = int(
            Decimal(str(proportion)) *
            Decimal(str(float(distributable_value))))
        if proportion_value == 0:
            proportion_value = 1
        proportion_value_list.append(proportion_value)
    return proportion_value_list


def adjust_proportion_value_list(proportion_value_list, distributable_value):
    """
    Adjust the calculated progress proportion value list to ensure that
    sum is equal to distributable_value.
    :param proportion_value_list: calculated progress proportion value list
    :param distributable_value: number of distributable progress values
    :return proportion_value_list: new progress proportion value list
    """
    if len(proportion_value_list) == 0:
        return []
    sum_proportion_value = sum(proportion_value_list)
    if sum_proportion_value > distributable_value:
        max_value = max(proportion_value_list)
        max_idx = proportion_value_list.index(max_value)
        proportion_value_list[max_idx] = \
            max_value - (sum_proportion_value - distributable_value)
    elif sum_proportion_value < distributable_value:
        min_value = min(proportion_value_list)
        min_idx = proportion_value_list.index(min_value)
        proportion_value_list[min_idx] = \
            min_value + (distributable_value - sum_proportion_value)
    return proportion_value_list


def create_script(prelude_script, verse_script,
                  refrain_script, ending_script):
    """
    Generate the script file.
    :param prelude_script: prelude script
    :param verse_script: verse script
    :param refrain_script: refrain script
    :param ending_script: ending script
    :return:
    """
    # Generate the prelude script.
    prelude_script.add_command("\n# ---- prelude ----\n")

    # Get the distribution progress.
    progress_value_dict = get_progress_value()
    if progress_value_dict is False:
        return False
    verse_script_content_list = verse_script.get_script()
    updater_content = []
    verse_script_content = '\n'.join(verse_script_content_list)

    for key, value in progress_value_dict.items():
        show_progress_content = \
            verse_script.show_progress((value[1] - value[0]) / 100, 0)
        verse_script_content = \
            re.sub(r'%s_WRITE_FLAG' % key, '%s' % show_progress_content,
                   verse_script_content, count=1)
    # Generate the verse script.
    write_script(verse_script_content, 'verse')
    # Generate the refrain script.
    refrain_script.add_command("\n# ---- refrain ----\n")
    # Generate the ending script.
    ending_script.add_command("\n# ---- ending ----\n")

    generate_total_script()