From c7f1b9c7c97faaac11f3466ef4e8978d73c4f2c8 Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Sat, 11 Mar 2023 00:41:22 -0500
Subject: [PATCH] Add explicit include/exclude by entity id

---
 check_home_assistant_state.py | 44 +++++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 12 deletions(-)

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),
         )