Minimal recut: ordered multi-URL dump input + pass-through

dumpyarabot/arq_jobs.py

@@ -33,6 +33,7 @@
 from dumpyarabot.process_utils import reset_current_job_id, set_current_job_id
 from dumpyarabot.property_extractor import PropertyExtractor
 from dumpyarabot.schemas import DumpJob
+from dumpyarabot import url_utils
 
 console = Console()
 
@@ -512,7 +513,6 @@ async def process_firmware_dump(ctx, job_data: Dict[str, Any]) -> Dict[str, Any]
             try:
                 # Initialize components (exact same as original)
                 await _raise_if_job_cancel_requested(job_id)
-                downloader = FirmwareDownloader(str(work_dir))
                 extractor = FirmwareExtractor(str(work_dir))
                 prop_extractor = PropertyExtractor(str(work_dir))
                 gitlab_manager = GitLabManager(str(work_dir))
@@ -523,8 +523,6 @@ async def process_firmware_dump(ctx, job_data: Dict[str, Any]) -> Dict[str, Any]
                 # Step 2: GitLab access validation (8%)
                 await update_progress_with_metadata(job_data, " Validating GitLab access...", 8.0)
                 await _validate_gitlab_access()
-                is_whitelisted = await gitlab_manager.check_whitelist(str(job_data["dump_args"]["url"]))
-
                 # Step 3: URL optimization and mirror selection (12%)
                 await update_progress_with_metadata(job_data, " Optimizing download URL and selecting mirrors...", 12.0)
 
@@ -533,23 +531,19 @@ async def process_firmware_dump(ctx, job_data: Dict[str, Any]) -> Dict[str, Any]
 
                 # Create DumpJob object for components that need it
                 dump_job = DumpJob.model_validate(job_data)
+                ordered_urls = [
+                    str(dump_job.dump_args.url),
+                    *[str(url) for url in dump_job.dump_args.delta_urls],
+                ]
+                if dump_job.dump_args.use_alt_dumper and len(ordered_urls) > 1:
+                    raise RuntimeError("Alternative dumper cannot be used with delta OTA chains")
+                is_whitelisted = all(
+                    url_utils.is_whitelisted_url(url) for url in ordered_urls
+                )
+                firmware_paths = []
 
                 # Download with live progress via aria2 RPC callback.
                 # Download progress is mapped into the 15%-50% band of overall job progress.
-                async def _on_download_progress(dp: DownloadProgress) -> None:
-                    dl_pct = dp.percentage  # 0-100 within download
-                    overall_pct = 15.0 + (dl_pct / 100.0) * 35.0  # map to 15%-50%
-                    dl_info = format_download_progress(dp)
-                    step_msg = f" Downloading firmware...\n{dl_info}"
-
-                    progress_data = {
-                        "current_step": step_msg,
-                        "percentage": overall_pct,
-                        "current_step_number": 4,
-                        "total_steps": 25,
-                    }
-                    await _send_status_update(job_data, step_msg, progress_data, job_data.get("metadata"))
-
                 # Use PeriodicTimerUpdate as a fallback for downloaders without
                 # live progress (Google Drive, MediaFire, MEGA, wget fallback).
                 # When aria2 RPC is active, the callback above sends updates instead.
@@ -560,9 +554,36 @@ async def _on_download_progress(dp: DownloadProgress) -> None:
                     "percentage": 15.0,
                 }
                 async with PeriodicTimerUpdate(job_data, " Downloading firmware...", download_progress):
-                    firmware_path, firmware_name = await downloader.download_firmware(
-                        dump_job, on_progress=_on_download_progress
-                    )
+                    for index, url in enumerate(ordered_urls):
+                        await _raise_if_job_cancel_requested(job_id)
+                        downloader = FirmwareDownloader(str(work_dir / "downloads" / f"{index:03d}"))
+
+                        async def _on_download_progress(
+                            dp: DownloadProgress, stage: int = index
+                        ) -> None:
+                            stage_progress = (stage + dp.percentage / 100.0) / len(ordered_urls)
+                            overall_pct = 15.0 + stage_progress * 35.0
+                            if len(ordered_urls) == 1:
+                                step_msg = f" Downloading firmware...\n{format_download_progress(dp)}"
+                            else:
+                                step_msg = (
+                                    f" Downloading firmware {stage + 1}/{len(ordered_urls)}...\n"
+                                    f"{format_download_progress(dp)}"
+                                )
+                            progress_data = {
+                                "current_step": step_msg,
+                                "percentage": overall_pct,
+                                "current_step_number": 4,
+                                "total_steps": 25,
+                            }
+                            await _send_status_update(
+                                job_data, step_msg, progress_data, job_data.get("metadata")
+                            )
+
+                        firmware_path, _ = await downloader.download_firmware(
+                            dump_job, on_progress=_on_download_progress, url=url
+                        )
+                        firmware_paths.append(firmware_path)
 
                 # Step 5: Download completed (50%)
                 await update_progress_with_metadata(job_data, " Firmware download completed", 50.0)
