Add main script

1 files changed, 63 insertions, 0 deletions

diff --git a/extract.py b/extract.py
new file mode 100755
index 0000000..266ef0a
--- /dev/null
+++ b/extract.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python2
+
+import argparse
+import errno
+import os
+
+import update_payload
+from update_payload import applier
+
+
+def list_content(payload_file_name):
+    with open(payload_file_name, 'rb') as payload_file:
+        payload = update_payload.Payload(payload_file)
+        payload.Init()
+
+        for part in payload.manifest.partitions:
+            print("{} ({} bytes)".format(part.partition_name,
+                                         part.new_partition_info.size))
+
+
+def extract(payload_file_name, output_dir="output", partition_names=None):
+    try:
+        os.makedirs(output_dir)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+    with open(payload_file_name, 'rb') as payload_file:
+        payload = update_payload.Payload(payload_file)
+        payload.Init()
+
+        if payload.IsDelta():
+            print("Delta payloads are not supported")
+            exit(1)
+
+        helper = applier.PayloadApplier(payload)
+        for part in payload.manifest.partitions:
+            if partition_names and part.partition_name not in partition_names:
+                continue
+            print("Extracting {}".format(part.partition_name))
+            output_file = os.path.join(output_dir, part.partition_name)
+            helper._ApplyToPartition(
+                part.operations, part.partition_name,
+                'install_operations', output_file,
+                part.new_partition_info)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("payload", metavar="payload.bin",
+                        help="Path to the payload.bin")
+    parser.add_argument("--output_dir", default="output",
+                        help="Output directory")
+    parser.add_argument("--partitions", type=str, nargs='+',
+                        help="Name of the partitions to extract")
+    parser.add_argument("--list_partitions", action="store_true",
+                        help="List the partitions included in the payload.bin")
+
+    args = parser.parse_args()
+    if args.list_partitions:
+        list_content(args.payload)
+    else:
+        extract(args.payload, args.output_dir, args.partitions)