Merge pull request #570 from onekey-sec/pattern-scan-speedup

Pattern scan speedup

Author: martonilles

flake.lock

@@ -23,7 +23,7 @@ python-lzo = "^1.14"
 plotext = ">=4.2.0,<6.0"
 pluggy = "^1.0.0"
 python-magic = "^0.4.27"
-pyperscan = "^0.2.0"
+pyperscan = "^0.2.2"
 lark = "^1.1.2"
 lz4 = "^4.0.0"
 lief = "^0.12.3"

poetry.lock

@@ -2,7 +2,7 @@
 import pytest
 from pyperscan import Scan
 
-from unblob.file_utils import InvalidInputFormat
+from unblob.file_utils import DEFAULT_BUFSIZE, InvalidInputFormat
 from unblob.finder import build_hyperscan_database, search_chunks
 from unblob.models import File, Handler, HexString, Regex, ValidChunk
 from unblob.parser import InvalidHexString
@@ -63,6 +63,17 @@ def calculate_chunk(self, file, start_offset: int):
         raise ValueError("Error")
 
 
+class TestHandlerL(Handler):
+    NAME = "handlerL"
+    PATTERNS = [Regex("L")]
+
+    def calculate_chunk(self, file, start_offset: int):
+        del file  # unused argument
+        return ValidChunk(
+            start_offset=start_offset, end_offset=start_offset + DEFAULT_BUFSIZE * 2
+        )
+
+
 def test_build_hyperscan_database():
     db = build_hyperscan_database((TestHandlerA, TestHandlerB))
     matches = []
@@ -137,6 +148,31 @@ def calculate_chunk(self, file, start_offset: int):
         pytest.param(
             b"EXCA2345", [ValidChunk(3, 8)], id="exception-ignored-scan-continues"
         ),
+        pytest.param(b"0", [], id="1-byte"),
+        pytest.param(b"1234567890", [], id="no-chunk"),
+        pytest.param(
+            b"A2345L1" + b"1" * DEFAULT_BUFSIZE * 2,
+            [ValidChunk(0, 5), ValidChunk(5, 5 + DEFAULT_BUFSIZE * 2)],
+            id="multi-large-chunk",
+        ),
+        pytest.param(
+            b"L" + b"1" * DEFAULT_BUFSIZE + b"A2345" + b"1" * DEFAULT_BUFSIZE,
+            [ValidChunk(0, DEFAULT_BUFSIZE * 2)],
+            id="large-small-inside-ignored",
+        ),
+        pytest.param(
+            b"0123456789L" + b"1" * DEFAULT_BUFSIZE + b"A2345" + b"1" * DEFAULT_BUFSIZE,
+            [ValidChunk(10, 10 + DEFAULT_BUFSIZE * 2)],
+            id="padding-large-small-inside-ignored",
+        ),
+        pytest.param(
+            b"L" + b"1" * (DEFAULT_BUFSIZE * 2 - 1) + b"A2345" + b"1" * DEFAULT_BUFSIZE,
+            [
+                ValidChunk(0, DEFAULT_BUFSIZE * 2),
+                ValidChunk(DEFAULT_BUFSIZE * 2, DEFAULT_BUFSIZE * 2 + 5),
+            ],
+            id="large-small",
+        ),
     ],
 )
 def test_search_chunks(content, expected_chunks, task_result):
@@ -149,6 +185,7 @@ def test_search_chunks(content, expected_chunks, task_result):
         TestHandlerEof,
         TestHandlerInvalid,
         TestHandlerExc,
+        TestHandlerL,
     )
 
     chunks = search_chunks(file, len(content), handlers, task_result)

pyproject.toml

@@ -9,7 +9,6 @@
 from typing import Iterator, List, Tuple, Union
 
 from dissect.cstruct import cstruct
-from pyperscan import Scan
 
 from .logging import format_hex
 
@@ -255,9 +254,7 @@ def iterate_file(
 
 def stream_scan(scanner, file: File):
     """Scan the whole file by increment of DEFAULT_BUFSIZE using Hyperscan's streaming mode."""
-    for i in range(0, file.size(), DEFAULT_BUFSIZE):
-        if scanner.scan(file[i : i + DEFAULT_BUFSIZE]) == Scan.Terminate:
-            break
+    scanner.scan(file, DEFAULT_BUFSIZE)
 
 
 class StructParser:

tests/test_finder.py

@@ -9,7 +9,7 @@
 from pyperscan import Flag, Pattern, Scan, StreamDatabase
 from structlog import get_logger
 
-from .file_utils import InvalidInputFormat, SeekError, stream_scan
+from .file_utils import DEFAULT_BUFSIZE, InvalidInputFormat, SeekError
 from .handlers import Handlers
 from .models import File, Handler, TaskResult, ValidChunk
 from .parser import InvalidHexString
@@ -24,6 +24,7 @@ class HyperscanMatchContext:
     file_size: int
     all_chunks: List
     task_result: TaskResult
+    start_offset: int
 
 
 def _calculate_chunk(
@@ -69,6 +70,7 @@ def _hyperscan_match(
     context: HyperscanMatchContext, handler: Handler, offset: int, end: int
 ) -> Scan:
     del end  # unused argument
+    offset += context.start_offset
     real_offset = offset + handler.PATTERN_MATCH_OFFSET
 
     if real_offset < 0:
@@ -90,6 +92,7 @@ def _hyperscan_match(
         start_offset=offset,
         real_offset=real_offset,
         _verbosity=2,
+        handler=handler.NAME,
     )
 
     chunk = _calculate_chunk(handler, context.file, real_offset, context.task_result)
@@ -103,15 +106,23 @@ def _hyperscan_match(
         return Scan.Continue
 
     chunk.handler = handler
-    logger.debug("Found valid chunk", chunk=chunk, handler=handler.NAME, _verbosity=2)
+    logger.debug("Found valid chunk", chunk=chunk, handler=handler.NAME, _verbosity=1)
     context.all_chunks.append(chunk)
+    context.start_offset = chunk.end_offset
 
-    # Terminate scan if we match till the end of the file
-    if chunk.end_offset == context.file_size:
-        logger.debug("Chunk covers till end of the file", chunk=chunk)
-        return Scan.Terminate
+    return Scan.Terminate
 
-    return Scan.Continue
+
+def stream_scan_chunks(scanner, file: File, context: HyperscanMatchContext):
+    """Scan the whole file by increment of DEFAULT_BUFSIZE using Hyperscan's streaming mode."""
+    i = context.start_offset
+    with memoryview(file) as data:
+        while i < file.size():
+            if scanner.scan(data[i : i + DEFAULT_BUFSIZE]) == Scan.Terminate:
+                scanner.reset()
+                i = context.start_offset
+            else:
+                i += DEFAULT_BUFSIZE
 
 
 def search_chunks(
@@ -136,12 +147,13 @@ def search_chunks(
         file_size=file_size,
         all_chunks=all_chunks,
         task_result=task_result,
+        start_offset=0,
     )
 
     scanner = hyperscan_db.build(hyperscan_context, _hyperscan_match)
 
     try:
-        stream_scan(scanner, file)
+        stream_scan_chunks(scanner, file, hyperscan_context)
     except Exception as e:
         logger.error(
             "Error scanning for patterns",