diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py index 3a237f5..164d706 100755 --- a/check_home_assistant_state.py +++ b/check_home_assistant_state.py @@ -56,6 +56,8 @@ class Entities(nagiosplugin.Resource): attribute: Optional[str] friendly_name: bool ignore_missing: bool + include: list[str] + exclude: list[str] min: Optional[float] = None max: Optional[float] = None @@ -118,8 +120,11 @@ class Entities(nagiosplugin.Resource): def probe(self): response = self.hass_get("/api/states") for state in response: - if state["attributes"].get("device_class") == self.device_class and all( - state["attributes"].get(k) == v for k, v in self.filters + if ( + state["attributes"].get("device_class") == self.device_class + and all(state["attributes"].get(k) == v for k, v in self.filters) + and (len(self.include) == 0 or state["entity_id"] in self.include) + and state["entity_id"] not in self.exclude ): yield from self._state_to_metric(state) @@ -302,6 +307,20 @@ def main(): nargs="*", help="filter by 'attribute=value' (may be specified multiple times)", ) + common_args.add_argument( + "-i", + "--include", + default=[], + nargs="*", + help="explicitly include entities by id. Other entities will not be considered if this is specified. Listed entities must also match the specified device class", + ) + common_args.add_argument( + "-e", + "--exclude", + default=[], + nargs="*", + help="exclude entities by id", + ) common_args.add_argument( "--friendly", action="store_true", @@ -366,19 +385,24 @@ def main(): ) args = argp.parse_args() + parsed_common_args = { + "device_class": args.device_class, + "attribute": args.attribute, + "filters": args.filter, + "include": args.include, + "exclude": args.exclude, + "friendly_name": args.friendly, + "ignore_missing": args.ignore_missing, + } if args.subparser_name == "scalar": check = nagiosplugin.Check( Entities( url=args.url, token=args.token, - device_class=args.device_class, numeric=True, min=args.min, max=args.max, - filters=args.filter, - attribute=args.attribute, - friendly_name=args.friendly, - ignore_missing=args.ignore_missing, + **parsed_common_args, ), ScalarOrUnknownContext(args.device_class, args.warning, args.critical), ) @@ -387,12 +411,8 @@ def main(): Entities( url=args.url, token=args.token, - device_class=args.device_class, numeric=False, - filters=args.filter, - attribute=args.attribute, - friendly_name=args.friendly, - ignore_missing=args.ignore_missing, + **parsed_common_args, ), RegexContext(args.device_class, args.ok, args.warning, args.critical), )