diff --git a/.gitignore b/.gitignore index 195708e2d..11a147a6e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ .DS_Store -*.pyc +__pycache__ *.swp debug .vscode/* diff --git a/frigate/config.py b/frigate/config.py index ad96cebef..af3ed0753 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np +import yaml from pydantic import ( BaseModel, ConfigDict, @@ -41,11 +42,11 @@ from frigate.ffmpeg_presets import ( ) from frigate.plus import PlusApi from frigate.util.builtin import ( + NoDuplicateKeysLoader, deep_merge, escape_special_characters, generate_color_palette, get_ffmpeg_arg_list, - load_config_with_no_duplicates, ) from frigate.util.config import StreamInfoRetriever, get_relative_coordinates from frigate.util.image import create_mask @@ -1764,7 +1765,7 @@ class FrigateConfig(FrigateBaseModel): raw_config = f.read() if config_file.endswith(YAML_EXT): - config = load_config_with_no_duplicates(raw_config) + config = yaml.load(raw_config, NoDuplicateKeysLoader) elif config_file.endswith(".json"): config = json.loads(raw_config) @@ -1772,5 +1773,5 @@ class FrigateConfig(FrigateBaseModel): @classmethod def parse_raw(cls, raw_config): - config = load_config_with_no_duplicates(raw_config) + config = yaml.load(raw_config, NoDuplicateKeysLoader) return cls.model_validate(config) diff --git a/frigate/test/test_config.py b/frigate/test/test_config.py index c703de893..1f0ff086d 100644 --- a/frigate/test/test_config.py +++ b/frigate/test/test_config.py @@ -4,13 +4,14 @@ import unittest from unittest.mock import patch import numpy as np +import yaml from pydantic import ValidationError from frigate.config import BirdseyeModeEnum, FrigateConfig from frigate.const import MODEL_CACHE_DIR from frigate.detectors import DetectorTypeEnum from frigate.plus import PlusApi -from frigate.util.builtin import deep_merge, load_config_with_no_duplicates +from frigate.util.builtin import NoDuplicateKeysLoader, deep_merge class TestConfig(unittest.TestCase): @@ -1537,7 +1538,7 @@ class TestConfig(unittest.TestCase): """ self.assertRaises( - ValueError, lambda: load_config_with_no_duplicates(raw_config) + ValueError, lambda: yaml.load(raw_config, NoDuplicateKeysLoader) ) def test_object_filter_ratios_work(self): diff --git a/frigate/util/builtin.py b/frigate/util/builtin.py index 2c3051fc0..413c8d8ce 100644 --- a/frigate/util/builtin.py +++ b/frigate/util/builtin.py @@ -89,32 +89,31 @@ def deep_merge(dct1: dict, dct2: dict, override=False, merge_lists=False) -> dic return merged -def load_config_with_no_duplicates(raw_config) -> dict: - """Get config ensuring duplicate keys are not allowed.""" +class NoDuplicateKeysLoader(yaml.loader.SafeLoader): + """A yaml SafeLoader that disallows duplicate keys""" - # https://stackoverflow.com/a/71751051 - # important to use SafeLoader here to avoid RCE - class PreserveDuplicatesLoader(yaml.loader.SafeLoader): - pass + def construct_mapping(self, node, deep=False): + mapping = super().construct_mapping(node, deep=deep) - def map_constructor(loader, node, deep=False): - keys = [loader.construct_object(node, deep=deep) for node, _ in node.value] - vals = [loader.construct_object(node, deep=deep) for _, node in node.value] - key_count = Counter(keys) - data = {} - for key, val in zip(keys, vals): - if key_count[key] > 1: - raise ValueError( - f"Config input {key} is defined multiple times for the same field, this is not allowed." + if len(node.value) != len(mapping): + # There's a duplicate key somewhere. Find it. + duplicate_keys = [ + key + for key, count in Counter( + self.construct_object(key, deep=deep) for key, _ in node.value ) - else: - data[key] = val - return data + if count > 1 + ] - PreserveDuplicatesLoader.add_constructor( - yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor - ) - return yaml.load(raw_config, PreserveDuplicatesLoader) + # This might be possible if PyYAML's construct_mapping() changes the node + # afterwards for some reason? I don't see why, but better safe than sorry. + assert len(duplicate_keys) > 0 + + raise ValueError( + f"Config field duplicates are not allowed, the following fields are duplicated in the config: {', '.join(duplicate_keys)}" + ) + + return mapping def clean_camera_user_pass(line: str) -> str: