Suport per a comptes empresa i per a permetre afegir multiples comptes encara que comparteixin el mateix NIF d'usuari autoritzat

custom_components/aigues_barcelona/__init__.py

@@ -13,6 +13,7 @@
 
 from .api import AiguesApiClient
 from .const import DOMAIN
+from .const import CONF_COMPANY_IDENTIFICATOR
 from .service import async_setup as setup_service
 
 # from homeassistant.exceptions import ConfigEntryNotReady
@@ -23,7 +24,11 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # TODO Change after fixing Recaptcha.
-    api = AiguesApiClient(entry.data[CONF_USERNAME], entry.data[CONF_PASSWORD])
+    api = AiguesApiClient(
+        entry.data[CONF_USERNAME],
+        entry.data[CONF_PASSWORD],
+        company_identification=entry.data.get(CONF_COMPANY_IDENTIFICATOR),
+    )
     api.set_token(entry.data.get(CONF_TOKEN))
 
     if api.is_token_expired():

custom_components/aigues_barcelona/api.py

@@ -16,7 +16,12 @@
 
 class AiguesApiClient:
     def __init__(
-        self, username, password, contract=None, session: requests.Session = None
+        self,
+        username,
+        password,
+        contract=None,
+        session: requests.Session = None,
+        company_identification=None,
     ):
         if session is None:
             session = requests.Session()
