Improve DHCP IP address parsing speed for journal

dissect/target/helpers/network_managers.py

@@ -7,12 +7,14 @@
 from io import StringIO
 from itertools import chain
 from re import compile, sub
-from typing import Any, Callable, Iterable, Match, Optional
+from typing import Any, Callable, Iterable, Iterator, Match, Optional
 
 from defusedxml import ElementTree
 
 from dissect.target.exceptions import PluginError
 from dissect.target.helpers.fsutil import TargetPath
+from dissect.target.plugins.os.unix.log.journal import JournalRecord
+from dissect.target.plugins.os.unix.log.messages import MessagesRecord
 from dissect.target.target import Target
 
 log = logging.getLogger(__name__)
@@ -509,14 +511,15 @@
         return values
 
 
-def parse_unix_dhcp_log_messages(target) -> list[str]:
+def parse_unix_dhcp_log_messages(target: Target, iter_all: bool = False) -> set[str]:
     """Parse local syslog, journal and cloud init-log files for DHCP lease IPs.
 
     Args:
         target: Target to discover and obtain network information from.
+        iter_all: Parse limited amount of journal messages (first 10000) or all of them.
 
     Returns:
-        List of DHCP ip addresses.
+        A set of found DHCP IP addresses.
     """
     ips = set()
     messages = set()
@@ -530,9 +533,19 @@
     if not messages:
         target.log.warning(f"Could not search for DHCP leases using {log_func}: No log entries found.")
 
-    for record in messages:
+    def records_enumerate(iterable: Iterable) -> Iterator[tuple[int, JournalRecord | MessagesRecord]]:
+        count = 0
+        for rec in iterable:
+            if rec._desc.name == "linux/log/journal":
+                count += 1
+            yield count, rec
+
+    for count, record in records_enumerate(messages):
         line = record.message
 
+        if not line:
+            continue
+
         # Ubuntu cloud-init
         if "Received dhcp lease on" in line:
             interface, ip, netmask = re.search(r"Received dhcp lease on (\w{0,}) for (\S+)\/(\S+)", line).groups()
@@ -576,9 +589,11 @@
             ips.add(ip)
             continue
 
-        # Journals and syslogs can be large and slow to iterate,
-        # so we stop if we have some results and have reached the journal plugin.
-        if len(ips) >= 2 and record._desc.name == "linux/log/journal":
+        # The journal parser is relatively slow, so we stop when we have read 10000 journal entries,
+        # or if we have found at least one ip address. When `iter_all` is `True` we continue searching.
+        if not iter_all and (ips or count > 10_000):
+            if not ips:
+                target.log.warning("No DHCP IP addresses found in first 10000 journal entries.")
             break
 
     return ips

dissect/target/plugins/os/unix/linux/_os.py

@@ -41,7 +41,7 @@ def ips(self) -> list[str]:
             for ip in ip_set:
                 ips.append(ip)
 
-        for ip in parse_unix_dhcp_log_messages(self.target):
+        for ip in parse_unix_dhcp_log_messages(self.target, iter_all=False):
             if ip not in ips:
                 ips.append(ip)
 

tests/plugins/os/unix/test_ips.py

@@ -1,27 +1,47 @@
 import textwrap
 from io import BytesIO
+from unittest.mock import patch
 
 import pytest
 
 from dissect.target import Target
 from dissect.target.filesystem import VirtualFilesystem
 from dissect.target.helpers.network_managers import NetworkManager
 from dissect.target.plugins.os.unix.linux._os import LinuxPlugin
+from dissect.target.tools.query import main as target_query
 from tests._utils import absolute_path
 
 
