Parallelize URL indicator checks (#844)

docs/iocs.md

@@ -37,8 +37,10 @@ export MVT_STIX2="/home/user/IOC1.stix2:/home/user/IOC2.stix2"
 ## Network Access
 
 When checking URL indicators, MVT follows recognized shortened URLs with an
-HTTP `HEAD` request. The following environment variables control these
-requests:
+HTTP `HEAD` request. URL checks are deduplicated and run concurrently, with at
+most 20 requests in progress at a time. Redirects within an individual URL
+chain are still followed sequentially. The following environment variables
+control these requests:
 
 - `MVT_NETWORK_ACCESS_ALLOWED` enables or disables network requests. It defaults
   to `true`. Set it to `false` to prevent MVT from attempting to resolve
@@ -73,4 +75,3 @@ You can automaticallly download the latest public indicator files with the comma
 
 Please [open an issue](https://github.com/mvt-project/mvt/issues/) to suggest new sources of STIX-formatted IOCs.
 
-

src/mvt/android/modules/backup/sms.py

@@ -36,6 +36,8 @@ class SMS(BackupModule):
         if not self.indicators:
             return
 
+        messages = []
+        url_batches = []
         for message in self.results:
             if "body" not in message:
                 continue
@@ -44,12 +46,16 @@ class SMS(BackupModule):
             if message_links == []:
                 message_links = check_for_links(message.get("text", ""))
 
-            ioc_match = self.indicators.check_urls(message_links)
+            messages.append(message)
+            url_batches.append(message_links)
+
+        for message, ioc_match in zip(
+            messages, self.indicators.check_url_batches(url_batches)
+        ):
             if ioc_match:
                 self.alertstore.critical(
                     ioc_match.message, "", message, matched_indicator=ioc_match.ioc
                 )
-                continue
 
     def run(self) -> None:
         sms_path = "apps/com.android.providers.telephony/d_f/*_sms_backup"

src/mvt/common/indicators.py

@@ -7,9 +7,10 @@ import glob
 import json
 import logging
 import os
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Sequence
 
 import ahocorasick
 from appdirs import user_data_dir
@@ -22,6 +23,8 @@ MVT_INDICATORS_FOLDER = os.path.join(MVT_DATA_FOLDER, "indicators")
 
 logger = logging.getLogger(__name__)
 
+URL_CHECK_MAX_WORKERS = 20
+
 
 @dataclass
 class Indicator:
@@ -490,12 +493,41 @@ class Indicators:
         if not urls:
             return None
 
-        for url in urls:
-            check = self.check_url(url)
-            if check:
-                return check
+        return self.check_url_batches([urls])[0]
 
-        return None
+    def check_url_batches(
+        self, url_batches: Sequence[Optional[Sequence[str]]]
+    ) -> List[Optional[IndicatorMatch]]:
+        """Check batches of URLs concurrently while preserving batch order.
+
+        URLs are deduplicated across batches before checking. Each returned item
+        is the first indicator match from the corresponding input batch, using
+        the original URL order.
+        """
+        batches = [list(urls) if urls else [] for urls in url_batches]
+        unique_urls = list(
+            dict.fromkeys(url for urls in batches for url in urls)
+        )
+
+        if not unique_urls:
+            return [None] * len(batches)
+
+        if settings.NETWORK_ACCESS_ALLOWED and len(unique_urls) > 1:
+            worker_count = min(URL_CHECK_MAX_WORKERS, len(unique_urls))
+            with ThreadPoolExecutor(max_workers=worker_count) as executor:
+                url_matches = dict(
+                    zip(unique_urls, executor.map(self.check_url, unique_urls))
+                )
+        else:
+            url_matches = {url: self.check_url(url) for url in unique_urls}
+
+        return [
+            next(
+                (url_matches[url] for url in urls if url_matches[url] is not None),
+                None,
+            )
+            for urls in batches
+        ]
 
     def check_process(self, process: str) -> Optional[IndicatorMatch]:
         """Check the provided process name against the list of process

src/mvt/ios/modules/mixed/shortcuts.py

@@ -76,8 +76,10 @@ class Shortcuts(IOSExtraction):
         if not self.indicators:
             return
 
-        for result in self.results:
-            ioc_match = self.indicators.check_urls(result["action_urls"])
+        url_batches = [result["action_urls"] for result in self.results]
+        for result, ioc_match in zip(
+            self.results, self.indicators.check_url_batches(url_batches)
+        ):
             if ioc_match:
                 self.alertstore.critical(
                     ioc_match.message, "", result, matched_indicator=ioc_match.ioc

src/mvt/ios/modules/mixed/sms.py

@@ -85,12 +85,17 @@ class SMS(IOSExtraction):
         if not self.indicators:
             return
 
+        url_batches = []
         for result in self.results:
             message_links = result.get("links", [])
             # Making sure not link was ignored
             if message_links == []:
                 message_links = check_for_links(result.get("text", ""))
-            ioc_match = self.indicators.check_urls(message_links)
+            url_batches.append(message_links)
+
+        for result, ioc_match in zip(
+            self.results, self.indicators.check_url_batches(url_batches)
+        ):
             if ioc_match:
                 self.alertstore.critical(
                     ioc_match.message, "", result, matched_indicator=ioc_match.ioc

src/mvt/ios/modules/mixed/webkit_session_resource_log.py

@@ -70,6 +70,8 @@ class WebkitSessionResourceLog(IOSExtraction):
         if not self.indicators:
             return
 
+        records = []
+        url_batches = []
         for _, entries in self.results.items():
             for entry in entries:
                 source_domains = self._extract_domains(entry["redirect_source"])
@@ -94,36 +96,42 @@ class WebkitSessionResourceLog(IOSExtraction):
                     )
                 )
 
-                ioc_match = self.indicators.check_urls(all_origins)
-                if ioc_match:
-                    self.alertstore.critical(
-                        ioc_match.message, "", entry, matched_indicator=ioc_match.ioc
-                    )
+                records.append((entry, source_domains, destination_domains))
+                url_batches.append(all_origins)
 
-                    redirect_path = ""
-                    if len(source_domains) > 0:
-                        redirect_path += "SOURCE: "
-                        for idx, item in enumerate(source_domains):
-                            source_domains[idx] = f'"{item}"'
+        for record, ioc_match in zip(
+            records, self.indicators.check_url_batches(url_batches)
+        ):
+            if ioc_match:
+                entry, source_domains, destination_domains = record
+                self.alertstore.critical(
+                    ioc_match.message, "", entry, matched_indicator=ioc_match.ioc
+                )
 
-                        redirect_path += ", ".join(source_domains)
-                        redirect_path += " -> "
+                redirect_path = ""
+                if len(source_domains) > 0:
+                    redirect_path += "SOURCE: "
+                    for idx, item in enumerate(source_domains):
+                        source_domains[idx] = f'"{item}"'
 
-                    redirect_path += f'ORIGIN: "{entry["origin"]}"'
+                    redirect_path += ", ".join(source_domains)
+                    redirect_path += " -> "
 
-                    if len(destination_domains) > 0:
-                        redirect_path += " -> "
-                        redirect_path += "DESTINATION: "
-                        for idx, item in enumerate(destination_domains):
-                            destination_domains[idx] = f'"{item}"'
+                redirect_path += f'ORIGIN: "{entry["origin"]}"'
 
-                        redirect_path += ", ".join(destination_domains)
+                if len(destination_domains) > 0:
+                    redirect_path += " -> "
+                    redirect_path += "DESTINATION: "
+                    for idx, item in enumerate(destination_domains):
+                        destination_domains[idx] = f'"{item}"'
 
-                    self.alertstore.high(
-                        f"Found HTTP redirect between suspicious domains: {redirect_path}",
-                        "",
-                        entry,
-                    )
+                    redirect_path += ", ".join(destination_domains)
+
+                self.alertstore.high(
+                    f"Found HTTP redirect between suspicious domains: {redirect_path}",
+                    "",
+                    entry,
+                )
 
     def _extract_browsing_stats(self, log_path):
         items = []

src/mvt/ios/modules/mixed/whatsapp.py

@@ -61,8 +61,10 @@ class Whatsapp(IOSExtraction):
         if not self.indicators:
             return
 
-        for result in self.results:
-            ioc_match = self.indicators.check_urls(result.get("links", []))
+        url_batches = [result.get("links", []) for result in self.results]
+        for result, ioc_match in zip(
+            self.results, self.indicators.check_url_batches(url_batches)
+        ):
             if ioc_match:
                 self.alertstore.critical(
                     ioc_match.message, "", result, matched_indicator=ioc_match.ioc

tests/common/test_indicators.py

@@ -5,7 +5,9 @@
 
 import logging
 import os
+import threading
 
+import requests
 
 from mvt.common.config import settings
 from mvt.common.indicators import Indicators
@@ -88,6 +90,119 @@ class TestIndicators:
         assert ind.check_url("https://goo.gl/maps/example") is None
         head_request.assert_not_called()
 
+    def test_check_url_batches_preserves_order(self, indicator_file):
+        ind = Indicators(log=logging)
+        ind.load_indicators_files([indicator_file], load_default=False)
+
+        matches = ind.check_url_batches(
+            [
+                [
+                    "https://github.com",
+                    "http://example.com/thisisbad",
+                    "https://www.example.org/foobar",
+                ],
+                ["https://github.com", "https://www.example.org/foobar"],
+                [],
+                None,
+            ]
+        )
+
+        assert matches[0]
+        assert matches[0].ioc.value == "http://example.com/thisisbad"
+        assert matches[1]
+        assert matches[1].ioc.value == "example.org"
+        assert matches[2] is None
+        assert matches[3] is None
+
+    def test_check_url_batches_deduplicates_and_limits_workers(
+        self, indicator_file, mocker
+    ):
+        ind = Indicators(log=logging)
+        ind.load_indicators_files([indicator_file], load_default=False)
+        mocker.patch("mvt.common.indicators.URL_CHECK_MAX_WORKERS", 2)
+
+        barrier = threading.Barrier(2)
+        lock = threading.Lock()
+        calls = []
+        active = 0
+        max_active = 0
+
+        def head_request(url, timeout):
+            nonlocal active, max_active
+            with lock:
+                calls.append(url)
+                active += 1
+                max_active = max(max_active, active)
+                call_number = len(calls)
+
+            try:
+                if call_number <= 2:
+                    barrier.wait(timeout=5)
+                return mocker.Mock(status_code=200, headers={})
+            finally:
+                with lock:
+                    active -= 1
+
+        mocker.patch("mvt.common.url.requests.head", side_effect=head_request)
+        urls = [
+            "https://bit.ly/one",
+            "https://tinyurl.com/two",
+            "https://t.co/three",
+        ]
+
+        assert ind.check_url_batches([urls, [urls[0]]]) == [None, None]
+        assert sorted(calls) == sorted(urls)
+        assert max_active == 2
+
+    def test_check_url_batches_respects_disabled_network(
+        self, indicator_file, mocker
+    ):
+        ind = Indicators(log=logging)
+        ind.load_indicators_files([indicator_file], load_default=False)
+        mocker.patch("mvt.common.indicators.settings.NETWORK_ACCESS_ALLOWED", False)
+        head_request = mocker.patch("mvt.common.url.requests.head")
+
+        assert ind.check_url_batches([["https://bit.ly/example"]]) == [None]
+        head_request.assert_not_called()
+
+    def test_check_url_batches_handles_nested_redirects_and_request_failures(
+        self, indicator_file, mocker
+    ):
+        ind = Indicators(log=logging)
+        ind.load_indicators_files([indicator_file], load_default=False)
+
+        def head_request(url, timeout):
+            if url == "https://bit.ly/failure":
+                raise requests.Timeout()
+            if url == "https://tinyurl.com/nested":
+                return mocker.Mock(
+                    status_code=301,
+                    headers={"Location": "https://t.co/nested"},
+                )
+            if url == "https://t.co/nested":
+                return mocker.Mock(
+                    status_code=302,
+                    headers={"Location": "https://www.example.org/landing"},
+                )
+            raise AssertionError(f"Unexpected URL: {url}")
+
+        head = mocker.patch(
+            "mvt.common.url.requests.head", side_effect=head_request
+        )
+
+        matches = ind.check_url_batches(
+            [["https://bit.ly/failure"], ["https://tinyurl.com/nested"]]
+        )
+
+        assert matches[0] is None
+        assert matches[1]
+        assert matches[1].ioc.value == "example.org"
+        assert {call.args[0] for call in head.call_args_list} == {
+            "https://bit.ly/failure",
+            "https://tinyurl.com/nested",
+            "https://t.co/nested",
+        }
+
     def test_check_file_hash(self, indicator_file):
         ind = Indicators(log=logging)
         ind.load_indicators_files([indicator_file], load_default=False)

tests/ios_backup/test_sms.py

@@ -29,3 +29,31 @@ class TestSMSModule:
         m.indicators = ind
         run_module(m)
         assert len(m.alertstore.alerts) == 1
+
+    def test_detection_batches_urls_and_preserves_event(self, indicator_file, mocker):
+        results = [
+            {
+                "text": "first",
+                "links": ["http://example.com/thisisbad"],
+            },
+            {
+                "text": "second",
+                "links": ["https://github.com"],
+            },
+        ]
+        m = SMS(results=results)
+        ind = Indicators(log=logging.getLogger())
+        ind.parse_stix2(indicator_file)
+        batch_check = mocker.spy(ind, "check_url_batches")
+        m.indicators = ind
+
+        m.check_indicators()
+
+        batch_check.assert_called_once_with(
+            [
+                ["http://example.com/thisisbad"],
+                ["https://github.com"],
+            ]
+        )
+        assert len(m.alertstore.alerts) == 1
+        assert m.alertstore.alerts[0].event is results[0]