Allow flashing all partitions with this tool

.gitignore

@@ -0,0 +1 @@
+__pycache__

fastbootflasher.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3.7
+
+from adb.fastboot import FastbootCommands, FastbootRemoteFailure, FastbootStateMismatch, FastbootInvalidResponse
+from adb.adb_commands import AdbCommands
+from adb.usb_exceptions import *
+import time
+
+def info_cb(msg):
+    if msg.header == b"FAIL":
+        print(msg)
+        raise RuntimeException("Flash Failed")
+    print(msg)
+
+def progress_cb(current, total):
+    print("flashing current img "+str(100*current/total))
+
+def flash(parts):
+    fdev = FastbootCommands()
+    while True:
+        try:
+            fdev.ConnectDevice()
+            break
+        except DeviceNotFoundError:
+            time.sleep(10)
+    # For faster flashing we take advantage of "ultraflash" which lets you stream the image
+    for partition, file in parts.items():
+        # File must be a filename, not a file()
+        fdev._SimpleCommand(b'ultraflash', arg=partition, info_cb=info_cb, timeout_ms=0)
+        fdev.Download(file, info_cb=info_cb, progress_callback=progress_cb)

idtconfig.py

@@ -15,16 +15,18 @@ def get_images(cfg: bytes):
     try:
         tree = etree.fromstring(cfg)
     except:
-        return get_simple(cfg)
+        raise
+        return get_simple(cfg.decode("utf-8")), {}
     ddr_images = tree.xpath("(//configurations/configuration)[1]/bootloaderimage_ddr/image")
     std_images = tree.xpath("(//configurations/configuration)[1]/bootloaderimage/image")
-    images = {int(img.get("address"), 0):img.text for img in ddr_images}
+    idt_images = {int(img.get("address"), 0):img.text for img in ddr_images}
     ddr_ids = [img.get("identifier") for img in ddr_images]
     for img in std_images:
         if not img.get("identifier") in ddr_ids:
-            images[int(img.get("address"), 0)] = img.text
-    print(images)
-    return images
+            idt_images[int(img.get("address"), 0)] = img.text
+    print(idt_images)
+    fastboot_images = {img.get("identifier", 0):img.text for img in tree.xpath("(//configurations/configuration)[1]/fastbootimage/image")}
+    return idt_images, fastboot_images
 
 def get_simple(cfg):
     ret = {}

main.py

@@ -9,20 +9,26 @@
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import imageflasher
+import fastbootflasher
 import idtconfig
 
 import argparse, os
 
-def main(config, device, chip=""):
+def main(config, device, full, chip=""):
     # First parse the config
-    images = idtconfig.get_images(config.read())
+    idt_images, fastboot_images = idtconfig.get_images(config.read())
     flasher = imageflasher.ImageFlasher(chip)
     if device != False: # We have to check not False rather than just True, because None evaluates to False, and None should be passed intact
         flasher.connect_serial(device)
-    for addr, fil in images.items():
+    for addr, fil in idt_images.items():
         flasher.download_from_disk(os.path.join(os.path.dirname(config.name), fil.replace("/", os.path.sep)), addr)
     print("Flash successful!")
-    print("Wait around 5 minutes with the USB plugged in and the device will enter fastboot. Use fastboot to reflash/recover your system.")
+    if full:
+        print(idt_images, fastboot_images)
+        print("Waiting for device to connect. Please do not disconnect the USB or reboot the device.")
+        fastbootflasher.flash(fastboot_images)
+    else:
+        print("Wait around 5 minutes with the USB plugged in and the device will enter fastboot. Use fastboot to reflash/recover your system.")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(epilog="""Copyright 2019 Penn Mackintosh
@@ -34,7 +40,8 @@ def main(config, device, chip=""):
 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.""")
     parser.add_argument("--norun", "-n", action="store_false", dest="run")
     parser.add_argument("--device", "-d")
+    parser.add_argument("--fastboot", "--full", "-f", help="Flashes all fastboot images from the IDT config as well. Not supported with hikey config", action="store_true", default=False, dest="full")
     parser.add_argument("config")
     args = parser.parse_args()
-    with open(args.config, "r") as f:
-        main(f, args.device if args.run else False)
+    with open(args.config, "rb") as f:
+        main(f, args.device if args.run else False, args.full)