chromite: Create cros_generate_quick_provision_payload script
Developers + external users need a way to generate quick provision
payloads without going through the build + artifacts flow.
BUG=b:217765909
TEST=run cros_generate_quick_provision_payload
Change-Id: Ieb1e42ed21d7aaec0f26ff2cb7b15229fd206ff0
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/3470856
Reviewed-by: Alex Klein <saklein@chromium.org>
Commit-Queue: Jae Hoon Kim <kimjae@chromium.org>
Tested-by: Jae Hoon Kim <kimjae@chromium.org>
diff --git a/scripts/cros_generate_quick_provision_payload.py b/scripts/cros_generate_quick_provision_payload.py
new file mode 100644
index 0000000..0d535f7
--- /dev/null
+++ b/scripts/cros_generate_quick_provision_payload.py
@@ -0,0 +1,74 @@
+# Copyright 2022 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Script to generate Chromium OS quick provision payloads."""
+
+import os
+
+from chromite.lib import commandline
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import osutils
+from chromite.lib import parallel
+from chromite.lib.paygen import partition_lib
+from chromite.lib.paygen import paygen_stateful_payload_lib
+
+
+def ParseArguments(argv):
+ """Returns a namespace for the CLI arguments."""
+ parser = commandline.ArgumentParser(description=__doc__)
+ parser.add_argument('--image', type='path', required=True,
+ help='The path to local image to build the quick '
+ 'provision payloads for.')
+ parser.add_argument('--output', type='path',
+ help='The output directory to generate quick '
+ 'provision payloads for.', default='.')
+
+ opts = parser.parse_args(argv)
+ # Check if output is valid directory.
+ if not os.path.isdir(opts.output):
+ parser.error('Please pass in a valid output directory.')
+
+ opts.Freeze()
+
+ return opts
+
+
+def CreateKernelQuickProvisionPayload(image, output):
+ with osutils.TempDir() as temp_dir:
+ # Extract kernel.
+ kern = os.path.join(temp_dir, 'kern')
+ partition_lib.ExtractKernel(image, os.path.join(temp_dir, kern))
+ # Compress kernel.
+ cros_build_lib.CompressFile(
+ kern,
+ os.path.join(output, constants.QUICK_PROVISION_PAYLOAD_KERNEL))
+
+
+def CreateRootQuickProvisionPayload(image, output):
+ with osutils.TempDir() as temp_dir:
+ # Extract root.
+ root = os.path.join(temp_dir, 'root')
+ partition_lib.ExtractRoot(image, os.path.join(temp_dir, root))
+ # Compress root.
+ cros_build_lib.CompressFile(
+ root,
+ os.path.join(output, constants.QUICK_PROVISION_PAYLOAD_ROOTFS))
+
+
+def CreateStatefulQuickProvisionPayload(image, output):
+ # Create stateful quick provision payload.
+ paygen_stateful_payload_lib.GenerateStatefulPayload(image, output)
+ # Change output ownership of file.
+
+
+def main(argv):
+ opts = ParseArguments(argv)
+
+ parallel.RunParallelSteps([
+ # Stateful generation is usually the slowest.
+ lambda: CreateStatefulQuickProvisionPayload(opts.image, opts.output),
+ lambda: CreateKernelQuickProvisionPayload(opts.image, opts.output),
+ lambda: CreateRootQuickProvisionPayload(opts.image, opts.output),
+ ])