@@ -572,7 +593,10 @@ async def _on_download_progress(dp: DownloadProgress) -> None:
 
                 # Use periodic timer for extraction operation
                 async with PeriodicTimerUpdate(job_data, " Extracting firmware partitions...", {"current_step": "Extract", "total_steps": 25, "current_step_number": 6, "percentage": 52.0}):
-                    await extractor.extract_firmware(dump_job, firmware_path)
+                    if dump_job.dump_args.delta_urls:
+                        await extractor.extract_delta_chain(dump_job, firmware_paths)
+                    else:
+                        await extractor.extract_firmware(dump_job, firmware_paths[0])
 
                 # Step 7: Firmware extraction completed (56%)
                 await update_progress_with_metadata(job_data, " Firmware extraction completed", 56.0)

dumpyarabot/firmware_downloader.py

@@ -30,6 +30,7 @@ async def download_firmware(
         self,
         job: DumpJob,
         on_progress: ProgressCallback | None = None,
+        url: str | None = None,
     ) -> Tuple[str, str]:
         """Download firmware and return (file_path, file_name).
 
@@ -38,7 +39,7 @@ async def download_firmware(
             on_progress: Optional async callback invoked with each DownloadProgress
                          snapshot during aria2 RPC downloads.
         """
-        url = str(job.dump_args.url)
+        url = url or str(job.dump_args.url)
 
         # Check if it's a local file
         if os.path.isfile(url):
@@ -262,4 +263,3 @@ async def _download_default(
 
         return str(latest_file)
 
-

dumpyarabot/firmware_extractor.py

@@ -1,8 +1,13 @@
 import asyncio
+import shutil
+from collections.abc import Sequence
 from pathlib import Path
+from typing import Callable, TypeVar
 
 from dumpyara.dumpyara import dumpyara
+from dumpyara.steps.extract_images import extract_images as dumpyara_extract_images
 from dumpyara.utils import multipartitions as dumpyara_multipartitions
+import otadump
 from rich.console import Console
 
 from dumpyarabot.file_utils import (
@@ -21,20 +26,156 @@
 from dumpyarabot.schemas import DumpJob
 
 console = Console()
+_T = TypeVar("_T")
 
 def _run_dumpyara(firmware_path: Path, output_path: Path) -> None:
     """Run Dumpyara with its reliable in-process payload parser."""
     dumpyara_multipartitions.OTADUMP_EXECUTABLE = None
     dumpyara(firmware_path, output_path)
 
 
+def _run_dumpyara_images(images_path: Path, output_path: Path) -> None:
+    dumpyara_extract_images(images_path, output_path)
+
+
+def _carry_forward_images(source_dir: Path, output_dir: Path) -> None:
+    for source_image in source_dir.glob("*.img"):
+        destination = output_dir / source_image.name
+        if not destination.exists():
+            shutil.copy2(source_image, destination)
+
+
+def _delta_stage_has_images(stage_dir: Path) -> bool:
+    return any(stage_dir.glob("*.img"))
+
+
+def _publish_stage_children(staging_dir: Path, destination_dir: Path) -> None:
+    children = list(staging_dir.iterdir())
+    if not children:
+        raise RuntimeError("delta chain final stage produced no files")
+    preexisting = [destination_dir / child.name for child in children if (destination_dir / child.name).exists()]
+    if preexisting:
+        raise RuntimeError(f"final output already exists: {preexisting[0]}")
+
+    moved: list[Path] = []
+    try:
+        for child in children:
+            target = destination_dir / child.name
+            shutil.move(str(child), str(target))
+            moved.append(target)
+    except BaseException:
+        for moved_child in moved:
+            if moved_child.exists():
+                if moved_child.is_dir():
+                    shutil.rmtree(moved_child)
+                else:
+                    moved_child.unlink()
+        raise
+
+
+async def _run_thread(
+    function: Callable[..., _T],
+    *args: object,
+    timeout: float,
+    token: otadump.CancellationToken,
+    **kwargs: object,
+) -> _T:
+    task = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs))
+    try:
+        return await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
+    except (asyncio.TimeoutError, asyncio.CancelledError):
+        token.cancel()
+        while True:
+            try:
+                await asyncio.shield(task)
+                break
+            except asyncio.CancelledError:
+                continue
+            except Exception:
+                break
+        raise
+
+
 class FirmwareExtractor:
     """Handles firmware extraction using both Python dumper and alternative methods."""
 
     def __init__(self, work_dir: str):
         self.work_dir = Path(work_dir)
         self.firmware_extractor_path = Path.home() / "Firmware_extractor"
 
