From fc992da255debdceb69a80f338e4caf496dd629a Mon Sep 17 00:00:00 2001
From: DoronZ <doron88@gmail.com>
Date: Wed, 27 Nov 2024 00:58:43 +0200
Subject: [PATCH] add remotezip support together with extract-kernel

---
 ipsw_parser/__main__.py       | 36 +++++++++++++++++++++++++----------
 ipsw_parser/build_identity.py | 12 ++++++++++++
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/ipsw_parser/__main__.py b/ipsw_parser/__main__.py
index 0e6cf6a..1c8aa9f 100644
--- a/ipsw_parser/__main__.py
+++ b/ipsw_parser/__main__.py
@@ -1,11 +1,12 @@
 #!/usr/bin/env python3
 import logging
 from pathlib import Path
-from typing import IO, Optional
+from typing import Optional
 from zipfile import ZipFile
 
 import click
 import coloredlogs
+from remotezip import RemoteZip
 
 from ipsw_parser.ipsw import IPSW
 
@@ -23,6 +24,14 @@
 
 PEM_DB_ENV_VAR = 'IPSW_PARSER_PEM_DB'
 
+
+def handle_ipsw_argument(ctx: click.Context, param: click.Argument, value: str) -> IPSW:
+    if value.startswith('http://') or value.startswith('https://'):
+        return IPSW(RemoteZip(value))
+    return IPSW(ZipFile(Path(value).expanduser()))
+
+
+ipsw_argument = click.argument('ipsw', callback=handle_ipsw_argument)
 pem_db_option = click.option('--pem-db', envvar=PEM_DB_ENV_VAR,
                              help='Path DB file url (can be either a filesystem path or an HTTP URL). '
                                   'Alternatively, use the IPSW_PARSER_PEM_DB envvar.')
@@ -35,10 +44,9 @@ def cli() -> None:
 
 
 @cli.command('info')
-@click.argument('file', type=click.Path(exists=True, file_okay=True, dir_okay=False))
-def info(file) -> None:
+@ipsw_argument
+def info(ipsw) -> None:
     """ Parse given .ipsw basic info """
-    ipsw = IPSW(ZipFile(file))
     print(f'SupportedProductTypes: {ipsw.build_manifest.supported_product_types}')
     print(f'ProductVersion: {ipsw.build_manifest.product_version}')
     print(f'ProductBuildVersion: {ipsw.build_manifest.product_build_version}')
@@ -51,28 +59,36 @@ def info(file) -> None:
 
 
 @cli.command('extract')
-@click.argument('file', type=click.Path(exists=True, file_okay=True, dir_okay=False))
+@ipsw_argument
 @click.argument('output', type=click.Path(exists=False))
 @pem_db_option
-def extract(file: IO, output: str, pem_db: Optional[str]) -> None:
+def extract(ipsw: IPSW, output: str, pem_db: Optional[str]) -> None:
     """ Extract .ipsw into filesystem layout """
     output = Path(output)
 
     if not output.exists():
         output.mkdir(parents=True, exist_ok=True)
 
-    ipsw = IPSW(ZipFile(file))
     ipsw.build_manifest.build_identities[0].extract(output, pem_db=pem_db)
     ipsw.archive.extractall(
         path=output, members=[f for f in ipsw.archive.filelist if f.filename.startswith('Firmware')])
 
 
+@cli.command('extract-kernel')
+@ipsw_argument
+@click.argument('output', type=click.Path(exists=False))
+@click.option('--arch', help='Arch name to extract using lipo')
+def extract_kernel(ipsw: IPSW, output: str, arch: Optional[str]) -> None:
+    """ Extract kernelcache from given .ipsw into given output filename """
+    Path(output).write_bytes(ipsw.build_manifest.build_identities[0].get_kernelcache_payload(arch=arch))
+
+
 @cli.command('device-support')
-@click.argument('file', type=click.Path(exists=True, file_okay=True, dir_okay=False))
+@ipsw_argument
 @pem_db_option
-def device_support(file: IO, pem_db: Optional[str]) -> None:
+def device_support(ipsw: IPSW, pem_db: Optional[str]) -> None:
     """ Create DeviceSupport directory """
-    IPSW(ZipFile(file)).create_device_support(pem_db=pem_db)
+    ipsw.create_device_support(pem_db=pem_db)
 
 
 if __name__ == '__main__':
diff --git a/ipsw_parser/build_identity.py b/ipsw_parser/build_identity.py
index 7ff9420..c7488cb 100644
--- a/ipsw_parser/build_identity.py
+++ b/ipsw_parser/build_identity.py
@@ -158,6 +158,18 @@ def extract_dsc(self, output: Path, pem_db: Optional[str] = None) -> None:
                      sub_path=Path('System'), pem_db=pem_db)
         _split_dsc(output)
 
+    def get_kernelcache_payload(self, arch: Optional[str] = None) -> bytes:
+        im4p = IM4P(self.build_manifest.build_identities[0].get_component('KernelCache').data)
+        im4p.payload.decompress()
+        payload = im4p.payload.output().data
+        if arch is None:
+            return payload
+
+        with TemporaryDirectory() as temp_dir:
+            kernel_output = Path(temp_dir) / 'kernel'
+            local['ipsw']('macho', 'lipo', '-a', arch, kernel_output)
+            return Path(next(kernel_output.parent.glob(f'*.{arch}'))).read_bytes()
+
     def extract(self, output: Path, pem_db: Optional[str] = None) -> None:
         logger.info(f'extracting into: {output}')