feat(handler): improve PE handler

Create a PEExtractor to handle the logic for extracting from PEs.

Author: jcrussell

python/unblob/handlers/executable/pe.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import Optional
 
 import io
@@ -8,6 +9,8 @@
 from unblob.extractors.command import Command
 
 from ...models import (
+    Extractor,
+    ExtractResult,
     File,
     Handler,
     HandlerDoc,
@@ -21,6 +24,36 @@
 
 logger = get_logger()
 
+
+class PEExtractor(Extractor):
+    """
+    PEExtractor extracts files embedded within PEs such as installers.
+    """
+
+    def extract(self, inpath: Path, outdir: Path) -> Optional[ExtractResult]:
+        binary = lief.PE.parse(inpath)
+        if not binary:
+            return None
+
+        if not self.is_nsis(binary):
+            return None
+
+        return Command("7z", "x", "-y", "{inpath}", "-o{outdir}").extract(inpath, outdir)
+
+
+    def is_nsis(self, binary: lief.PE.Binary) -> bool:
+        """
+        Test if binary appears to be a Nullsoft Installer self-extracting archive
+
+        TODO: this series of tests is possibly too strict
+        """
+
+        return binary.has_resources and \
+            binary.resources_manager.has_manifest and \
+            "Nullsoft" in binary.resources_manager.manifest
+
+
+
 class PEHandler(Handler):
     NAME = "pe"
 
@@ -30,12 +63,18 @@ class PEHandler(Handler):
             // MZ header
             4d 5a
             """
+        ),
+        HexString(
+            """
+            // PE header
+            50 45 00 00
+            """
         )
     ]
 
-    EXTRACTOR = Command(
-        "7z", "x", "-y", "{inpath}", "-o{outdir}"
-    )
+
+    EXTRACTOR = PEExtractor()
+
 
     DOC = HandlerDoc(
         name="pe",
@@ -55,21 +94,6 @@ class PEHandler(Handler):
         limitations=[],
     )
 
-    def is_nsis(self, binary: lief.PE.Binary) -> bool:
-        """
-        Test if binary appears to be a Nullsoft Installer self-extracting archive
-
-        TODO: this series of tests is possibly too strict
-        """
-
-        if not binary.has_resources:
-            return False
-
-        if not binary.resources_manager.has_manifest:
-            return False
-
-        return "Nullsoft" in binary.resources_manager.manifest
-
 
     def calculate_chunk(self, file: File, start_offset: int) -> Optional[ValidChunk]:
         file.seek(start_offset, io.SEEK_SET)
@@ -78,9 +102,6 @@ def calculate_chunk(self, file: File, start_offset: int) -> Optional[ValidChunk]
         if not binary:
             return None
 
-        if not self.is_nsis(binary):
-            return None
-
         return ValidChunk(
             start_offset = start_offset,
             end_offset = start_offset + binary.original_size,