+    async def extract_delta_chain(self, job: DumpJob, firmware_paths: Sequence[str]) -> str:
+        if len(firmware_paths) < 2:
+            raise ValueError("A delta chain requires a base and at least one delta OTA")
+        if job.dump_args.use_alt_dumper:
+            raise ValueError("Alternative dumper is not supported for delta OTA chains")
+
+        staging_root = self.work_dir / ".ota_staging"
+        shutil.rmtree(staging_root, ignore_errors=True)
+        staging_root.mkdir(parents=True, exist_ok=True)
+        cancellation_token = otadump.CancellationToken()
+
+        try:
+            current_stage = staging_root / "stage_000"
+            await _run_thread(
+                otadump.extract,
+                Path(firmware_paths[0]),
+                current_stage,
+                timeout=ONE_HOUR,
+                token=cancellation_token,
+                cancellation_token=cancellation_token,
+            )
+            if not _delta_stage_has_images(current_stage):
+                raise RuntimeError(
+                    f"delta stage produced no images: {Path(firmware_paths[0]).name}"
+                )
+            for index, firmware_path in enumerate(firmware_paths[1:], start=1):
+                next_stage = staging_root / f"stage_{index:03d}"
+                await _run_thread(
+                    otadump.extract,
+                    Path(firmware_path),
+                    next_stage,
+                    timeout=ONE_HOUR,
+                    token=cancellation_token,
+                    source_dir=current_stage,
+                    cancellation_token=cancellation_token,
+                )
+                if not _delta_stage_has_images(next_stage):
+                    raise RuntimeError(
+                        f"delta stage produced no images: {Path(firmware_path).name}"
+                    )
+                await _run_thread(
+                    _carry_forward_images,
+                    current_stage,
+                    next_stage,
+                    timeout=ONE_HOUR,
+                    token=cancellation_token,
+                )
+                shutil.rmtree(current_stage)
+                current_stage = next_stage
+
+            final_output = staging_root / "final"
+            final_output.mkdir(parents=True, exist_ok=True)
+            await _run_thread(
+                _run_dumpyara_images,
+                current_stage,
+                final_output,
+                timeout=ONE_HOUR,
+                token=cancellation_token,
+            )
+            system_path = final_output / "system"
+            if not system_path.exists():
+                raise RuntimeError("dumpyara output missing required non-empty system")
+            if system_path.is_file() and system_path.stat().st_size == 0:
+                raise RuntimeError("dumpyara output missing required non-empty system")
+            if system_path.is_dir() and not any(system_path.iterdir()):
+                raise RuntimeError("dumpyara output missing required non-empty system")
+            _publish_stage_children(final_output, self.work_dir)
+            return str(self.work_dir)
+        finally:
+            shutil.rmtree(staging_root, ignore_errors=True)
+            for firmware_path in firmware_paths:
+                safe_remove_file(firmware_path)
+
     async def extract_firmware(self, job: DumpJob, firmware_path: str) -> str:
         """Extract firmware and return extraction directory."""
         console.print(f"[blue]Extracting firmware: {firmware_path}[/blue]")

dumpyarabot/gitlab_manager.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from typing import Any, Dict, Tuple
-from urllib.parse import urlparse
 
 import httpx
 from rich.console import Console
@@ -318,30 +317,3 @@ async def send_channel_notification(
                 console.print("[green]Channel notification sent successfully[/green]")
             else:
                 console.print(f"[yellow]Failed to send channel notification: {response.text}[/yellow]")
-
-    async def check_whitelist(self, url: str) -> bool:
-        """Check if URL is in whitelist."""
-        whitelist_file = Path.home() / "dumpbot" / "whitelist.txt"
-
-        if not whitelist_file.exists():
-            console.print("[yellow]Whitelist file not found[/yellow]")
-            return False
-
-        try:
-            with open(whitelist_file, 'r') as f:
-                whitelist_domains = [line.strip() for line in f if line.strip()]
-
-            hostname = (urlparse(url).hostname or "").lower()
-
-            for domain in whitelist_domains:
-                normalized_domain = domain.lower()
-                if hostname == normalized_domain or hostname.endswith(f".{normalized_domain}"):
-                    console.print(f"[green]URL is whitelisted (domain: {domain})[/green]")
-                    return True
-
-            console.print("[yellow]URL is not whitelisted[/yellow]")
-            return False
-
-        except Exception as e:
-            console.print(f"[red]Error checking whitelist: {e}[/red]")
-            return False