Recover complete files from truncated sysdiagnose archives

src/mvt/ios/cmd_check_sysdiagnose.py

@@ -3,6 +3,7 @@
 # Use of this software is governed by the MVT License 1.1 that can be found at
 #   https://license.mvt.re/1.1/
 
+import gzip
 import json
 import logging
 import os
@@ -11,7 +12,7 @@
 import tarfile
 import zlib
 from pathlib import Path, PurePosixPath
-from tempfile import TemporaryDirectory
+from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import Any, Optional
 
 from mvt.common.command import Command
@@ -23,6 +24,34 @@
 log = logging.getLogger(__name__)
 
 
+class _StreamingGzipFile(gzip.GzipFile):
+    """Deliver available bytes before reporting a missing gzip trailer.
+
+    GzipFile.read(n) tries to fill the entire request. At a truncated end it
+    can raise EOFError even after decoding bytes for complete tar members.
+    read1() returns those bytes first and raises on the following read. The
+    tar stream accepts short reads, so it can retain every complete member.
+    """
+
+    def read(self, size: Optional[int] = -1) -> bytes:
+        return self.read1(-1 if size is None else size)
+
+
+class _SysdiagnoseTarInfo(tarfile.TarInfo):
+    @classmethod
+    def fromtarfile(cls, archive: tarfile.TarFile) -> "_SysdiagnoseTarInfo":
+        try:
+            return super().fromtarfile(archive)
+        except (
+            tarfile.EmptyHeaderError,  # type: ignore[attr-defined]
+            tarfile.TruncatedHeaderError,  # type: ignore[attr-defined]
+            tarfile.InvalidHeaderError,  # type: ignore[attr-defined]
+        ) as exc:
+            # TarFile otherwise treats a short/invalid header after its first
+            # member as an ordinary end of archive, hiding the truncation.
+            raise tarfile.ReadError(str(exc)) from exc
+
+
 class CmdIOSCheckSysdiagnose(Command):
     def __init__(
         self,
@@ -63,6 +92,8 @@ def __init__(
         self.ips_files: list[dict[str, Any]] = []
         self.temp_sysdiagnose_dir: Optional[TemporaryDirectory[str]] = None
         self.extracted_sysdiagnose_path: Optional[str] = None
+        self.sysdiagnose_tar_members: list[tarfile.TarInfo] = []
+        self.archive_incomplete = False
 
     @staticmethod
     def _parse_bugtype_header(data: bytes) -> Optional[int]:
@@ -104,82 +135,140 @@ def init(self) -> None:
         self.log.info("Parsing sysdiagnose archive. This might take a while...")
         self.sysdiagnose_format = "tar"
         try:
-            self.sysdiagnose_archive = tarfile.open(self.target_path, "r:gz")
-            self._extract_sysdiagnose_archive()
+            with _StreamingGzipFile(self.target_path, "rb") as compressed:
+                self.sysdiagnose_archive = tarfile.open(
+                    fileobj=compressed, mode="r|", tarinfo=_SysdiagnoseTarInfo
+                )
+                self._extract_sysdiagnose_archive()
+                # Tar stops at its end marker, which can precede the gzip
+                # trailer. Drain the compressed stream to check its EOF/CRC.
+                if not self.archive_incomplete:
+                    try:
+                        while compressed.read1(1024 * 1024):
+                            pass
+                    except (EOFError, gzip.BadGzipFile, zlib.error) as exc:
+                        self._warn_incomplete_archive(exc)
         except (tarfile.ReadError, EOFError, zlib.error, OSError) as exc:
-            # A truncated archive ends in EOFError from gzip, which Click would
-            # otherwise report as a bare "Aborted!" with no reason.
             self.log.critical(
                 "Unable to read the sysdiagnose archive %s: %s. "
                 "The file may be truncated or not a gzip-compressed tarball.",
                 self.target_path,
                 exc,
             )
+            self.finish()
             sys.exit(1)
 
+    def _warn_incomplete_archive(self, exc: Exception) -> None:
+        self.archive_incomplete = True
+        self.log.warning(
+            "The sysdiagnose archive is truncated or damaged: %s. Recovered "
+            "%d complete files; incomplete files are excluded. Analysis will "
+            "be partial and archive integrity cannot be verified. Obtain a "
+            "complete copy for a full analysis.",
+            exc,
+            len(self.sysdiagnose_files),
+        )
+
     def _extract_sysdiagnose_archive(self) -> None:
         archive = self.sysdiagnose_archive
         if archive is None:
             raise RuntimeError("Sysdiagnose archive has not been initialized")
 
         self.temp_sysdiagnose_dir = TemporaryDirectory()
         extraction_root = Path(self.temp_sysdiagnose_dir.name).resolve()
-        archive_roots = set()
-
-        for member in archive:
-            member_path = PurePosixPath(member.name.replace("\\", "/"))
-            if member_path.is_absolute() or ".." in member_path.parts:
-                self.log.warning("Skipping unsafe sysdiagnose path %r", member.name)
-                continue
-
-            destination = extraction_root.joinpath(*member_path.parts).resolve()
-            if not destination.is_relative_to(extraction_root):
-                self.log.warning("Skipping unsafe sysdiagnose path %r", member.name)
-                continue
-
-            if not member_path.parts:
-                continue
-            # AppleDouble sidecars (._name) carry a file's extended attributes,
-            # not sysdiagnose content. Device archives hold hundreds of them;
-            # bsdtar hides them from listings, tarfile does not.
-            if member_path.name.startswith("._"):
-                continue
-            archive_roots.add(member_path.parts[0])
-
-            if member.isdir():
-                destination.mkdir(parents=True, exist_ok=True)
-                continue
-
-            # Modules only need directories and regular files. Do not materialize
-            # links or device nodes from an untrusted sysdiagnose archive.
-            if not member.isfile():
-                self.log.warning("Skipping unsafe sysdiagnose member %r", member.name)
-                continue
-
-            normalized_name = member_path.as_posix()
-            self.sysdiagnose_files.append(normalized_name)
-
-            source = archive.extractfile(member)
-            if source is None:
-                continue
-
-            destination.parent.mkdir(parents=True, exist_ok=True)
-            with source, destination.open("wb") as output:
-                shutil.copyfileobj(source, output)
-
-            if normalized_name.endswith(".ips"):
-                self._add_ips_file(str(destination), destination.read_bytes())
+        archive_roots: set[str] = set()
+
+        try:
+            for member in archive:
+                self._extract_member(archive, member, extraction_root, archive_roots)
+        except (tarfile.ReadError, EOFError, gzip.BadGzipFile, zlib.error) as exc:
+            self._warn_incomplete_archive(exc)
+            if not self.sysdiagnose_files:
+                raise tarfile.ReadError(
+                    "No complete sysdiagnose files could be recovered"
+                ) from exc
 
         if len(archive_roots) != 1:
+            self.finish()
             raise ValueError("Sysdiagnose archive must contain one top-level directory")
 
         self.extracted_sysdiagnose_path = str(extraction_root / archive_roots.pop())
 
+    def _extract_member(
+        self,
+        archive: tarfile.TarFile,
+        member: tarfile.TarInfo,
+        extraction_root: Path,
+        archive_roots: set[str],
+    ) -> None:
+        member_path = PurePosixPath(member.name.replace("\\", "/"))
+        if member_path.is_absolute() or ".." in member_path.parts:
+            self.log.warning("Skipping unsafe sysdiagnose path %r", member.name)
+            return
+
+        destination = extraction_root.joinpath(*member_path.parts).resolve()
+        if not destination.is_relative_to(extraction_root):
+            self.log.warning("Skipping unsafe sysdiagnose path %r", member.name)
+            return
+
+        if not member_path.parts:
+            return
+        # AppleDouble sidecars (._name) carry a file's extended attributes,
+        # not sysdiagnose content. Device archives hold hundreds of them;
+        # bsdtar hides them from listings, tarfile does not.
+        if member_path.name.startswith("._"):
+            return
+        archive_roots.add(member_path.parts[0])
+
+        if member.isdir():
+            destination.mkdir(parents=True, exist_ok=True)
+            self.sysdiagnose_tar_members.append(member)
+            return
+
+        # Modules only need directories and regular files. Do not materialize
+        # links or device nodes from an untrusted sysdiagnose archive.
+        if not member.isfile():
+            self.log.warning("Skipping unsafe sysdiagnose member %r", member.name)
+            return
+
+        normalized_name = member_path.as_posix()
+
+        source = archive.extractfile(member)
+        if source is None:
+            return
+
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        # Stage each file independently. A truncated member must never be
+        # indexed or left for a module that scans the extracted directory.
+        # Staging also preserves an earlier complete entry with the same name.
+        with NamedTemporaryFile(dir=destination.parent, delete=False) as output:
+            staged_path = Path(output.name)
+            try:
+                with source:
+                    shutil.copyfileobj(source, output)
+                if output.tell() != member.size:
+                    raise tarfile.ReadError("Unexpected end of sysdiagnose member")
+            except BaseException:
+                output.close()
+                staged_path.unlink(missing_ok=True)
+                raise
+        try:
+            staged_path.replace(destination)
+        finally:
+            staged_path.unlink(missing_ok=True)
+        self.sysdiagnose_files.append(normalized_name)
+        self.sysdiagnose_tar_members.append(member)
+
+        if normalized_name.endswith(".ips"):
+            self._add_ips_file(str(destination), destination.read_bytes())
+
     def module_init(self, module) -> None:
         module.ips_files = self.ips_files
         if self.sysdiagnose_format == "tar":
             if self.extracted_sysdiagnose_path is None:
                 raise RuntimeError("Sysdiagnose archive has not been extracted")
+            module.sysdiagnose_tar_members = self.sysdiagnose_tar_members
+            module.sysdiagnose_archive_incomplete = self.archive_incomplete
             module.from_sysdiagnose_folder(
                 self.extracted_sysdiagnose_path, self.sysdiagnose_files
             )

src/mvt/ios/modules/sysdiagnose/base.py

@@ -40,6 +40,10 @@ def __init__(
         self.tar: Optional[tarfile.TarFile] = None
         self.tar_files: list[str] = []
         self.ips_files: list[dict[str, object]] = []
+        # Metadata for safely extracted, complete archive members. None for
+        # directory input or commands that do not supply archive metadata.
+        self.sysdiagnose_tar_members: Optional[list[tarfile.TarInfo]] = None
+        self.sysdiagnose_archive_incomplete = False
 
     def from_sysdiagnose_folder(
         self, target_path: str, sysdiagnose_files: list[str]
@@ -56,27 +60,29 @@ def from_sysdiagnose_tar(
     def _extract_timezone(self):
         """Determine the sysdiagnose timezone from its diagnostic log."""
         file_paths = self._get_files_by_pattern("*/sysdiagnose.log")
-        if not file_paths:
-            self.log.info(
-                "Unable to determine the timezone in which the sysdiagnose was "
-                "generated. Assuming UTC for logs without a timezone."
+        filenames = []
+        if file_paths:
+            content = self._get_file_content(file_paths[0]).decode(
+                "utf-8", errors="replace"
             )
-            return timezone.utc
-
-        content = self._get_file_content(file_paths[0]).decode(
-            "utf-8", errors="replace"
-        )
-        filenames = re.findall(r"sysdiagnose_\S+?\.tar\.gz", content)
+            filenames = re.findall(r"sysdiagnose_\S+?\.tar\.gz", content)
+        if not filenames and self.target_path:
+            # A recovered archive can be missing sysdiagnose.log. Its original
+            # filename still records the collection offset, when not renamed.
+            filename = Path(self.target_path).name
+            if re.match(
+                r"sysdiagnose_\d{4}\.\d{2}\.\d{2}_\d{2}-\d{2}-\d{2}[+-]\d{4}(?:_|\.)",
+                filename,
+            ):
+                filenames = [filename]
         if not filenames:
             self.log.info(
                 "Unable to determine the timezone in which the sysdiagnose was "
                 "generated. Assuming UTC for logs without a timezone."
             )
             return timezone.utc
 
-        timestamp = "_".join(
-            filenames[0].removesuffix(".tar.gz").split("_")[1:3]
-        )
+        timestamp = "_".join(filenames[0].removesuffix(".tar.gz").split("_")[1:3])
         sysdiagnose_timezone = datetime.strptime(
             timestamp, "%Y.%m.%d_%H-%M-%S%z"
         ).tzinfo

tests/test_check_ios_sysdiagnose.py

@@ -1,11 +1,14 @@
 import logging
+import json
 import os
 import tarfile
 
 from click.testing import CliRunner
 
 from mvt.ios.cli import check_sysdiagnose
 
+from .test_cmd_check_sysdiagnose import _recovery_archive
+
 
 CUSTOM_MODULE = """
 from mvt.ios.modules.sysdiagnose import SysdiagnoseExtraction
@@ -90,3 +93,27 @@ def test_check_sysdiagnose_reports_a_truncated_archive(tmp_path, caplog):
     assert "Unable to read the sysdiagnose archive" in caplog.text
     assert "truncated" in caplog.text
     assert "Aborted!" not in result.output
+
+
+def test_check_sysdiagnose_recovers_and_saves_warning(tmp_path, caplog):
+    archive = _recovery_archive(tmp_path, "gzip-body")
+    module_path = tmp_path / "custom_sysdiagnose.py"
+    module_path.write_text(CUSTOM_MODULE, encoding="utf-8")
+    output_path = tmp_path / "output"
+    with caplog.at_level(logging.WARNING, logger="mvt"):
+        result = CliRunner().invoke(
+            check_sysdiagnose,
+            [
+                "--load-module",
+                str(module_path),
+                "--output",
+                str(output_path),
+                str(archive),
+            ],
+        )
+    assert result.exit_code == 0, result.output
+    assert json.loads((output_path / "custom_sysdiagnose_module.json").read_text()) == [
+        {"content": "artifact"}
+    ]
+    assert "truncated or damaged" in caplog.text
+    assert "Recovered 3 complete files" in (output_path / "command.log").read_text()