Convert mkuserimg_mke2fs.sh to python

The mkuserimg_mke2fs.sh was a shell script, which expects all the
arguments in order. This cl changes the script to python which leverages
python's argumentparser and unittest modules.

The script usage is unchanged; and this tool will be packed into the
otatools.zip in the follow up cls. Then we'll change the caller site
gradually; and remove the old shell script after that.

Bug: 112555072
Bug: 63866463
Test: run unit tests & build a userdata image
Change-Id: Ie6b687da3de31a3481363f01d2b5c12df91ca5ce

3 files changed, 394 insertions, 0 deletions

diff --git a/ext4_utils/Android.bp b/ext4_utils/Android.bp
index 0deca762..6b2f863f 100644
--- a/ext4_utils/Android.bp
+++ b/ext4_utils/Android.bp
@@ -59,3 +59,27 @@ cc_library {
         },
     },
 }
+
+python_binary_host {
+    name: "mkuserimg_mke2fs",
+    srcs: [
+        "mkuserimg_mke2fs.py",
+    ],
+
+    version: {
+        py2: {
+            enabled: true,
+            embedded_launcher: true,
+        },
+        py3: {
+            enabled: false,
+            embedded_launcher: false,
+        },
+    },
+
+    required: [
+        "mke2fs",
+        "e2fsdroid",
+    ],
+
+}
diff --git a/ext4_utils/mkuserimg_mke2fs.py b/ext4_utils/mkuserimg_mke2fs.py
new file mode 100644
index 00000000..e56e5b71
--- /dev/null
+++ b/ext4_utils/mkuserimg_mke2fs.py
@@ -0,0 +1,235 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2018 The Android Open Source Project
+#
+# 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.
+
+import argparse
+import logging
+import os
+import subprocess
+import sys
+
+
+def RunCommand(cmd, env):
+  """Runs the given command.
+
+  Args:
+    cmd: the command represented as a list of strings.
+    env: a dictionary of additional environment variables.
+  Returns:
+    A tuple of the output and the exit code.
+  """
+  env_copy = os.environ.copy()
+  env_copy.update(env)
+
+  logging.info("Env: %s", env)
+  logging.info("Running: " + " ".join(cmd))
+
+  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                       env=env_copy)
+  output, _ = p.communicate()
+
+  return output, p.returncode
+
+
+def ParseArguments(argv):
+  """Parses the input arguments to the program."""
+
+  parser = argparse.ArgumentParser(
+      description=__doc__,
+      formatter_class=argparse.RawDescriptionHelpFormatter)
+
+  parser.add_argument("src_dir", help="The source directory for user image.")
+  parser.add_argument("output_file", help="The path of the output image file.")
+  parser.add_argument("ext_variant", choices=["ext2", "ext4"],
+                      help="Variant of the extended filesystem.")
+  parser.add_argument("mount_point", help="The mount point for user image.")
+  parser.add_argument("fs_size", help="Size of the file system.")
+  parser.add_argument("file_contexts", nargs='?',
+                      help="The selinux file context.")
+
+  parser.add_argument("--android_sparse", "-s", action="store_true",
+                      help="Outputs an android sparse image (mke2fs).")
+  parser.add_argument("--journal_size", "-j",
+                      help="Journal size (mke2fs).")
+  parser.add_argument("--timestamp", "-T",
+                      help="Fake timetamp for the output image.")
+  parser.add_argument("--fs_config", "-C",
+                      help="Path to the fs config file (e2fsdroid).")
+  parser.add_argument("--product_out", "-D",
+                      help="Path to the directory with device specific fs"
+                           " config files (e2fsdroid).")
+  parser.add_argument("--block_list_file", "-B",
+                      help="Path to the block list file (e2fsdroid).")
+  parser.add_argument("--base_alloc_file_in", "-d",
+                      help="Path to the input base fs file (e2fsdroid).")
+  parser.add_argument("--base_alloc_file_out", "-A",
+                      help="Path to the output base fs file (e2fsdroid).")
+  parser.add_argument("--label", "-L",
+                      help="The mount point (mke2fs).")
+  parser.add_argument("--inodes", "-i",
+                      help="The extfs inodes count (mke2fs).")
+  parser.add_argument("--reserved_percent", "-M",
+                      help="The reserved blocks percentage (mke2fs).")
+  parser.add_argument("--flash_erase_block_size", "-e",
+                      help="The flash erase block size (mke2fs).")
+  parser.add_argument("--flash_logical_block_size", "-o",
+                      help="The flash logical block size (mke2fs).")
+  parser.add_argument("--mke2fs_uuid", "-U",
+                      help="The mke2fs uuid (mke2fs) .")
+  parser.add_argument("--mke2fs_hash_seed", "-S",
+                      help="The mke2fs hash seed (mke2fs).")
+  parser.add_argument("--share_dup_blocks", "-c", action="store_true",
+                      help="ext4 share dup blocks (e2fsdroid).")
+
+  args, remainder = parser.parse_known_args(argv)
+  # The current argparse doesn't handle intermixed arguments well. Checks
+  # manually whether the file_contexts exists as the last argument.
+  # TODO(xunchang) use parse_intermixed_args() when we switch to python 3.7.
+  if len(remainder) == 1 and remainder[0] == argv[-1]:
+    args.file_contexts = remainder[0]
+  elif remainder:
+    parser.print_usage()
+    sys.exit(1)
+
+  return args
+
+
+def ConstructE2fsCommands(args):
+  """Builds the mke2fs & e2fsdroid command based on the input arguments.
+
+  Args:
+    args: The result of ArgumentParser after parsing the command line arguments.
+  Returns:
+    A tuple of two lists that serve as the command for mke2fs and e2fsdroid.
+  """
+
+  BLOCKSIZE = 4096
+
+  e2fsdroid_opts = []
+  mke2fs_extended_opts = []
+  mke2fs_opts = []
+
+  if args.android_sparse:
+    mke2fs_extended_opts.append("android_sparse")
+  else:
+    e2fsdroid_opts.append("-e")
+  if args.timestamp:
+    e2fsdroid_opts += ["-T", args.timestamp]
+  if args.fs_config:
+    e2fsdroid_opts += ["-C", args.fs_config]
+  if args.product_out:
+    e2fsdroid_opts += ["-p", args.product_out]
+  if args.block_list_file:
+    e2fsdroid_opts += ["-B", args.block_list_file]
+  if args.base_alloc_file_in:
+    e2fsdroid_opts += ["-d", args.base_alloc_file_in]
+  if args.base_alloc_file_out:
+    e2fsdroid_opts += ["-D", args.base_alloc_file_out]
+  if args.share_dup_blocks:
+    e2fsdroid_opts.append("-s")
+  if args.file_contexts:
+    e2fsdroid_opts += ["-S", args.file_contexts]
+
+  if args.flash_erase_block_size:
+    mke2fs_extended_opts.append("stripe_width={}".format(
+        int(args.flash_erase_block_size) / BLOCKSIZE))
+  if args.flash_logical_block_size:
+    # stride should be the max of 8kb and the logical block size
+    stride = max(int(args.flash_logical_block_size), 8192)
+    mke2fs_extended_opts.append("stride={}".format(stride / BLOCKSIZE))
+  if args.mke2fs_hash_seed:
+    mke2fs_extended_opts.append("hash_seed=" + args.mke2fs_hash_seed)
+
+  if args.journal_size:
+    if args.journal_size == "0":
+      mke2fs_opts += ["-O", "^has_journal"]
+    else:
+      mke2fs_opts += ["-J", "size=" + args.journal_size]
+  if args.label:
+    mke2fs_opts += ["-L", args.label]
+  if args.inodes:
+    mke2fs_opts += ["-N", args.inodes]
+  if args.reserved_percent:
+    mke2fs_opts += ["-m", args.reserved_percent]
+  if args.mke2fs_uuid:
+    mke2fs_opts += ["-U", args.mke2fs_uuid]
+  if mke2fs_extended_opts:
+    mke2fs_opts += ["-E", ','.join(mke2fs_extended_opts)]
+
+  # Round down the filesystem length to be a multiple of the block size
+  blocks = int(args.fs_size) / BLOCKSIZE
+  mke2fs_cmd = (["mke2fs"] + mke2fs_opts +
+                ["-t", args.ext_variant, "-b", str(BLOCKSIZE), args.output_file,
+                 str(blocks)])
+
+  e2fsdroid_cmd = (["e2fsdroid"] + e2fsdroid_opts +
+                   ["-f", args.src_dir, "-a", args.mount_point,
+                    args.output_file])
+
+  return mke2fs_cmd, e2fsdroid_cmd
+
+
+def main(argv):
+  logging_format = '%(asctime)s %(filename)s %(levelname)s: %(message)s'
+  logging.basicConfig(level=logging.INFO, format=logging_format,
+                      datefmt='%H:%M:%S')
+
+  args = ParseArguments(argv)
+  if not os.path.isdir(args.src_dir):
+    logging.error("Can not find directory %s", args.src_dir)
+    sys.exit(2)
+  if not args.mount_point:
+    logging.error("Mount point is required")
+    sys.exit(2)
+  if args.mount_point[0] != '/':
+    args.mount_point = '/' + args.mount_point
+  if not args.fs_size:
+    logging.error("Size of the filesystem is required")
+    sys.exit(2)
+
+  mke2fs_cmd, e2fsdroid_cmd = ConstructE2fsCommands(args)
+
+  # truncate output file since mke2fs will keep verity section in existing file
+  with open(args.output_file, 'w') as output:
+    output.truncate()
+
+  # run mke2fs
+  mke2fs_env = {"MKE2FS_CONFIG" : "./system/extras/ext4_utils/mke2fs.conf"}
+  if args.timestamp:
+    mke2fs_env["E2FSPROGS_FAKE_TIME"] = args.timestamp
+
+  output, ret = RunCommand(mke2fs_cmd, mke2fs_env)
+  print(output)
+  if ret != 0:
+    logging.error("Failed to run mke2fs: " + output)
+    sys.exit(4)
+
+  # run e2fsdroid
+  e2fsdroid_env = {}
+  if args.timestamp:
+    e2fsdroid_env["E2FSPROGS_FAKE_TIME"] = args.timestamp
+
+  output, ret = RunCommand(e2fsdroid_cmd, e2fsdroid_env)
+  # The build script is parsing the raw output of e2fsdroid; keep the pattern
+  # unchanged for now.
+  print(output)
+  if ret != 0:
+    logging.error("Failed to run e2fsdroid_cmd: " + output)
+    os.remove(args.output_file)
+    sys.exit(4)
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])
diff --git a/ext4_utils/test_mkuserimg_mke2fs.py b/ext4_utils/test_mkuserimg_mke2fs.py
new file mode 100644
index 00000000..d5a68c5c
--- /dev/null
+++ b/ext4_utils/test_mkuserimg_mke2fs.py
@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2018 The Android Open Source Project
+#
+# 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.
+
+import unittest
+
+import mkuserimg_mke2fs
+
+class MkuserimgMke2fsTest(unittest.TestCase):
+  def setUp(self):
+    self.optional_arguments = {
+        "-j": "10", "-T": "1230768000.0", "-C": "fs_config",
+        "-D": "product_out", "-B": "block_list_file",
+        "-d": "base_alloc_file_in", "-A": "base_alloc_file_out",
+        "-L": "label", "-i": "20", "-M": "30", "-e": "8192",
+        "-o": "16384", "-U": "mke2fs_uuid", "-S": "mke2fs_hash_seed",
+    }
+
+  def test_parse_arguments_smoke(self):
+    args_list = ["source_directory", "output_file", "ext4", "data", "8192"]
+    for key, value in self.optional_arguments.items():
+      args_list += [key, value]
+    args_list.append("-c")
+
+    args = mkuserimg_mke2fs.ParseArguments(args_list)
+
+    self.assertEqual("source_directory", args.src_dir)
+    self.assertEqual("output_file", args.output_file)
+    self.assertEqual("ext4", args.ext_variant)
+    self.assertEqual("data", args.mount_point)
+    self.assertEqual("8192", args.fs_size)
+
+    self.assertFalse(args.android_sparse)
+    self.assertEqual("10", args.journal_size)
+    self.assertEqual("1230768000.0", args.timestamp)
+    self.assertEqual("fs_config", args.fs_config)
+    self.assertEqual("product_out", args.product_out)
+    self.assertEqual("block_list_file", args.block_list_file)
+    self.assertEqual("base_alloc_file_in", args.base_alloc_file_in)
+    self.assertEqual("base_alloc_file_out", args.base_alloc_file_out)
+    self.assertEqual("label", args.label)
+    self.assertEqual("20", args.inodes)
+    self.assertEqual("30", args.reserved_percent)
+    self.assertEqual("8192", args.flash_erase_block_size)
+    self.assertEqual("16384", args.flash_logical_block_size)
+    self.assertEqual("mke2fs_uuid", args.mke2fs_uuid)
+    self.assertEqual("mke2fs_hash_seed", args.mke2fs_hash_seed)
+    self.assertTrue(args.share_dup_blocks)
+
+  def test_parse_arguments_with_filecontext(self):
+    args_list = ["-s", "source_directory", "output_file", "ext4", "data",
+                 "8192"]
+    for key, value in self.optional_arguments.items():
+      args_list += [key, value]
+    args_list += ["-c", "file_contexts.bin"]
+
+    args = mkuserimg_mke2fs.ParseArguments(args_list)
+
+    self.assertEqual("file_contexts.bin", args.file_contexts)
+
+    self.assertEqual("source_directory", args.src_dir)
+    self.assertEqual("output_file", args.output_file)
+    self.assertEqual("ext4", args.ext_variant)
+    self.assertEqual("data", args.mount_point)
+    self.assertEqual("8192", args.fs_size)
+
+    self.assertTrue(args.android_sparse)
+    self.assertEqual("10", args.journal_size)
+    self.assertEqual("1230768000.0", args.timestamp)
+    self.assertEqual("fs_config", args.fs_config)
+    self.assertEqual("product_out", args.product_out)
+    self.assertEqual("block_list_file", args.block_list_file)
+    self.assertEqual("base_alloc_file_in", args.base_alloc_file_in)
+    self.assertEqual("base_alloc_file_out", args.base_alloc_file_out)
+    self.assertEqual("label", args.label)
+    self.assertEqual("20", args.inodes)
+    self.assertEqual("30", args.reserved_percent)
+    self.assertEqual("8192", args.flash_erase_block_size)
+    self.assertEqual("16384", args.flash_logical_block_size)
+    self.assertEqual("mke2fs_uuid", args.mke2fs_uuid)
+    self.assertEqual("mke2fs_hash_seed", args.mke2fs_hash_seed)
+    self.assertTrue(args.share_dup_blocks)
+
+  def test_parse_arguments_not_enough_arguments(self):
+    args_list = ["-s", "source_directory", "output_file", "ext4", "data",]
+    for key, value in self.optional_arguments.items():
+      args_list += [key, value]
+
+    with self.assertRaises(SystemExit):
+      mkuserimg_mke2fs.ParseArguments(args_list)
+
+  def test_construct_e2fs_opts_smoke(self):
+    args_list = ["-s", "source_directory", "output_file", "ext4", "data",
+                 "8192"]
+    for key, value in self.optional_arguments.items():
+      args_list += [key, value]
+    args_list += ["-c", "file_contexts.bin"]
+
+    args = mkuserimg_mke2fs.ParseArguments(args_list)
+
+    mke2fs_cmd, e2fsdroid_cmd = mkuserimg_mke2fs.ConstructE2fsCommands(
+        args)
+
+    expected_mke2fs_extended_opts = (
+        "android_sparse,stripe_width=2,stride=4,hash_seed={}".format(
+            args.mke2fs_hash_seed))
+    expected_mke2fs_cmd = [
+        "mke2fs", "-J", "size=10", "-L", args.label, "-N", args.inodes, "-m",
+        args.reserved_percent, "-U", args.mke2fs_uuid, "-E",
+        expected_mke2fs_extended_opts, "-t", args.ext_variant, "-b", "4096",
+        args.output_file, str(int(args.fs_size) / 4096)]
+
+    expected_e2fsdroid_cmd = [
+        "e2fsdroid", "-T", args.timestamp, "-C", args.fs_config, "-p",
+        args.product_out, "-B", args.block_list_file, "-d",
+        args.base_alloc_file_in, "-D", args.base_alloc_file_out, "-s", "-S",
+        args.file_contexts, "-f", args.src_dir, "-a", args.mount_point,
+        args.output_file]
+
+    self.assertEqual(' '.join(mke2fs_cmd), ' '.join(expected_mke2fs_cmd))
+
+    self.assertEqual(' '.join(e2fsdroid_cmd),
+                     ' '.join(expected_e2fsdroid_cmd))