Skip to content

Commit 2beeadb

Browse files
committed
feat: read DTL JSON slug indexes, keep pickle fallback
Upstream devicetype-library dropped tests/known-*.pickle in favor of tests/known-*.json (GHSA-492p-5wp7-2w7c, PR #4238). The slug fast path looked only for the pickle files, so against current DTL master it always missed and fell back to a full file scan. Prefer known-*.json (same [slug, relpath] / [model, vendor_dir] shape), falling back to the legacy pickle for older pinned checkouts. Loaders are dispatched by extension; JSON carries no code-execution risk so it skips the restricted unpickler. Shape/size validation is shared across both. PRs #4239 and #4240 only harden the DTL test harness (package layout, URL allowlisting) and do not affect this consumer.
1 parent a266408 commit 2beeadb

3 files changed

Lines changed: 226 additions & 47 deletions

File tree

core/repo.py

Lines changed: 93 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Git repository helpers for cloning, updating, and parsing the device-type library."""
22

3+
import json
34
import os
45
import pickle
56
from glob import glob
@@ -28,22 +29,39 @@ def find_class(self, module, name):
2829
)
2930

3031

31-
_PICKLE_MAX_BYTES = 10 * 1024 * 1024 # 10 MiB — DTL pickles are typically <500 KiB
32+
_INDEX_MAX_BYTES = 10 * 1024 * 1024 # 10 MiB — DTL index files are typically <1 MiB
3233

3334

34-
def _vendor_slugs_from_pickle(
35-
pickle_path: str, slugs_lower: list, slug_format, subdir_filter: "Optional[str]" = None
35+
def _resolve_index_path(base_dir: str, stem: str) -> "Optional[str]":
36+
"""Return the path to a DTL index file, preferring the JSON form over the legacy pickle.
37+
38+
Upstream DTL replaced ``tests/known-*.pickle`` with ``tests/known-*.json`` (GHSA-492p-5wp7-2w7c).
39+
We prefer ``<stem>.json`` when present and fall back to ``<stem>.pickle`` so users pinned to an
40+
older DTL checkout still get the slug fast path. Returns ``None`` when neither exists.
41+
"""
42+
json_path = os.path.join(base_dir, stem + ".json")
43+
if os.path.exists(json_path):
44+
return json_path
45+
pickle_path = os.path.join(base_dir, stem + ".pickle")
46+
if os.path.exists(pickle_path):
47+
return pickle_path
48+
return None
49+
50+
51+
def _vendor_slugs_from_index(
52+
index_path: "Optional[str]", slugs_lower: list, slug_format, subdir_filter: "Optional[str]" = None
3653
) -> "Optional[set]":
37-
"""Load a (model_name, vendor_dir) pickle and return the set of vendor slugs matching *slugs_lower*.
54+
"""Load a (model_name, vendor_dir) index and return the set of vendor slugs matching *slugs_lower*.
3855
39-
*subdir_filter*, if given, requires the vendor_dir to contain that substring.
40-
Returns ``None`` when the pickle is unavailable (missing or unreadable), so callers
56+
*index_path* may point at a ``.json`` (current) or ``.pickle`` (legacy) file; the loader is
57+
chosen by extension. *subdir_filter*, if given, requires the vendor_dir to contain that
58+
substring. Returns ``None`` when the index is unavailable (missing or unreadable), so callers
4159
can distinguish "no matches" (empty set) from "hint unavailable" (None).
4260
"""
43-
if not os.path.exists(pickle_path):
61+
if index_path is None:
4462
return None
4563
try:
46-
entries = _safe_pickle_load(pickle_path)
64+
entries = _safe_index_load(index_path)
4765
except Exception:
4866
return None
4967
result = set()
@@ -63,32 +81,67 @@ def _safe_abs_path(repo_root: str, relpath: str) -> "Optional[str]":
6381
return abs_path if abs_path.startswith(os.path.normpath(repo_root) + os.sep) else None
6482

6583

66-
def _safe_pickle_load(path: str):
67-
"""Load a DTL upstream pickle using the restricted unpickler.
84+
def _validate_index_data(data, source: str):
85+
"""Validate that *data* is a set/list of (str, str) 2-element pairs.
6886
69-
Enforces a hard size cap before unpickling and validates the loaded object
70-
is a set/list of (str, str) 2-tuples so malformed/oversized pickles cannot
71-
cause resource exhaustion. Returns the loaded set on success or raises
72-
``ValueError`` on shape violations (callers should catch and fall back).
87+
Shared by the JSON and pickle loaders. JSON arrays decode to lists rather than tuples, so
88+
both container shapes are accepted. Raises ``ValueError`` on any shape violation (callers
89+
should catch and fall back to the normal glob path).
7390
"""
74-
size = os.path.getsize(path)
75-
if size > _PICKLE_MAX_BYTES:
76-
raise ValueError(f"Pickle file {path!r} is {size} bytes (limit {_PICKLE_MAX_BYTES}); refusing to load.")
77-
with open(path, "rb") as fh:
78-
data = _RestrictedUnpickler(fh).load()
7991
if not isinstance(data, (set, list, frozenset)):
80-
raise ValueError(f"Unexpected pickle root type {type(data).__name__!r}; expected set/list.")
92+
raise ValueError(f"Unexpected {source} root type {type(data).__name__!r}; expected set/list.")
8193
for item in data:
8294
if (
83-
not isinstance(item, tuple)
95+
not isinstance(item, (tuple, list))
8496
or len(item) != 2
8597
or not isinstance(item[0], str)
8698
or not isinstance(item[1], str)
8799
):
88-
raise ValueError(f"Unexpected item shape in pickle: {item!r}")
100+
raise ValueError(f"Unexpected item shape in {source}: {item!r}")
89101
return data
90102

91103

104+
def _check_index_size(path: str):
105+
"""Enforce the hard size cap before reading an index file, guarding against resource exhaustion."""
106+
size = os.path.getsize(path)
107+
if size > _INDEX_MAX_BYTES:
108+
raise ValueError(f"Index file {path!r} is {size} bytes (limit {_INDEX_MAX_BYTES}); refusing to load.")
109+
110+
111+
def _safe_json_load(path: str):
112+
"""Load a DTL upstream ``known-*.json`` index file.
113+
114+
The current DTL format is a JSON array of ``[str, str]`` pairs. Enforces a hard size cap and
115+
validates the decoded shape. JSON carries no code-execution risk, so no restricted loader is
116+
needed. Returns the loaded list on success or raises ``ValueError`` on shape violations.
117+
"""
118+
_check_index_size(path)
119+
with open(path, encoding="utf-8") as fh:
120+
data = json.load(fh)
121+
return _validate_index_data(data, "JSON index")
122+
123+
124+
def _safe_pickle_load(path: str):
125+
"""Load a legacy DTL ``known-*.pickle`` index using the restricted unpickler.
126+
127+
Enforces a hard size cap before unpickling and validates the loaded object is a set/list of
128+
(str, str) 2-tuples so malformed/oversized pickles cannot cause resource exhaustion. Returns
129+
the loaded set on success or raises ``ValueError`` on shape violations (callers should catch
130+
and fall back).
131+
"""
132+
_check_index_size(path)
133+
with open(path, "rb") as fh:
134+
data = _RestrictedUnpickler(fh).load()
135+
return _validate_index_data(data, "pickle")
136+
137+
138+
def _safe_index_load(path: str):
139+
"""Load a DTL slug index, dispatching on file extension (``.json`` current, ``.pickle`` legacy)."""
140+
if path.endswith(".json"):
141+
return _safe_json_load(path)
142+
return _safe_pickle_load(path)
143+
144+
92145
def validate_git_url(url):
93146
"""Determine whether a Git remote URL is allowed (HTTPS, SSH, or file://).
94147
@@ -484,16 +537,18 @@ def get_devices(self, base_path, vendors: list = None):
484537
return files, discovered_vendors
485538

486539
def resolve_slug_files(self, slugs):
487-
"""Use the upstream pickle indexes to resolve YAML file paths for slug/model matches.
540+
"""Use the upstream slug indexes to resolve YAML file paths for slug/model matches.
488541
489-
The DTL repo ships three pickle files under ``tests/``:
542+
The DTL repo ships three index files under ``tests/``. Upstream replaced the original
543+
``known-*.pickle`` files with ``known-*.json`` (GHSA-492p-5wp7-2w7c); both carry the same
544+
data shape, and we prefer JSON, falling back to pickle for older pinned checkouts:
490545
491-
* ``known-slugs.pickle`` — set of ``(manufacturer_prefixed_slug, relpath)`` for
546+
* ``known-slugs.json`` — list of ``(manufacturer_prefixed_slug, relpath)`` for
492547
device types. ``relpath`` is relative to the repo root, e.g.
493548
``device-types/Nokia/7750-SR-7s.yaml``.
494-
* ``known-modules.pickle`` — set of ``(model_name, vendor_dir)`` for module
549+
* ``known-modules.json`` — list of ``(model_name, vendor_dir)`` for module
495550
types. Only the vendor directory is stored, not the file name.
496-
* ``known-racks.pickle`` — same format as known-modules.
551+
* ``known-racks.json`` — same format as known-modules.
497552
498553
Matching uses a **case-insensitive substring** check identical to the runtime
499554
:meth:`parse_files` filter so that partial slug/model searches work the same way.
@@ -502,32 +557,33 @@ def resolve_slug_files(self, slugs):
502557
slugs (list[str]): User-supplied slug/model substrings (``--slugs``).
503558
504559
Returns:
505-
dict or None: ``None`` when the device pickle is unavailable (caller falls back
560+
dict or None: ``None`` when the device index is unavailable (caller falls back
506561
to the normal glob path). Otherwise a dict with the keys:
507562
508563
``"device_files"``
509-
``{vendor_slug: [abs_path, ...]}`` for devices resolved via pickle.
564+
``{vendor_slug: [abs_path, ...]}`` for devices resolved via the index.
510565
``"module_vendors"``
511566
``{vendor_slug}`` — set of vendor slugs that may contain matching module
512-
types, or ``None`` when the module pickle was unavailable (caller should
567+
types, or ``None`` when the module index was unavailable (caller should
513568
fall back to full glob+parse instead of skipping).
514569
``"rack_vendors"``
515570
``{vendor_slug}`` — same for rack types; ``None`` means unavailable.
516571
"""
517572
repo_root = self.get_absolute_path()
518-
device_pickle = os.path.join(repo_root, "tests", "known-slugs.pickle")
519-
module_pickle = os.path.join(repo_root, "tests", "known-modules.pickle")
520-
rack_pickle = os.path.join(repo_root, "tests", "known-racks.pickle")
573+
tests_dir = os.path.join(repo_root, "tests")
574+
device_index = _resolve_index_path(tests_dir, "known-slugs")
575+
module_index = _resolve_index_path(tests_dir, "known-modules")
576+
rack_index = _resolve_index_path(tests_dir, "known-racks")
521577

522-
if not os.path.exists(device_pickle):
578+
if device_index is None:
523579
return None
524580

525581
slugs_lower = [s.casefold() for s in slugs]
526582

527583
# --- device types --------------------------------------------------
528584
device_files = {} # vendor_slug -> [abs_path]
529585
try:
530-
known_slugs = _safe_pickle_load(device_pickle)
586+
known_slugs = _safe_index_load(device_index)
531587
except Exception:
532588
return None
533589

@@ -545,10 +601,10 @@ def resolve_slug_files(self, slugs):
545601
device_files.setdefault(vendor_slug, []).append(abs_path)
546602

547603
# --- module types --------------------------------------------------
548-
module_vendors = _vendor_slugs_from_pickle(module_pickle, slugs_lower, self.slug_format)
604+
module_vendors = _vendor_slugs_from_index(module_index, slugs_lower, self.slug_format)
549605

550606
# --- rack types ----------------------------------------------------
551-
rack_vendors = _vendor_slugs_from_pickle(rack_pickle, slugs_lower, self.slug_format, subdir_filter="rack-types")
607+
rack_vendors = _vendor_slugs_from_index(rack_index, slugs_lower, self.slug_format, subdir_filter="rack-types")
552608

553609
return {
554610
"device_files": device_files,

nb-dt-import.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -854,20 +854,20 @@ def _validate_argument_combinations(parser, args):
854854

855855

856856
def _apply_slug_fast_path(dtl_repo, args, vendors_to_process, handle):
857-
"""Use upstream pickle indexes to pre-resolve files and narrow the vendor list.
857+
"""Use upstream slug indexes to pre-resolve files and narrow the vendor list.
858858
859-
When ``args.slugs`` is set and the DTL pickle files are present, resolves
860-
exactly which device-type files match the requested slugs and restricts
861-
``vendors_to_process`` to only those vendors. Returns a ``(vendors, resolved)``
862-
pair where *resolved* is the dict returned by :meth:`DTLRepo.resolve_slug_files`
863-
(or ``None`` when the pickle is absent/unavailable).
859+
When ``args.slugs`` is set and the DTL index files (``known-*.json``, or legacy
860+
``known-*.pickle``) are present, resolves exactly which device-type files match the
861+
requested slugs and restricts ``vendors_to_process`` to only those vendors. Returns a
862+
``(vendors, resolved)`` pair where *resolved* is the dict returned by
863+
:meth:`DTLRepo.resolve_slug_files` (or ``None`` when the index is absent/unavailable).
864864
"""
865865
if not args.slugs:
866866
return vendors_to_process, None
867867

868868
slug_resolved = dtl_repo.resolve_slug_files(args.slugs)
869869
if slug_resolved is None:
870-
handle.verbose_log("Slug pickle unavailable; falling back to full file scan.")
870+
handle.verbose_log("Slug index unavailable; falling back to full file scan.")
871871
return vendors_to_process, None
872872

873873
matched_vendor_slugs = (
@@ -876,11 +876,11 @@ def _apply_slug_fast_path(dtl_repo, args, vendors_to_process, handle):
876876
| (slug_resolved["rack_vendors"] or set())
877877
)
878878
if not matched_vendor_slugs:
879-
handle.verbose_log("Slug pickle returned no matches; falling back to full file scan.")
879+
handle.verbose_log("Slug index returned no matches; falling back to full file scan.")
880880
return vendors_to_process, None
881881
narrowed = [v for v in vendors_to_process if v["slug"] in matched_vendor_slugs]
882882
handle.verbose_log(
883-
f"Slug pickle resolved {sum(len(f) for f in slug_resolved['device_files'].values())} "
883+
f"Slug index resolved {sum(len(f) for f in slug_resolved['device_files'].values())} "
884884
f"device file(s) across {len(matched_vendor_slugs)} vendor(s)."
885885
)
886886
return narrowed, slug_resolved

tests/test_repo.py

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11
import pytest
22
from unittest.mock import MagicMock, call, mock_open, patch
33
from git import exc as git_exc
4-
from core.repo import DTLRepo, _safe_pickle_load, validate_git_url, normalize_port_mappings
4+
from core.repo import (
5+
DTLRepo,
6+
_resolve_index_path,
7+
_safe_index_load,
8+
_safe_json_load,
9+
_safe_pickle_load,
10+
validate_git_url,
11+
normalize_port_mappings,
12+
)
513

614

715
class TestValidateGitUrl:
@@ -1512,3 +1520,118 @@ def test_rack_pickle_only_uses_rack_type_entries(self, tmp_path):
15121520
result = repo.resolve_slug_files(["rack"])
15131521

15141522
assert result["rack_vendors"] == {"apc"}
1523+
1524+
1525+
class TestIndexLoaders:
1526+
"""Tests for the format-dispatching index loaders and path resolver."""
1527+
1528+
def test_resolve_prefers_json_over_pickle(self, tmp_path):
1529+
(tmp_path / "known-slugs.json").write_text("[]")
1530+
(tmp_path / "known-slugs.pickle").write_bytes(b"x")
1531+
assert _resolve_index_path(str(tmp_path), "known-slugs").endswith("known-slugs.json")
1532+
1533+
def test_resolve_falls_back_to_pickle(self, tmp_path):
1534+
(tmp_path / "known-slugs.pickle").write_bytes(b"x")
1535+
assert _resolve_index_path(str(tmp_path), "known-slugs").endswith("known-slugs.pickle")
1536+
1537+
def test_resolve_returns_none_when_absent(self, tmp_path):
1538+
assert _resolve_index_path(str(tmp_path), "known-slugs") is None
1539+
1540+
def test_safe_json_load_reads_pairs(self, tmp_path):
1541+
path = tmp_path / "known-slugs.json"
1542+
path.write_text('[["nokia-7750", "device-types/Nokia/7750.yaml"]]')
1543+
data = _safe_json_load(str(path))
1544+
assert data == [["nokia-7750", "device-types/Nokia/7750.yaml"]]
1545+
1546+
def test_safe_json_load_rejects_bad_shape(self, tmp_path):
1547+
path = tmp_path / "known-slugs.json"
1548+
path.write_text('[["only-one-element"]]')
1549+
with pytest.raises(ValueError, match="Unexpected item shape"):
1550+
_safe_json_load(str(path))
1551+
1552+
def test_safe_index_load_dispatches_on_extension(self, tmp_path):
1553+
import pickle
1554+
1555+
json_path = tmp_path / "a.json"
1556+
json_path.write_text('[["x", "y"]]')
1557+
pickle_path = tmp_path / "a.pickle"
1558+
pickle_path.write_bytes(pickle.dumps({("p", "q")}))
1559+
assert _safe_index_load(str(json_path)) == [["x", "y"]]
1560+
assert _safe_index_load(str(pickle_path)) == {("p", "q")}
1561+
1562+
1563+
class TestResolveSlugFilesJson:
1564+
"""Tests for the JSON slug fast path (current upstream format)."""
1565+
1566+
def _make_repo(self):
1567+
return TestResolveSlugFiles._make_repo(self)
1568+
1569+
def _write_json(self, tests_dir, stem, entries):
1570+
import json
1571+
1572+
(tests_dir / f"{stem}.json").write_text(json.dumps([list(e) for e in entries]))
1573+
1574+
def test_returns_device_files_for_matching_slug(self, tmp_path):
1575+
tests_dir = tmp_path / "tests"
1576+
tests_dir.mkdir()
1577+
self._write_json(tests_dir, "known-slugs", [("nokia-7750-sr-7s", "device-types/Nokia/7750-SR-7s.yaml")])
1578+
self._write_json(tests_dir, "known-modules", [])
1579+
self._write_json(tests_dir, "known-racks", [])
1580+
1581+
repo = self._make_repo()
1582+
repo.repo_path = str(tmp_path)
1583+
repo.cwd = ""
1584+
1585+
result = repo.resolve_slug_files(["nokia-7750-sr-7s"])
1586+
1587+
assert result is not None
1588+
expected_path = str(tmp_path / "device-types" / "Nokia" / "7750-SR-7s.yaml")
1589+
assert expected_path in result["device_files"]["nokia"]
1590+
1591+
def test_matching_module_json_adds_vendor_slug(self, tmp_path):
1592+
tests_dir = tmp_path / "tests"
1593+
tests_dir.mkdir()
1594+
self._write_json(tests_dir, "known-slugs", [])
1595+
self._write_json(tests_dir, "known-modules", [("7750 line card", "module-types/Nokia")])
1596+
self._write_json(tests_dir, "known-racks", [])
1597+
1598+
repo = self._make_repo()
1599+
repo.repo_path = str(tmp_path)
1600+
repo.cwd = ""
1601+
1602+
result = repo.resolve_slug_files(["7750"])
1603+
1604+
assert result["module_vendors"] == {"nokia"}
1605+
1606+
def test_json_preferred_over_pickle(self, tmp_path):
1607+
"""When both formats exist, JSON wins — pickle here would yield a different match."""
1608+
import pickle
1609+
1610+
tests_dir = tmp_path / "tests"
1611+
tests_dir.mkdir()
1612+
self._write_json(tests_dir, "known-slugs", [("nokia-7750-sr-7s", "device-types/Nokia/7750-SR-7s.yaml")])
1613+
(tests_dir / "known-slugs.pickle").write_bytes(
1614+
pickle.dumps({("cisco-catalyst-9200", "device-types/Cisco/Catalyst-9200.yaml")})
1615+
)
1616+
self._write_json(tests_dir, "known-modules", [])
1617+
self._write_json(tests_dir, "known-racks", [])
1618+
1619+
repo = self._make_repo()
1620+
repo.repo_path = str(tmp_path)
1621+
repo.cwd = ""
1622+
1623+
result = repo.resolve_slug_files(["nokia"])
1624+
1625+
assert "nokia" in result["device_files"]
1626+
assert "cisco" not in result["device_files"]
1627+
1628+
def test_corrupted_device_json_returns_none(self, tmp_path):
1629+
tests_dir = tmp_path / "tests"
1630+
tests_dir.mkdir()
1631+
(tests_dir / "known-slugs.json").write_text("{not valid json")
1632+
1633+
repo = self._make_repo()
1634+
repo.repo_path = str(tmp_path)
1635+
repo.cwd = ""
1636+
1637+
assert repo.resolve_slug_files(["nokia"]) is None

0 commit comments

Comments
 (0)