build_identity: improve performance of _extract_dmg() usages

ipsw_parser/build_identity.py

@@ -14,22 +14,26 @@
 
 logger = logging.getLogger(__name__)
 
+AEA_MAGIC = b'AEA1'
 
-def _extract_dmg(buf: bytes, output: Path, sub_path: Optional[Path] = None, pem_db: Optional[str] = None) -> None:
+
+def _extract_dmg(dmg: Path, output: Path, sub_path: Optional[Path] = None, pem_db: Optional[str] = None) -> None:
     ipsw = local['ipsw']
     hdiutil = local['hdiutil']
     # darwin system statistically have problems cleaning up after detaching the mountpoint
     with TemporaryDirectory() as temp_dir:
         temp_dir = Path(temp_dir)
         mnt = temp_dir / 'mnt'
         mnt.mkdir()
-        dmg = temp_dir / 'image.dmg'
 
-        if buf.startswith(b'AEA1'):
+        with dmg.open('rb') as f:
+            magic = f.read(len(AEA_MAGIC))
+        if magic == AEA_MAGIC:
             logger.debug('Found Apple Encrypted Archive. Decrypting...')
-            dmg_aea = Path(str(dmg) + '.aea')
-            dmg_aea.write_bytes(buf)
-            args = ['fw', 'aea', dmg_aea, '-o', temp_dir]
+            aea_dmg = str(dmg.absolute()) + '.aea'
+            dmg.rename(aea_dmg)
+            dmg = temp_dir / Path(aea_dmg).name.rsplit('.', 1)[0]
+            args = ['fw', 'aea', aea_dmg, '-o', temp_dir]
             if pem_db is not None:
                 if '://' in pem_db:
                     # create a local file containing it
@@ -38,8 +42,6 @@ def _extract_dmg(buf: bytes, output: Path, sub_path: Optional[Path] = None, pem_
                     pem_db = temp_pem_db
                 args += ['--pem-db', pem_db]
             ipsw(args)
-        else:
-            dmg.write_bytes(buf)
 
         hdiutil('attach', '-mountpoint', mnt, dmg)
 
@@ -154,8 +156,11 @@ def extract_dsc(self, output: Path, pem_db: Optional[str] = None) -> None:
         device_support_symbols_path = output / 'private/preboot/Cryptexes/OS/System'
         device_support_symbols_path.mkdir(parents=True, exist_ok=True)
 
-        _extract_dmg(build_identity.get_component('Cryptex1,SystemOS').data, device_support_symbols_path,
-                     sub_path=Path('System'), pem_db=pem_db)
+        with TemporaryDirectory() as temp_dir:
+            system_os = Path(temp_dir) / 'system_os.dmg'
+            system_os.write_bytes(build_identity.get_component('Cryptex1,SystemOS').data)
+            _extract_dmg(system_os, device_support_symbols_path,
+                         sub_path=Path('System'), pem_db=pem_db)
         _split_dsc(output)
 
     def extract(self, output: Path, pem_db: Optional[str] = None) -> None:
@@ -164,7 +169,10 @@ def extract(self, output: Path, pem_db: Optional[str] = None) -> None:
         build_identity = self.build_manifest.build_identities[0]
 
         logger.info(f'extracting OS into: {output}')
-        _extract_dmg(build_identity.get_component('OS').data, output, pem_db=pem_db)
+        with TemporaryDirectory() as temp_dir:
+            os_dmg = Path(temp_dir) / 'os.dmg'
+            os_dmg.write_bytes(build_identity.get_component('OS').data)
+            _extract_dmg(os_dmg, output, pem_db=pem_db)
 
         kernel_component = build_identity.get_component('KernelCache')
         kernel_path = Path(kernel_component.path)
@@ -194,6 +202,9 @@ def extract(self, output: Path, pem_db: Optional[str] = None) -> None:
             cryptex_path.mkdir(parents=True, exist_ok=True)
 
             logger.info(f'extracting {name} into: {cryptex_path}')
-            _extract_dmg(build_identity.get_component(name).data, cryptex_path, pem_db=pem_db)
+            with TemporaryDirectory() as temp_dir:
+                component = (Path(temp_dir) / name).with_suffix('.dmg')
+                component.write_bytes(build_identity.get_component(name).data)
+                _extract_dmg(component, cryptex_path, pem_db=pem_db)
 
         _split_dsc(output)