-def test_ips_dhcp(target_unix_users: Target, fs_unix: VirtualFilesystem) -> None:
+@pytest.mark.parametrize(
+    "expected_ips, messages",
+    [
+        (
+            ["10.13.37.1"],
+            "Jan  1 13:37:01 hostname NetworkManager[1]: <info>  [1600000000.0000] dhcp4 (eth0): option ip_address           => '10.13.37.1'",  # noqa: E501
+        ),
+        (["10.13.37.2"], "Feb  2 13:37:02 test systemd-networkd[2]: eth0: DHCPv4 address 10.13.37.2/24 via 10.13.37.0"),
+        (
+            ["10.13.37.3"],
+            "Mar  3 13:37:03 localhost NetworkManager[3]: <info>  [1600000000.0003] dhcp4 (eth0):   address 10.13.37.3",
+        ),
+        (
+            ["10.13.37.4"],
+            "Apr  4 13:37:04 localhost dhclient[4]: bound to 10.13.37.4 -- renewal in 1337 seconds.",
+        ),
+        (
+            ["2001:db8::"],
+            (
+                "Jun  6 13:37:06 test systemd-networkd[5]: eth0: DHCPv6 address 2001:db8::/64 via 2001:db8:ffff:ffff:ffff:ffff:ffff:ffff\n"  # noqa: E501
+                "May  5 13:37:05 test systemd-networkd[5]: eth0: DHCPv6 lease lost\n"
+            ),
+        ),
+    ],
+)
+def test_ips_dhcp(
+    target_unix_users: Target, fs_unix: VirtualFilesystem, expected_ips: list[str], messages: str
+) -> None:
     """Test DHCP lease messages from /var/log/syslog."""
 
-    messages = """
-    Jan  1 13:37:01 hostname NetworkManager[1]: <info>  [1600000000.0000] dhcp4 (eth0): option ip_address           => '10.13.37.1'
-    Feb  2 13:37:02 test systemd-networkd[2]: eth0: DHCPv4 address 10.13.37.2/24 via 10.13.37.0
-    Mar  3 13:37:03 localhost NetworkManager[3]: <info>  [1600000000.0003] dhcp4 (eth0):   address 10.13.37.3
-    Apr  4 13:37:04 localhost dhclient[4]: bound to 10.13.37.4 -- renewal in 1337 seconds.
-    May  5 13:37:05 test systemd-networkd[5]: eth0: DHCPv6 lease lost
-    Jun  6 13:37:06 test systemd-networkd[5]: eth0: DHCPv6 address 2001:db8::/64 via 2001:db8:ffff:ffff:ffff:ffff:ffff:ffff
-    """  # noqa E501
-
     fs_unix.map_file_fh(
         "/var/log/syslog",
         BytesIO(textwrap.dedent(messages).encode()),
@@ -30,8 +50,53 @@ def test_ips_dhcp(target_unix_users: Target, fs_unix: VirtualFilesystem) -> None
     target_unix_users.add_plugin(LinuxPlugin)
     results = target_unix_users.ips
     results.reverse()
-    assert len(results) == 5
-    assert sorted(results) == ["10.13.37.1", "10.13.37.2", "10.13.37.3", "10.13.37.4", "2001:db8::"]
+    assert len(results) == len(expected_ips)
+    assert sorted(results) == expected_ips
+
+
+@pytest.mark.parametrize(
+    "flag, expected_out",
+    [
+        (None, "['10.13.37.2']"),
+        # ("--dhcp-all", "['10.13.37.2', '10.13.37.1']"),
+        # Temporarily disabled behaviour, for discussion see:
+        # https://github.com/fox-it/dissect.target/pull/687#discussion_r1698515269
+    ],
+)
+def test_ips_dhcp_arg(
+    target_unix: Target,
+    fs_unix: VirtualFilesystem,
+    flag: str,
+    expected_out: str,
+    capsys: pytest.CaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test --dhcp-all flag behaviour"""
+
+    fs_unix.map_file_fh("/etc/timezone", BytesIO(b"Europe/Amsterdam"))
+
+    messages = """
+    Apr  1 13:37:01 localhost dhclient[4]: bound to 10.13.37.1 -- renewal in 1337 seconds.
+    Apr  2 13:37:02 localhost foo[1]: some other message.
+    Apr  3 13:37:03 localhost dhclient[4]: bound to 10.13.37.2 -- renewal in 1337 seconds.
+    """
+
+    fs_unix.map_file_fh(
+        "/var/log/syslog",
+        BytesIO(textwrap.dedent(messages).encode()),
+    )
+    target_unix.add_plugin(LinuxPlugin)
+
+    argv = ["target-query", "foo", "-f", "ips"]
+    if flag:
+        argv.append(flag)
+
+    with patch("dissect.target.Target.open_all", return_value=[target_unix]):
+        with monkeypatch.context() as m:
+            m.setattr("sys.argv", argv)
+            target_query()
+            out, _ = capsys.readouterr()
+            assert expected_out in out
 
 
 def test_ips_cloud_init(target_unix_users: Target, fs_unix: VirtualFilesystem) -> None: