From 8d88430577ff4a4ca9259b9bf794f5364b28dd8c Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Sat, 26 Nov 2022 19:44:08 -0500
Subject: [PATCH] Add cli arg for filtering by arbitrary attributes

---
 check_home_assistant_state.conf |  6 +++++-
 check_home_assistant_state.py   | 19 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/check_home_assistant_state.conf b/check_home_assistant_state.conf
index 9ad09be..242f34a 100644
--- a/check_home_assistant_state.conf
+++ b/check_home_assistant_state.conf
@@ -29,5 +29,9 @@ object CheckCommand "home_assistant_state" {
       value = "$home_assistant_state_max$"
       description = "Max for performance data"
     }
+    "--filter" = {
+      value = "$home_assistant_state_filter$"
+      description = "Filter by 'attribute=value' (can be an array)"
+    }
   }
-}
\ No newline at end of file
+}
diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index 3439b9a..7fc7688 100755
--- a/check_home_assistant_state.py
+++ b/check_home_assistant_state.py
@@ -19,6 +19,7 @@ class Entities(nagiosplugin.Resource):
     device_class: str
     min: float
     max: float
+    filters: list[str]
 
     def hass_get(self, endpoint: str) -> requests.Response:
         headers = {
@@ -43,7 +44,9 @@ 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:
+            if state["attributes"].get("device_class") == self.device_class and all(
+                state["attributes"].get(k) == v for k, v in self.filters
+            ):
                 if state["state"] == "unavailable":
                     _log.info(f"{state['entity_id']} unavailable")
                 yield nagiosplugin.Metric(
@@ -56,6 +59,11 @@ class Entities(nagiosplugin.Resource):
                 )
 
 
+def key_value(arg: str):
+    k, _, v = arg.partition("=")
+    return k, v
+
+
 @nagiosplugin.guarded
 def main():
     argp = argparse.ArgumentParser(description=__doc__)
@@ -98,6 +106,14 @@ def main():
         type=float,
         help="Max for performance data",
     )
+    argp.add_argument(
+        "-f",
+        "--filter",
+        type=key_value,
+        default=[],
+        nargs="*",
+        help="Filter by 'attribute=value'",
+    )
 
     args = argp.parse_args()
 
@@ -108,6 +124,7 @@ def main():
             args.device_class,
             args.min,
             args.max,
+            args.filter,
         ),
         nagiosplugin.ScalarContext(args.device_class, args.warning, args.critical),
     )