fix refresh expiring

Author: NelsonDane

public_invest_api/endpoints.py

@@ -51,8 +51,7 @@ def get_pending_orders_url(self, account_uuid):
     def cancel_pending_order_url(self, account_uuid, order_id):
         return f"{self.ordergateway}/accounts/{account_uuid}/orders/{order_id}"
 
-    @staticmethod
-    def build_headers(auth=None, prodApi=False):
+    def build_headers(self, auth=None, prodApi=False):
         headers = {
             "authority": "public.com",
             "accept": "*/*",
@@ -76,7 +75,7 @@ def build_headers(auth=None, prodApi=False):
         if auth is not None:
             headers["authorization"] = auth
         if prodApi:
-            headers["authority"] = "prod-api.154310543964.hellopublic.com"
+            headers["authority"] = self.prodapi.replace("https://", "")
             headers["sec-fetch-site"] = "cross-site"
         return headers
 

public_invest_api/public.py

@@ -28,6 +28,25 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+def refresh_check(func):
+    def wrapper(self, *args, **kwargs):
+        """
+        A wrapper function that checks if the access token needs to be refreshed
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+        Raises:
+            Exception: If the access token refresh fails (i.e., response status code is not 200).
+        Returns:
+            The result of the function `func` if the access token does not need to be refreshed.
+        """
+        if self.expires_at is not None and datetime.now().timestamp() > self.expires_at:
+            self._refresh_token()
+        return func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class Public:
     def __init__(self, filename=None, path=None):
         self.session = requests.Session()
@@ -38,6 +57,7 @@ def __init__(self, filename=None, path=None):
         self.account_number = None
         self.all_login_info = None
         self.timeout = 10
+        self.expires_at = None
         self.filename = "public_credentials.pkl"
         if filename is not None:
             self.filename = filename
@@ -110,6 +130,15 @@ def login(self, username=None, password=None, wait_for_2fa=True, code=None) -> d
                 data=payload,
                 timeout=self.timeout,
             )
+            if response.status_code != 200:
+                # Perhaps cookies are expired
+                self._clear_cookies()
+                response = self.session.post(
+                    self.endpoints.login_url(),
+                    headers=headers,
+                    data=payload,
+                    timeout=self.timeout,
+                )
             if response.status_code != 200:
                 print(response.text)
                 raise Exception("Login failed, check credentials")
@@ -145,12 +174,13 @@ def login(self, username=None, password=None, wait_for_2fa=True, code=None) -> d
         self.access_token = response["loginResponse"]["accessToken"]
         self.account_uuid = response["loginResponse"]["accounts"][0]["accountUuid"]
         self.account_number = response["loginResponse"]["accounts"][0]["account"]
+        self.expires_at = (int(response["loginResponse"]["serverTime"]) / 1000) + int(response["loginResponse"]["expiresIn"])
         self.all_login_info = response
         self._save_cookies()
         return response
 
     @login_required
-    def refresh_token(self) -> dict:
+    def _refresh_token(self) -> dict:
         """
         Refreshes the access token by making a POST request to the refresh URL.
         Returns:
@@ -166,10 +196,13 @@ def refresh_token(self) -> dict:
             raise Exception("Token refresh failed")
         response = response.json()
         self.access_token = response["accessToken"]
+        self.expires_at = (int(response["serverTime"]) / 1000) + int(response["expiresIn"])
+        self.account_uuid = response["accounts"][0]["accountUuid"]
         self._save_cookies()
         return response
 
     @login_required
+    @refresh_check
     def get_portfolio(self) -> dict:
         """
         Gets the user's portfolio by making a GET request to the portfolio URL.
@@ -178,7 +211,7 @@ def get_portfolio(self) -> dict:
         Raises:
             Exception: If the portfolio request fails (i.e., response status code is not 200).
         """
-        headers = self.endpoints.build_headers(self.access_token)
+        headers = self.endpoints.build_headers(self.access_token, prodApi=True)
         portfolio = self.session.get(
             self.endpoints.portfolio_url(self.account_uuid),
             headers=headers,
@@ -222,6 +255,7 @@ def _history_filter_date(date: str) -> dict:
         }
 
     @login_required
+    @refresh_check
     def get_account_history(
         self,
         date="all",
@@ -423,6 +457,7 @@ def get_account_cash(self) -> float:
         return account_info["equity"]["cash"]
 
     @login_required
+    @refresh_check
     def get_symbol_price(self, symbol) -> float:
         """
         Gets the price of a stock by making a GET request to the quote URL.
@@ -445,6 +480,7 @@ def get_symbol_price(self, symbol) -> float:
         return response.json()["price"]
 
     @login_required
+    @refresh_check
     def get_order_quote(self, symbol) -> dict:
         """
         Gets a quote for an order by making a GET request to the order quote URL.
@@ -466,6 +502,7 @@ def get_order_quote(self, symbol) -> dict:
         return response.json()
 
     @login_required
+    @refresh_check
     def place_order(
         self,
         symbol,
@@ -587,6 +624,7 @@ def place_order(
         return check_response
 
     @login_required
+    @refresh_check
     def get_pending_orders(self) -> dict:
         """
         Gets the user's pending orders by making a GET request to the pending orders URL.
@@ -606,6 +644,7 @@ def get_pending_orders(self) -> dict:
         return response.json()
 
     @login_required
+    @refresh_check
     def cancel_order(self, order_id) -> dict:
         """
         Cancels an order by making a DELETE request to the cancel order URL.

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="public_invest_api",
-    version="1.2.1",
+    version="1.2.2",
     description="Unofficial Public.com Invest API written in Python Requests",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",