1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3
4# Copyright (c) 2023 Huawei Device Co., Ltd.
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import os
18import sys
19import argparse
20import subprocess
21import secrets
22
23MOUNT_POINT = "/module_update"
24BLOCK_SIZE = 4096
25HASH_SIZE = 32
26IMG_MARGIN_SIZE = 16777216
27
28
29def run_cmd(cmd):
30    res = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE,
31                           stderr=subprocess.PIPE)
32    sout, serr = res.communicate(timeout=120)
33
34    return res.pid, res.returncode, sout, serr
35
36
37def verify_ret(res):
38    if res[1]:
39        print(" ".join(["pid ", str(res[0]), " ret ", str(res[1]), "\n",
40                        res[2].decode(), res[3].decode()]))
41        print("MkImages failed errno: %s" % str(res[1]))
42        sys.exit(2)
43
44
45def sparse_img2simg(is_sparse, device):
46    # we don't support sparse in mktools.
47    if "sparse" in is_sparse:
48        tmp_device = "%s.tmp" % device
49        run_cmd(" ".join(["img2simg ", device, " ", tmp_device]))
50        os.rename(tmp_device, device)
51
52
53def getdirsize(target_dir):
54    size = 0
55    for root, dirs, files in os.walk(target_dir):
56        size += sum([roundup(os.path.getsize(os.path.join(root, name)), BLOCK_SIZE) for name in files])
57    return size
58
59
60def roundup(size, unit):
61    return (size + unit - 1) & (~(unit - 1))
62
63
64def resize_img(device):
65    run_cmd(" ".join(["resize2fs", "-M", device]))
66
67
68def get_hash_tree_size(img_size):
69    treesize = 0
70    size = img_size
71    while size > BLOCK_SIZE:
72        blocknum = size // BLOCK_SIZE
73        level_size = roundup(blocknum * HASH_SIZE, BLOCK_SIZE)
74        treesize += level_size
75        size = level_size
76    return treesize
77
78
79def sign_img(device, pubkey, privkey, output):
80    hvb_tools = "hvbtool.py"
81    source_image = device
82    partition = "update"
83    img_size = os.path.getsize(device)
84    partition_size = img_size + get_hash_tree_size(img_size) + BLOCK_SIZE
85    salt = secrets.token_hex(20)
86    algorithm = "SHA256_RSA2048"
87    rollback_index = "0"
88    rollback_location = "0"
89    run_cmd(" ".join([hvb_tools, "make_hashtree_footer", "--image", device,
90        "--partition", partition,
91        "--partition_size", str(partition_size),
92        "--salt", salt,
93        "--pubkey", pubkey,
94        "--privkey", privkey,
95        "--algorithm", algorithm,
96        "--rollback_index", rollback_index,
97        "--rollback_location", rollback_location,
98        "--output", output]))
99
100
101def build_image(args):
102    image_type = "raw"
103    if args.sparse_image:
104        image_type = "sparse"
105    if args.build_image_tools_path:
106        env_path = ':'.join(args.build_image_tools_path)
107        os.environ['PATH'] = '{}:{}'.format(env_path, os.environ.get('PATH'))
108    src_dir = args.input_path
109    device = args.output_image_path
110    is_sparse = image_type
111    if "ext4" == args.fs_type:
112        fs_type = "ext4"
113        mkfs_tools = "mkextimage.py"
114    elif "f2fs" == args.fs_type:
115        fs_type = "f2fs"
116        mkfs_tools = "mkf2fsimage.py"
117    elif "cpio" == args.fs_type:
118        fs_type = "cpio"
119        mkfs_tools = "mkcpioimage.py"
120    else:
121        print("not support filesystem type!!")
122        sys.exit(1)
123
124    fs_size = getdirsize(src_dir) + IMG_MARGIN_SIZE
125    fs_type_config = "--fs_type=" + fs_type
126    dac_config = args.dac_config
127    file_context = args.file_context
128
129    mk_configs = " ".join([src_dir, device, MOUNT_POINT, str(fs_size), fs_type_config,
130        "--dac_config", dac_config, "--file_context", file_context])
131    res = run_cmd(" ".join([mkfs_tools, mk_configs]))
132    verify_ret(res)
133    sparse_img2simg(is_sparse, device)
134    resize_img(device)
135    if args.pubkey:
136        sign_img(device, args.pubkey, args.privkey, args.output_sign_image_path)
137    else:
138        print("not sign image")
139
140
141def main(argv):
142    parser = argparse.ArgumentParser()
143    parser.add_argument('--image-name', required=True)
144    parser.add_argument('--fs-type', required=True)
145    parser.add_argument('--dac-config', required=True)
146    parser.add_argument('--file-context', required=True)
147    parser.add_argument('--input-path', required=True)
148    parser.add_argument('--output-image-path', required=True)
149    parser.add_argument('--sparse-image',
150                        dest="sparse_image",
151                        action='store_true')
152    parser.set_defaults(sparse_image=False)
153    parser.add_argument('--build-image-tools-path', nargs='*', required=False)
154    parser.add_argument('--target-cpu', required=False)
155    parser.add_argument('--pubkey', required=False)
156    parser.add_argument('--privkey', required=False)
157    parser.add_argument('--output-sign-image-path', required=False)
158    args = parser.parse_args(argv)
159
160    if os.path.exists(args.output_image_path):
161        os.remove(args.output_image_path)
162    if os.path.isdir(args.input_path):
163        build_image(args)
164    return 0
165
166
167if __name__ == '__main__':
168    sys.exit(main(sys.argv[1:]))
169