Merge pull request #622 from onekey-sec/621-add-timeout

Add execution timeout to Command Extractor

Author: qkaiser

unblob/extractors/command.py

@@ -6,11 +6,19 @@
 from structlog import get_logger
 
 from unblob.models import DirectoryExtractor, ExtractError, Extractor
-from unblob.report import ExtractCommandFailedReport, ExtractorDependencyNotFoundReport
+from unblob.report import (
+    ExtractCommandFailedReport,
+    ExtractorDependencyNotFoundReport,
+    ExtractorTimedOut,
+)
 
 if TYPE_CHECKING:
     import io
 
+# value that is high enough not to block long running execution such as extraction of large
+# disk images, but small enough to make sure unblob finish its execution at some point.
+COMMAND_TIMEOUT = 12 * 60 * 60
+
 logger = get_logger()
 
 
@@ -45,6 +53,7 @@ def no_op():
                 cmd,
                 stdout=stdout_file,
                 stderr=subprocess.PIPE,
+                timeout=COMMAND_TIMEOUT,
             )
             if res.returncode != 0:
                 error_report = ExtractCommandFailedReport(
@@ -65,6 +74,13 @@ def no_op():
                 **error_report.asdict(),
             )
             raise ExtractError(error_report) from None
+        except subprocess.TimeoutExpired as e:
+            error_report = ExtractorTimedOut(cmd=e.cmd, timeout=e.timeout)
+            logger.error(
+                "Extract command timed out.",
+                **error_report.asdict(),
+            )
+            raise ExtractError(error_report) from None
         finally:
             cleanup()
 

unblob/report.py

@@ -103,6 +103,15 @@ class ExtractorDependencyNotFoundReport(ErrorReport):
     dependencies: List[str]
 
 
+@attr.define(kw_only=True, frozen=True)
+class ExtractorTimedOut(ErrorReport):
+    """Describes an error when the extractor execution timed out."""
+
+    severity: Severity = Severity.ERROR
+    cmd: str
+    timeout: float
+
+
 @attr.define(kw_only=True, frozen=True)
 class MaliciousSymlinkRemoved(ErrorReport):
     """Describes an error when malicious symlinks have been removed from disk."""