@@ -33,6 +38,7 @@ def __init__(
         self._username = username
         self._password = password
         self._contract = contract
+        self._company_identification = company_identification
         self.last_response = None
 
     def _generate_url(self, path, query) -> str:
@@ -44,11 +50,9 @@ def _generate_url(self, path, query) -> str:
     def _return_token_field(self, key):
         token = self.cli.cookies.get_dict().get(API_COOKIE_TOKEN)
         if not token:
-            _LOGGER.warning("Token login missing")
             return False
 
         data = token.split(".")[1]
-        _LOGGER.debug(data)
         # add padding to avoid failures
         data = base64.urlsafe_b64decode(data + "==")
 
@@ -59,35 +63,48 @@ def _query(self, path, query=None, json=None, headers=None, method="GET"):
             headers = dict()
         headers = {**self.headers, **headers}
 
-        resp = self.cli.request(
-            method=method,
-            url=self._generate_url(path, query),
-            json=json,
-            headers=headers,
-            timeout=TIMEOUT,
-        )
-        _LOGGER.debug(f"Query done with code {resp.status_code}")
-        msg = resp.text
-        self.last_response = resp.text
-        if len(msg) > 5 and (msg.startswith("{") or msg.startswith("[")):
-            msg = resp.json()
-            if isinstance(msg, list) and len(msg) == 1:
-                msg = msg[0]
-            self.last_response = msg.copy()
-            msg = msg.get("message", resp.text)
-
-        if resp.status_code == 500:
-            raise Exception(f"Server error: {msg}")
-        if resp.status_code == 404:
-            raise Exception(f"Not found: {msg}")
-        if resp.status_code == 401:
-            raise Exception(f"Denied: {msg}")
-        if resp.status_code == 400:
-            raise Exception(f"Bad response: {msg}")
-        if resp.status_code == 429:
-            raise Exception(f"Rate-Limited: {msg}")
-
-        return resp
+        try:
+            resp = self.cli.request(
+                method=method,
+                url=self._generate_url(path, query),
+                json=json,
+                headers=headers,
+                timeout=TIMEOUT,
+            )
+            _LOGGER.debug(f"Query done with code {resp.status_code}")
+
+            # Store raw response first
+            self.last_response = resp.text
+
+            # Try to parse JSON response if possible
+            try:
+                if resp.text and len(resp.text) > 0:
+                    msg = resp.json()
+                    if isinstance(msg, list) and len(msg) == 1:
+                        msg = msg[0]
+                    self.last_response = msg
+            except json.JSONDecodeError:
+                msg = resp.text
+                _LOGGER.debug(f"Response is not JSON: {msg}")
+
+            if resp.status_code == 503:
+                raise Exception("Service temporarily unavailable")
+            if resp.status_code == 500:
+                raise Exception(f"Server error: {msg}")
+            if resp.status_code == 404:
+                raise Exception(f"Not found: {msg}")
+            if resp.status_code == 401:
+                raise Exception(f"Denied: {msg}")
+            if resp.status_code == 400:
+                raise Exception(f"Bad response: {msg}")
+            if resp.status_code == 429:
+                raise Exception(f"Rate-Limited: {msg}")
+
+            return resp
+
+        except requests.exceptions.RequestException as e:
+            _LOGGER.error(f"Request failed: {str(e)}")
+            raise Exception(f"Request failed: {str(e)}")
 
     def login(self, user=None, password=None, recaptcha=None):
         if user is None:
@@ -102,7 +119,7 @@ def login(self, user=None, password=None, recaptcha=None):
         query = {"lang": "ca", "recaptchaClientResponse": recaptcha}
         body = {
             "scope": "ofex",
-            "companyIdentification": "",
+            "companyIdentification": self._company_identification or "",
             "userIdentification": user,
             "password": password,
         }
@@ -113,15 +130,12 @@ def login(self, user=None, password=None, recaptcha=None):
 
         r = self._query(path, query, body, headers, method="POST")
 
-        _LOGGER.debug(r)
         error = r.json().get("errorMessage", None)
         if error:
-            _LOGGER.warning(error)
             return False
 
         access_token = r.json().get("access_token", None)
         if not access_token:
-            _LOGGER.warning("Access token missing")
             return False
 
         return True
@@ -168,16 +182,28 @@ def profile(self, user=None):
         assert r.json().get("user_data"), "User data missing"
         return r.json()
 
-    def contracts(self, user=None, status=["ASSIGNED", "PENDING"]):
+    def contracts(self, user=None, status=None):
         if user is None:
             user = self._return_token_field("name")
-        if isinstance(status, str):
-            status = [status]
+        if status is None:
+            status = ["ASSIGNED", "PENDING"]
 
         path = "/ofex-contracts-api/contracts"
-        query = {"lang": "ca", "userId": user, "clientId": user}
-        for idx, stat in enumerate(status):
-            query[f"assignationStatus[{str(idx)}]"] = stat.upper()
+        query = {
+            "userId": user,
+            "clientId": self._company_identification or user,
+            "lang": "ca",
+        }
+
+        # Add each status as a separate query parameter
+        for stat in status:
+            if "assignationStatus" not in query:
+                query["assignationStatus"] = stat.upper()
+            else:
+                # Append additional status values
+                query["assignationStatus"] = (
+                    f"{query['assignationStatus']}&assignationStatus={stat.upper()}"
+                )
 
         r = self._query(path, query)
 
@@ -221,31 +247,22 @@ def invoices_debt(self, contract=None, user=None):
         return self.invoices(contract, user, last_months=0, mode="DEBT")
 
     def consumptions(
-        self, date_from, date_to=None, contract=None, user=None, frequency="HOURLY"
+        self, date_from, date_to, contract=None, user=None, frequency="HOURLY"
     ):
         if user is None:
-            user = self._return_token_field("name")
+            user = self._username
         if contract is None:
-            contract = self.first_contract
-        if frequency not in ["HOURLY", "DAILY"]:
-            raise ValueError(f"Invalid {frequency=}")
-
-        if date_to is None:
-            date_to = date_from + datetime.timedelta(days=1)
-        if isinstance(date_from, datetime.date):
-            date_from = date_from.strftime("%d-%m-%Y")
-        if isinstance(date_to, datetime.date):
-            date_to = date_to.strftime("%d-%m-%Y")
+            contract = self._contract
 
         path = "/ofex-water-consumptions-api/meter/consumptions"
         query = {
             "consumptionFrequency": frequency,
             "contractNumber": contract,
-            "clientId": user,
+            "clientId": self._company_identification or user,
             "userId": user,
             "lang": "ca",
-            "fromDate": date_from,
-            "toDate": date_to,
+            "fromDate": date_from.strftime("%d-%m-%Y"),
+            "toDate": date_to.strftime("%d-%m-%Y"),
             "showNegativeValues": "false",
         }