140681896Sopenharmony_ci#!/usr/bin/env python3
240681896Sopenharmony_ci# -*- coding: utf-8 -*-
340681896Sopenharmony_ci
440681896Sopenharmony_ci# Copyright (c) 2021 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_ciDescription: create actions_list and transfer package
1840681896Sopenharmony_ci"""
1940681896Sopenharmony_ci
2040681896Sopenharmony_ciimport os
2140681896Sopenharmony_ciimport re
2240681896Sopenharmony_cifrom collections import OrderedDict
2340681896Sopenharmony_cifrom enum import Enum
2440681896Sopenharmony_ci
2540681896Sopenharmony_cifrom blocks_manager import BlocksManager
2640681896Sopenharmony_cifrom log_exception import UPDATE_LOGGER
2740681896Sopenharmony_cifrom utils import FILE_MAP_ZERO_KEY
2840681896Sopenharmony_cifrom utils import FILE_MAP_COPY_KEY
2940681896Sopenharmony_ci
3040681896Sopenharmony_ciVERSION_NAME_RE = r"[0-9]+"
3140681896Sopenharmony_ciREPLACE_CONTENT = "#"
3240681896Sopenharmony_ci
3340681896Sopenharmony_ci
3440681896Sopenharmony_ciclass ActionType(Enum):
3540681896Sopenharmony_ci    NEW = 0
3640681896Sopenharmony_ci    ZERO = 1
3740681896Sopenharmony_ci    DIFFERENT = 2
3840681896Sopenharmony_ci    MOVE = 3
3940681896Sopenharmony_ci
4040681896Sopenharmony_ci
4140681896Sopenharmony_ciclass ActionInfo:
4240681896Sopenharmony_ci    def __init__(self, type_str, tgt_name, src_name, tgt_block_set,
4340681896Sopenharmony_ci                 src_block_set):
4440681896Sopenharmony_ci        self.type_str = type_str
4540681896Sopenharmony_ci        self.tgt_name = tgt_name
4640681896Sopenharmony_ci        self.src_name = src_name
4740681896Sopenharmony_ci        self.tgt_block_set = tgt_block_set
4840681896Sopenharmony_ci        if src_block_set is not None:
4940681896Sopenharmony_ci            self.src_block_set = src_block_set
5040681896Sopenharmony_ci        else:
5140681896Sopenharmony_ci            self.src_block_set = BlocksManager()
5240681896Sopenharmony_ci        self.child = OrderedDict()
5340681896Sopenharmony_ci        self.parent = OrderedDict()
5440681896Sopenharmony_ci        self.stash_before = []
5540681896Sopenharmony_ci        self.use_stash = []
5640681896Sopenharmony_ci
5740681896Sopenharmony_ci    def get_max_block_number(self):
5840681896Sopenharmony_ci        if self.src_block_set and self.src_block_set.size() != 0:
5940681896Sopenharmony_ci            return max(self.src_block_set.range_data)
6040681896Sopenharmony_ci        else:
6140681896Sopenharmony_ci            return 0
6240681896Sopenharmony_ci
6340681896Sopenharmony_ci    def net_stash_change(self):
6440681896Sopenharmony_ci        return (sum(sr.size() for (_, sr) in self.stash_before) -
6540681896Sopenharmony_ci                sum(sr.size() for (_, sr) in self.use_stash))
6640681896Sopenharmony_ci
6740681896Sopenharmony_ci
6840681896Sopenharmony_ciclass TransfersManager(object):
6940681896Sopenharmony_ci    def __init__(self, partition, tgt_img_obj, src_img_obj,
7040681896Sopenharmony_ci                 disable_img_diff=False):
7140681896Sopenharmony_ci        self.tgt_img_obj = tgt_img_obj
7240681896Sopenharmony_ci        self.src_img_obj = src_img_obj
7340681896Sopenharmony_ci        self.partition = partition
7440681896Sopenharmony_ci        self.disable_img_diff = disable_img_diff
7540681896Sopenharmony_ci
7640681896Sopenharmony_ci        self.action_file_list = []
7740681896Sopenharmony_ci
7840681896Sopenharmony_ci    @staticmethod
7940681896Sopenharmony_ci    def simplify_file_name(file_name):
8040681896Sopenharmony_ci        base_name = os.path.basename(file_name)
8140681896Sopenharmony_ci        no_version_name = re.sub(VERSION_NAME_RE, REPLACE_CONTENT, base_name)
8240681896Sopenharmony_ci        return base_name, no_version_name
8340681896Sopenharmony_ci
8440681896Sopenharmony_ci    def arrange_source_file(self):
8540681896Sopenharmony_ci        base_names = {}
8640681896Sopenharmony_ci        version_patterns = {}
8740681896Sopenharmony_ci        for file_name in self.src_img_obj.file_map.keys():
8840681896Sopenharmony_ci            base_name, no_version_name = self.simplify_file_name(file_name)
8940681896Sopenharmony_ci            base_names[base_name] = file_name
9040681896Sopenharmony_ci            version_patterns[no_version_name] = file_name
9140681896Sopenharmony_ci        return base_names, version_patterns
9240681896Sopenharmony_ci
9340681896Sopenharmony_ci    def find_process_needs(self):
9440681896Sopenharmony_ci        """
9540681896Sopenharmony_ci        generate action_list
9640681896Sopenharmony_ci        """
9740681896Sopenharmony_ci        src_base_names, src_version_patterns = self.arrange_source_file()
9840681896Sopenharmony_ci        max_size = -1
9940681896Sopenharmony_ci
10040681896Sopenharmony_ci        for tgt_file_name, tgt_blocks in \
10140681896Sopenharmony_ci                self.tgt_img_obj.file_map.items():
10240681896Sopenharmony_ci            if FILE_MAP_ZERO_KEY == tgt_file_name:
10340681896Sopenharmony_ci                UPDATE_LOGGER.print_log("Apply ZERO type!")
10440681896Sopenharmony_ci                self.action_file_list.append(
10540681896Sopenharmony_ci                    ActionInfo(
10640681896Sopenharmony_ci                        ActionType.ZERO, tgt_file_name, FILE_MAP_ZERO_KEY,
10740681896Sopenharmony_ci                        tgt_blocks, self.src_img_obj.
10840681896Sopenharmony_ci                        file_map.get(FILE_MAP_ZERO_KEY, None)))
10940681896Sopenharmony_ci                continue
11040681896Sopenharmony_ci            if FILE_MAP_COPY_KEY == tgt_file_name:
11140681896Sopenharmony_ci                UPDATE_LOGGER.print_log("Apply COPY type!")
11240681896Sopenharmony_ci                self.action_file_list.append(
11340681896Sopenharmony_ci                    ActionInfo(ActionType.NEW, tgt_file_name,
11440681896Sopenharmony_ci                               None, tgt_blocks, None))
11540681896Sopenharmony_ci                continue
11640681896Sopenharmony_ci            if tgt_file_name in self.src_img_obj.file_map:
11740681896Sopenharmony_ci                UPDATE_LOGGER.print_log("Apply DIFF type!")
11840681896Sopenharmony_ci                action_info = ActionInfo(
11940681896Sopenharmony_ci                    ActionType.DIFFERENT, tgt_file_name, tgt_file_name,
12040681896Sopenharmony_ci                    tgt_blocks,
12140681896Sopenharmony_ci                    self.src_img_obj.file_map[tgt_file_name])
12240681896Sopenharmony_ci                max_size = action_info.get_max_block_number() \
12340681896Sopenharmony_ci                    if action_info.get_max_block_number() > \
12440681896Sopenharmony_ci                    max_size else max_size
12540681896Sopenharmony_ci                self.action_file_list.append(action_info)
12640681896Sopenharmony_ci                continue
12740681896Sopenharmony_ci            src_file_name = self.get_file_name(
12840681896Sopenharmony_ci                src_base_names, src_version_patterns, tgt_file_name)
12940681896Sopenharmony_ci            if src_file_name:
13040681896Sopenharmony_ci                action_info = ActionInfo(
13140681896Sopenharmony_ci                    ActionType.DIFFERENT, tgt_file_name, src_file_name,
13240681896Sopenharmony_ci                    tgt_blocks,
13340681896Sopenharmony_ci                    self.src_img_obj.file_map[src_file_name])
13440681896Sopenharmony_ci                max_size = action_info.get_max_block_number() if \
13540681896Sopenharmony_ci                    action_info.get_max_block_number() > max_size else max_size
13640681896Sopenharmony_ci                self.action_file_list.append(action_info)
13740681896Sopenharmony_ci                continue
13840681896Sopenharmony_ci            self.action_file_list.append(
13940681896Sopenharmony_ci                ActionInfo(ActionType.NEW, tgt_file_name,
14040681896Sopenharmony_ci                           None, tgt_blocks, None))
14140681896Sopenharmony_ci
14240681896Sopenharmony_ci        return max_size
14340681896Sopenharmony_ci
14440681896Sopenharmony_ci    def get_file_name(self, src_base_names, src_version_patterns,
14540681896Sopenharmony_ci                      tgt_file_name):
14640681896Sopenharmony_ci        tgt_base_name, tgt_version_patterns = \
14740681896Sopenharmony_ci            self.simplify_file_name(tgt_file_name)
14840681896Sopenharmony_ci        has_diff_name = True if tgt_base_name in src_base_names else False
14940681896Sopenharmony_ci        has_diff_version = \
15040681896Sopenharmony_ci            True if tgt_version_patterns in src_version_patterns else False
15140681896Sopenharmony_ci        src_file_name = \
15240681896Sopenharmony_ci            src_base_names[tgt_base_name] if has_diff_name else \
15340681896Sopenharmony_ci            src_version_patterns[tgt_version_patterns] if \
15440681896Sopenharmony_ci            has_diff_version else None
15540681896Sopenharmony_ci        return src_file_name
15640681896Sopenharmony_ci
15740681896Sopenharmony_ci    def get_action_list(self):
15840681896Sopenharmony_ci        return self.action_file_list
159