extract_utils: add built-in extract fns to standalone extractor

Change-Id: I16bdb59566011703db9c6c1702738bd76debf7a1

Author: Demon000

extract.py

@@ -8,10 +8,21 @@
 
 from extract_utils.extract import (
     ExtractCtx,
+    extract_fns_type,
     extract_image,
     filter_already_extracted_partitions,
     get_dump_dir,
 )
+from extract_utils.extract_pixel import (
+    extract_pixel_factory_image,
+    extract_pixel_firmware,
+    pixel_factory_image_regex,
+    pixel_firmware_regex,
+)
+from extract_utils.extract_star import (
+    extract_star_firmware,
+    star_firmware_regex,
+)
 
 DEFAULT_EXTRACTED_PARTITIONS = [
     'odm',
@@ -41,6 +52,24 @@
     action='store_true',
     help='Extract all files from archive',
 )
+parser.add_argument(
+    '--pixel-factory',
+    nargs='*',
+    type=str,
+    help='Files to extract as pixel factory image',
+)
+parser.add_argument(
+    '--pixel-firmware',
+    nargs='*',
+    type=str,
+    help='Files to extract as pixel firmware',
+)
+parser.add_argument(
+    '--star-firmware',
+    nargs='*',
+    type=str,
+    help='Files to extract as star firmware',
+)
 
 parser.add_argument(
     'source',
@@ -52,13 +81,43 @@
 if __name__ == '__main__':
     args = parser.parse_args()
 
+    if args.pixel_factory is not None and not args.pixel_factory:
+        args.pixel_factory = [pixel_factory_image_regex]
+
+    if args.pixel_firmware is not None and not args.pixel_firmware:
+        args.pixel_firmware = [pixel_firmware_regex]
+
+    if args.star_firmware is not None and not args.star_firmware:
+        args.star_firmware = [star_firmware_regex]
+
+    extract_fns: extract_fns_type = {}
+
+    if args.pixel_factory:
+        for extract_pattern in args.pixel_factory:
+            extract_fns.setdefault(extract_pattern, []).append(
+                extract_pixel_factory_image,
+            )
+
+    if args.pixel_firmware:
+        for extract_pattern in args.pixel_firmware:
+            extract_fns.setdefault(extract_pattern, []).append(
+                extract_pixel_firmware,
+            )
+
+    if args.star_firmware:
+        for extract_pattern in args.star_firmware:
+            extract_fns.setdefault(extract_pattern, []).append(
+                extract_star_firmware,
+            )
+
     extract_partitions = args.partitions
     if args.extra_partitions is not None:
         extract_partitions += args.extra_partitions
 
     ctx = ExtractCtx(
         keep_dump=True,
         extract_partitions=extract_partitions,
+        extract_fns=extract_fns,
         extract_all=args.all,
     )