From 94c25844154b51e785b9e827cf98f7af5f86cbba Mon Sep 17 00:00:00 2001
From: doronz88 <doron88@gmail.com>
Date: Wed, 30 Oct 2024 14:11:28 +0200
Subject: [PATCH] cli: add `--pem-db` option

---
 ipsw_parser/__main__.py       |  5 +++--
 ipsw_parser/build_identity.py | 17 +++++++++++++----
 requirements.txt              |  1 +
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/ipsw_parser/__main__.py b/ipsw_parser/__main__.py
index 297e23a..b3693d2 100644
--- a/ipsw_parser/__main__.py
+++ b/ipsw_parser/__main__.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import logging
 from pathlib import Path
-from typing import IO
+from typing import IO, Optional
 from zipfile import ZipFile
 
 import click
@@ -47,7 +47,8 @@ def info(file) -> None:
 @cli.command('extract')
 @click.argument('file', type=click.Path(exists=True, file_okay=True, dir_okay=False))
 @click.argument('output', type=click.Path(exists=False))
-def extract(file: IO, output: str) -> None:
+@click.option('--pem-db', help='Path DB file url (can be either a filesystem path or an HTTP URL)')
+def extract(file: IO, output: str, pem_db: Optional[str]) -> None:
     """ Extract .ipsw into filesystem layout """
     output = Path(output)
 
diff --git a/ipsw_parser/build_identity.py b/ipsw_parser/build_identity.py
index dd95628..408d99e 100644
--- a/ipsw_parser/build_identity.py
+++ b/ipsw_parser/build_identity.py
@@ -5,6 +5,7 @@
 from tempfile import TemporaryDirectory
 from typing import Optional
 
+import requests
 from cached_property import cached_property
 from plumbum import ProcessExecutionError, local
 from pyimg4 import IM4P
@@ -14,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _extract_dmg(buf: bytes, output: Path, sub_path: Optional[Path] = None) -> None:
+def _extract_dmg(buf: bytes, 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
@@ -28,7 +29,15 @@ def _extract_dmg(buf: bytes, output: Path, sub_path: Optional[Path] = None) -> N
             logger.debug('Found Apple Encrypted Archive. Decrypting...')
             dmg_aea = Path(str(dmg) + '.aea')
             dmg_aea.write_bytes(buf)
-            ipsw('fw', 'aea', dmg_aea, '-o', temp_dir)
+            args = ['fw', 'aea', dmg_aea, '-o', temp_dir]
+            if pem_db is not None:
+                if '://' in pem_db:
+                    # create a local file containing it
+                    temp_pem_db = temp_dir / 'pem-db.json'
+                    temp_pem_db.write_text(requests.get(pem_db).text)
+                    pem_db = temp_pem_db
+                args += ['--pem-db', pem_db]
+            ipsw(args)
         else:
             dmg.write_bytes(buf)
 
@@ -149,13 +158,13 @@ def extract_dsc(self, output: Path) -> None:
                      sub_path=Path('System'))
         _split_dsc(output)
 
-    def extract(self, output: Path) -> None:
+    def extract(self, output: Path, pem_db: Optional[str] = None) -> None:
         logger.info(f'extracting into: {output}')
 
         build_identity = self.build_manifest.build_identities[0]
 
         logger.info(f'extracting OS into: {output}')
-        _extract_dmg(build_identity.get_component('OS').data, output)
+        _extract_dmg(build_identity.get_component('OS').data, output, pem_db=pem_db)
 
         kernel_component = build_identity.get_component('KernelCache')
         kernel_path = Path(kernel_component.path)
diff --git a/requirements.txt b/requirements.txt
index 5d59b93..2fe8ee8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,4 @@ coloredlogs
 cached_property
 plumbum
 pyimg4>=0.8.6
+requests
\ No newline at end of file