Load and Save token to disk

src/iec_api_client.py

@@ -24,15 +24,15 @@ def __init__(self, user_id, automatically_login: bool = False):
         user_id (str): The user ID (SSN) to be associated with the instance.
         automatically_login (bool): Whether to automatically log in the user. Default is False.
         """
-        self._state_token = None  # Token for maintaining the state of the user's session
-        self._factor_id = None  # Factor ID for multi-factor authentication
-        self._session_token = None  # Token for maintaining the user's session
-        self.logged_in = None  # Flag to indicate if the user is logged in
-        self._token = None  # Token for authentication
-        self._user_id = user_id  # User ID associated with the instance
-        self._login_response = None  # Response from the login attempt
-        self._bp_number = None  # BP Number associated with the instance
-        self._contract_id = None  # Contract ID associated with the instance
+        self._state_token: str | None = None  # Token for maintaining the state of the user's session
+        self._factor_id: str | None = None  # Factor ID for multi-factor authentication
+        self._session_token: str | None = None  # Token for maintaining the user's session
+        self.logged_in: bool = False  # Flag to indicate if the user is logged in
+        self._token: JWT = JWT(access_token="", refresh_token="", token_type="", expires_in=0, scope="", id_token="")  # Token for authentication
+        self._user_id: str = user_id  # User ID associated with the instance
+        self._login_response: str | None = None  # Response from the login attempt
+        self._bp_number: str | None = None  # BP Number associated with the instance
+        self._contract_id: str | None = None  # Contract ID associated with the instance
         if automatically_login:
             self.login_with_id()  # Attempt to log in automatically if specified
 
@@ -72,8 +72,8 @@ def override_token(self, token):
         :return: None
         """
         logger.debug("Overriding jwt.py token: %s", token)
-        self._token = JWT()
-        self._token.access_token = token
+        self._token = JWT(access_token=token, refresh_token="", token_type="", expires_in=0, scope="", id_token="")
+        self._token.id_token = token
         self.logged_in = True
 
     def get_customer(self) -> Customer:
@@ -99,13 +99,13 @@ def get_contracts(self, bp_number: str = None) -> list[Contract]:
         if not bp_number:
             bp_number = self._bp_number
 
-        get_contract_response = data.get_contract(self._token, bp_number)
+        get_contract_response = data.get_contract(self._token.id_token, bp_number)
         if get_contract_response and get_contract_response.data:
             contracts = get_contract_response.data.contracts
             if contracts and len(contracts) > 0:
                 self._contract_id = contracts[0].contract_id
             return contracts
-        return list()
+        return []
 
     def get_last_meter_reading(self, bp_number: str, contract_id: str) -> MeterReadings | None:
         """
@@ -124,7 +124,7 @@ def get_last_meter_reading(self, bp_number: str, contract_id: str) -> MeterReadi
         if not contract_id:
             contract_id = self._contract_id
 
-        response = data.get_last_meter_reading(self._token, bp_number, contract_id)
+        response = data.get_last_meter_reading(self._token.id_token, bp_number, contract_id)
         if response and response.data:
             return response.data
         return None
@@ -147,11 +147,12 @@ def get_electric_bill(self, bp_number: str, contract_id: str) -> Invoices | None
         if not contract_id:
             contract_id = self._contract_id
 
-        response = data.get_electric_bill(self._token, bp_number, contract_id)
+        response = data.get_electric_bill(self._token.id_token, bp_number, contract_id)
         if response.data:
             return response.data
         return None
 
+
     def get_remote_reading(self, meter_serial_number: str, meter_code: int, last_invoice_date: str, from_date: str,
                            resolution: int) -> RemoteReadingResponse:
         """
@@ -167,16 +168,29 @@ def get_remote_reading(self, meter_serial_number: str, meter_code: int, last_inv
             RemoteReadingResponse: The response containing the remote reading.
         """
         self.check_token()
-        return data.get_remote_reading(self._token, meter_serial_number, meter_code, last_invoice_date, from_date,
+        return data.get_remote_reading(self._token.id_token, meter_serial_number, meter_code, last_invoice_date, from_date,
                                        resolution)
 
     def check_token(self):
         """
         Check the validity of the jwt.py token and refresh in the case of expired signature errors.
         """
         try:
-            jwt.decode(self._token.access_token, options={"verify_signature": False}, algorithms=["RS256"])
+            jwt.decode(self._token.id_token, options={"verify_signature": False}, algorithms=["RS256"])
         except jwt.exceptions.ExpiredSignatureError:
             logger.debug("jwt.py token expired, retrying login")
             self.logged_in = False
             login.refresh_token(self._token)
+
+    def load_token(self, file_path: str = "token.json"):
+        """
+        Load token from file.
+        """
+        self._token = login.load_token_from_file(file_path)
+        self.logged_in = True
+
+    def save_token(self, file_path: str = "token.json"):
+        """
+        Save token to file.
+        """
+        login.save_token_to_file(self._token, file_path)

src/login.py

@@ -5,6 +5,7 @@
 import string
 from typing import Tuple
 
+import jwt
 import pkce
 import requests
 
@@ -161,6 +162,20 @@ def refresh_token(token: JWT) -> JWT | None:
     return None
 
 
+def save_token_to_file(token: JWT, path: str = "token.json") -> None:
+    """Save token to file."""
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(token.to_dict(), f)
+
+
+def load_token_from_file(path: str = "token.json") -> JWT:
+    """Load token from file."""
+    with open(path, "r", encoding="utf-8") as f:
+        jwt_data = JWT.from_dict(json.load(f))
+        jwt.decode(jwt_data.access_token, options={"verify_signature": False}, algorithms=["RS256"])
+        return jwt_data
+
+
 class IECLoginError(Exception):
     """Exception raised for errors in the IEC Login.
 

src/main.py

@@ -16,7 +16,17 @@
     try:
         # Example of usage
         client = IecApiClient("123456789")
-        client.manual_login()
+
+        token_json_file = "token.json"
+        if os.path.exists(token_json_file):
+            client.load_token(token_json_file)
+        else:
+            client.login_with_id()
+            otp = input("Enter your ID Number: ")
+            client.verify_otp(otp)
+            client.save_token(token_json_file)
+
+        # client.manual_login()
         customer = client.get_customer()
         print(customer)