From cedc8a73e728b55c3d66232f42e6f5648f7a2ba0 Mon Sep 17 00:00:00 2001
From: Sam Kleiner <sam-kleiner@users.noreply.github.com>
Date: Sat, 2 Dec 2023 01:20:48 -0500
Subject: [PATCH] allow select device

---
 waterfurnace/waterfurnace.py | 50 +++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 3 deletions(-)

diff --git a/waterfurnace/waterfurnace.py b/waterfurnace/waterfurnace.py
index 034b0e8..aee28e9 100644
--- a/waterfurnace/waterfurnace.py
+++ b/waterfurnace/waterfurnace.py
@@ -2,8 +2,8 @@
 
 """Main module."""
 import copy
-import logging
 import json
+import logging
 import threading
 import time
 
@@ -97,9 +97,11 @@ class WFError(WFException):
 
 class WaterFurnace(object):
 
-    def __init__(self, user, passwd, max_fails=5):
+    def __init__(self, user, passwd, max_fails=5, location=0, device=0):
         self.user = user
         self.passwd = passwd
+        self.location = location
+        self.device = device
         self.gwid = None
         self.sessionid = None
         self.tid = 0
@@ -148,7 +150,49 @@ def _login_ws(self):
         recv = self.ws.recv()
         data = json.loads(recv)
         _LOGGER.debug("Login response: %s" % data)
-        self.gwid = data["locations"][0]["gateways"][0]["gwid"]
+
+        locations = data["locations"]
+        location = None
+
+        if isinstance(self.location, int):
+            try:
+                location = locations[self.location]
+            except Exception:
+                raise WFError("Location index out of range. Max index is {}".format(len(locations) - 1))
+        elif isinstance(self.location, str):
+            for index, location_data in enumerate(locations):
+                location_description = location_data.get("description")
+                if location_description == self.location:
+                    location = locations[index]
+                    break
+
+            if not location:
+                raise WFError("Unable to find location: {}".format(self.location))
+        else:
+            raise WFError("Unknown location type ({}): {}. Should be int or str".format(type(self.location), self.location))
+
+        gateways = location["gateways"]
+        device = None
+
+        if isinstance(self.device, int):
+            try:
+                device = gateways[self.device]
+            except Exception:
+                raise WFError("Device index out of range. Max index is {}".format(len(gateways) - 1))
+        elif isinstance(self.device, str):
+            for index, gateway_data in enumerate(gateways):
+                gateway_gwid = gateway_data.get("gwid")
+                gateway_description = gateway_data.get("description")
+                if gateway_gwid == self.device or gateway_description == self.device:
+                    device = gateways[index]
+                    break
+
+            if not device:
+                raise WFError("Unable to find device: {}".format(self.device))
+        else:
+            raise WFError("Unknown device type ({}): {}. Should be int or str".format(type(self.location), self.location))
+
+        self.gwid = device["gwid"]
         self.next_tid()
 
     def login(self):