firstrade 0.0.21__py3-none-any.whl → 0.0.30__py3-none-any.whl

firstrade/account.py

@@ -1,17 +1,66 @@
 import os
 import pickle
-import re
 
+import pyotp
 import requests
-from bs4 import BeautifulSoup
 
 from firstrade import urls
+from firstrade.exceptions import (
+    AccountResponseError,
+    LoginRequestError,
+    LoginResponseError,
+)
 
 
 class FTSession:
-    """Class creating a session for Firstrade."""
+    """
+    Class creating a session for Firstrade.
 
-    def __init__(self, username, password, pin, profile_path=None):
+    This class handles the creation and management of a session for logging into the Firstrade platform.
+    It supports multi-factor authentication (MFA) and can save session cookies for persistent logins.
+
+    Attributes:
+        username (str): Firstrade login username.
+        password (str): Firstrade login password.
+        pin (str): Firstrade login pin.
+        email (str, optional): Firstrade MFA email.
+        phone (str, optional): Firstrade MFA phone number.
+        mfa_secret (str, optional): Secret key for generating MFA codes.
+        profile_path (str, optional): The path where the user wants to save the cookie pkl file.
+        t_token (str, optional): Token used for MFA.
+        otp_options (dict, optional): Options for OTP (One-Time Password) if MFA is enabled.
+        login_json (dict, optional): JSON response from the login request.
+        session (requests.Session): The requests session object used for making HTTP requests.
+
+    Methods:
+        __init__(username, password, pin=None, email=None, phone=None, mfa_secret=None, profile_path=None):
+            Initializes a new instance of the FTSession class.
+        login():
+            Validates and logs into the Firstrade platform.
+        login_two(code):
+            Finishes the login process to the Firstrade platform. When using email or phone mfa.
+        delete_cookies():
+            Deletes the session cookies.
+        _load_cookies():
+            Checks if session cookies were saved and loads them.
+        _save_cookies():
+            Saves session cookies to a file.
+        _mask_email(email):
+            Masks the email for use in the API.
+        _handle_mfa():
+            Handles multi-factor authentication.
+    """
+
+    def __init__(
+        self,
+        username,
+        password,
+        pin=None,
+        email=None,
+        phone=None,
+        mfa_secret=None,
+        profile_path=None,
+    ):
         """
         Initializes a new instance of the FTSession class.
 
@@ -19,76 +68,110 @@ class FTSession:
             username (str): Firstrade login username.
             password (str): Firstrade login password.
             pin (str): Firstrade login pin.
-            persistent_session (bool, optional): Whether the user wants to save the session cookies.
+            email (str, optional): Firstrade MFA email.
+            phone (str, optional): Firstrade MFA phone number.
             profile_path (str, optional): The path where the user wants to save the cookie pkl file.
         """
         self.username = username
         self.password = password
         self.pin = pin
+        self.email = FTSession._mask_email(email) if email is not None else None
+        self.phone = phone
+        self.mfa_secret = mfa_secret
         self.profile_path = profile_path
+        self.t_token = None
+        self.otp_options = None
+        self.login_json = None
         self.session = requests.Session()
-        self.login()
 
     def login(self):
-        """Method to validate and login to the Firstrade platform."""
-        headers = urls.session_headers()
-        cookies = self.load_cookies()
-        cookies = requests.utils.cookiejar_from_dict(cookies)
-        self.session.cookies.update(cookies)
-        response = self.session.get(
-            url=urls.get_xml(), headers=urls.session_headers(), cookies=cookies
-        )
-        if response.status_code != 200:
-            raise Exception(
-                "Login failed. Check your credentials or internet connection."
-            )
-        if "/cgi-bin/sessionfailed?reason=6" in response.text:
-            self.session.get(url=urls.login(), headers=headers)
-            data = {
-                "redirect": "",
-                "ft_locale": "en-us",
-                "login.x": "Log In",
-                "username": r"" + self.username,
-                "password": r"" + self.password,
-                "destination_page": "home",
-            }
+        """
+        Validates and logs into the Firstrade platform.
 
-            self.session.post(
-                url=urls.login(),
-                headers=headers,
-                cookies=self.session.cookies,
-                data=data,
-            )
-            data = {
-                "destination_page": "home",
-                "pin": self.pin,
-                "pin.x": "++OK++",
-                "sring": "0",
-                "pin": self.pin,
-            }
+        This method sets up the session headers, loads cookies if available, and performs the login request.
+        It handles multi-factor authentication (MFA) if required.
+
+        Raises:
+            LoginRequestError: If the login request fails with a non-200 status code.
+            LoginResponseError: If the login response contains an error message.
+        """
+        self.session.headers = urls.session_headers()
+        ftat = self._load_cookies()
+        if ftat != "":
+            self.session.headers["ftat"] = ftat
+        response = self.session.get(url="https://api3x.firstrade.com/", timeout=10)
+        self.session.headers["access-token"] = urls.access_token()
+
+        data = {
+            "username": r"" + self.username,
+            "password": r"" + self.password,
+        }
 
-            self.session.post(
-                url=urls.pin(), headers=headers, cookies=self.session.cookies, data=data
-            )
-            self.save_cookies()
+        response = self.session.post(
+            url=urls.login(),
+            data=data,
+        )
+        self.login_json = response.json()
         if (
-            "/cgi-bin/sessionfailed?reason=6"
-            in self.session.get(
-                url=urls.get_xml(), headers=urls.session_headers(), cookies=cookies
-            ).text
+            "mfa" not in self.login_json
+            and "ftat" in self.login_json
+            and self.login_json["error"] == ""
         ):
-            raise Exception("Login failed. Check your credentials.")
+            self.session.headers["sid"] = self.login_json["sid"]
+            return False
+        self.t_token = self.login_json.get("t_token")
+        if self.mfa_secret is None:
+            self.otp_options = self.login_json.get("otp")
+        if response.status_code != 200:
+            raise LoginRequestError(response.status_code)
+        if self.login_json["error"] != "":
+            raise LoginResponseError(self.login_json["error"])
+        need_code = self._handle_mfa()
+        if self.login_json["error"] != "":
+            raise LoginResponseError(self.login_json["error"])
+        if need_code:
+            return True
+        self.session.headers["ftat"] = self.login_json["ftat"]
+        self.session.headers["sid"] = self.login_json["sid"]
+        self._save_cookies()
+        return False
+
+    def login_two(self, code):
+        """Method to finish login to the Firstrade platform."""
+        data = {
+            "otpCode": code,
+            "verificationSid": self.session.headers["sid"],
+            "remember_for": "30",
+            "t_token": self.t_token,
+        }
+        response = self.session.post(urls.verify_pin(), data=data)
+        self.login_json = response.json()
+        if self.login_json["error"] != "":
+            raise LoginResponseError(self.login_json["error"])
+        self.session.headers["ftat"] = self.login_json["ftat"]
+        self.session.headers["sid"] = self.login_json["sid"]
+        self._save_cookies()
 
-    def load_cookies(self):
+    def delete_cookies(self):
+        """Deletes the session cookies."""
+        if self.profile_path is not None:
+            path = os.path.join(self.profile_path, f"ft_cookies{self.username}.pkl")
+        else:
+            path = f"ft_cookies{self.username}.pkl"
+        os.remove(path)
+
+    def _load_cookies(self):
         """
         Checks if session cookies were saved.
 
         Returns:
-            Dict: Dictionary of cookies. Nom Nom
+            str: The saved session token.
         """
-        cookies = {}
-        directory = os.path.abspath(self.profile_path) if self.profile_path is not None else "."
-        
+
+        ftat = ""
+        directory = (
+            os.path.abspath(self.profile_path) if self.profile_path is not None else "."
+        )
         if not os.path.exists(directory):
             os.makedirs(directory)
 
@@ -96,10 +179,10 @@ class FTSession:
             if filename.endswith(f"{self.username}.pkl"):
                 filepath = os.path.join(directory, filename)
                 with open(filepath, "rb") as f:
-                    cookies = pickle.load(f)
-        return cookies
+                    ftat = pickle.load(f)
+        return ftat
 
-    def save_cookies(self):
+    def _save_cookies(self):
         """Saves session cookies to a file."""
         if self.profile_path is not None:
             directory = os.path.abspath(self.profile_path)
@@ -109,15 +192,76 @@ class FTSession:
         else:
             path = f"ft_cookies{self.username}.pkl"
         with open(path, "wb") as f:
-            pickle.dump(self.session.cookies.get_dict(), f)
-    
-    def delete_cookies(self):    
-        """Deletes the session cookies."""
-        if self.profile_path is not None:
-            path = os.path.join(self.profile_path, f"ft_cookies{self.username}.pkl")
-        else:
-            path = f"ft_cookies{self.username}.pkl"
-        os.remove(path)
+            ftat = self.session.headers.get("ftat")
+            pickle.dump(ftat, f)
+
+    @staticmethod
+    def _mask_email(email):
+        """
+        Masks the email for use in the API.
+
+        Args:
+            email (str): The email address to be masked.
+
+        Returns:
+            str: The masked email address.
+        """
+        local, domain = email.split("@")
+        masked_local = local[0] + "*" * 4
+        domain_name, tld = domain.split(".")
+        masked_domain = domain_name[0] + "*" * 4
+        return f"{masked_local}@{masked_domain}.{tld}"
+
+    def _handle_mfa(self):
+        """
+        Handles multi-factor authentication.
+
+        This method processes the MFA requirements based on the login response and user-provided details.
+
+        Raises:
+            LoginRequestError: If the MFA request fails with a non-200 status code.
+            LoginResponseError: If the MFA response contains an error message.
+        """
+        if not self.login_json["mfa"] and self.pin is not None:
+            data = {
+                "pin": self.pin,
+                "remember_for": "30",
+                "t_token": self.t_token,
+            }
+            response = self.session.post(urls.verify_pin(), data=data)
+            self.login_json = response.json()
+        elif not self.login_json["mfa"] and (
+            self.email is not None or self.phone is not None
+        ):
+            for item in self.otp_options:
+                if item["channel"] == "sms" and self.phone is not None:
+                    if self.phone in item["recipientMask"]:
+                        data = {
+                            "recipientId": item["recipientId"],
+                            "t_token": self.t_token,
+                        }
+                elif item["channel"] == "email" and self.email is not None:
+                    if self.email == item["recipientMask"]:
+                        data = {
+                            "recipientId": item["recipientId"],
+                            "t_token": self.t_token,
+                        }
+            response = self.session.post(urls.request_code(), data=data)
+        elif self.login_json["mfa"] and self.mfa_secret is not None:
+            mfa_otp = pyotp.TOTP(self.mfa_secret).now()
+            data = {
+                "mfaCode": mfa_otp,
+                "remember_for": "30",
+                "t_token": self.t_token,
+            }
+            response = self.session.post(urls.verify_pin(), data=data)
+        self.login_json = response.json()
+        if self.login_json["error"] == "":
+            if self.pin or self.mfa_secret is not None:
+                self.session.headers["sid"] = self.login_json["sid"]
+                return False
+            self.session.headers["sid"] = self.login_json["verificationSid"]
+            return True
 
     def __getattr__(self, name):
         """
@@ -146,74 +290,28 @@ class FTAccountData:
         self.session = session
         self.all_accounts = []
         self.account_numbers = []
-        self.account_statuses = []
-        self.account_balances = []
-        self.securities_held = {}
-        all_account_info = []
-        html_string = self.session.get(
-            url=urls.account_list(),
-            headers=urls.session_headers(),
-            cookies=self.session.cookies,
-        ).text
-        regex_accounts = re.findall(r"([0-9]+)-", html_string)
-
-        for match in regex_accounts:
-            self.account_numbers.append(match)
-
-        for account in self.account_numbers:
-            # reset cookies to base login cookies to run scripts
-            self.session.cookies.clear()
-            self.session.cookies.update(self.session.load_cookies())
-            # set account to get data for
-            data = {"accountId": account}
-            self.session.post(
-                url=urls.account_status(),
-                headers=urls.session_headers(),
-                cookies=self.session.cookies,
-                data=data,
-            )
-            # request to get account status data
-            data = {"req": "get_status"}
-            account_status = self.session.post(
-                url=urls.status(),
-                headers=urls.session_headers(),
-                cookies=self.session.cookies,
-                data=data,
-            ).json()
-            self.account_statuses.append(account_status["data"])
-            data = {"page": "bal", "account_id": account}
-            account_soup = BeautifulSoup(
-                self.session.post(
-                    url=urls.get_xml(),
-                    headers=urls.session_headers(),
-                    cookies=self.session.cookies,
-                    data=data,
-                ).text,
-                "xml",
-            )
-            balance = account_soup.find("total_account_value").text
-            self.account_balances.append(balance)
-            all_account_info.append(
-                {
-                    account: {
-                        "Balance": balance,
-                        "Status": {
-                            "primary": account_status["data"]["primary"],
-                            "domestic": account_status["data"]["domestic"],
-                            "joint": account_status["data"]["joint"],
-                            "ira": account_status["data"]["ira"],
-                            "hasMargin": account_status["data"]["hasMargin"],
-                            "opLevel": account_status["data"]["opLevel"],
-                            "p_country": account_status["data"]["p_country"],
-                            "mrgnStatus": account_status["data"]["mrgnStatus"],
-                            "opStatus": account_status["data"]["opStatus"],
-                            "margin_id": account_status["data"]["margin_id"],
-                        },
-                    }
-                }
-            )
-
-        self.all_accounts = all_account_info
+        self.account_balances = {}
+        response = self.session.get(url=urls.user_info())
+        self.user_info = response.json()
+        response = self.session.get(urls.account_list())
+        if response.status_code != 200 or response.json()["error"] != "":
+            raise AccountResponseError(response.json()["error"])
+        self.all_accounts = response.json()
+        for item in self.all_accounts["items"]:
+            self.account_numbers.append(item["account"])
+            self.account_balances[item["account"]] = item["total_value"]
+
+    def get_account_balances(self, account):
+        """Gets account balances for a given account.
+
+        Args:
+            account (str): Account number of the account you want to get balances for.
+
+        Returns:
+            dict: Dict of the response from the API.
+        """
+        response = self.session.get(urls.account_balances(account))
+        return response.json()
 
     def get_positions(self, account):
         """Gets currently held positions for a given account.
@@ -222,36 +320,53 @@ class FTAccountData:
             account (str): Account number of the account you want to get positions for.
 
         Returns:
-            self.securities_held {dict}:
-            Dict of held positions with the pos. ticker as the key.
+            dict: Dict of the response from the API.
         """
+
+        response = self.session.get(urls.account_positions(account))
+        return response.json()
+
+    def get_account_history(self, account):
+        """Gets account history for a given account.
+
+        Args:
+            account (str): Account number of the account you want to get history for.
+
+        Returns:
+            dict: Dict of the response from the API.
+        """
+        response = self.session.get(urls.account_history(account))
+        return response.json()
+
+    def get_orders(self, account):
+        """
+        Retrieves existing order data for a given account.
+
+        Args:
+            ft_session (FTSession): The session object used for making HTTP requests to Firstrade.
+            account (str): Account number of the account to retrieve orders for.
+
+        Returns:
+            list: A list of dictionaries, each containing details about an order.
+        """
+
+        response = self.session.get(url=urls.order_list(account))
+        return response.json()
+
+    def cancel_order(self, order_id):
+        """
+        Cancels an existing order.
+
+        Args:
+            order_id (str): The order ID to cancel.
+
+        Returns:
+            dict: A dictionary containing the response data.
+        """
+
         data = {
-            "page": "pos",
-            "accountId": str(account),
+            "order_id": order_id,
         }
-        position_soup = BeautifulSoup(
-            self.session.post(
-                url=urls.get_xml(),
-                headers=urls.session_headers(),
-                data=data,
-                cookies=self.session.cookies,
-            ).text,
-            "xml",
-        )
 
-        tickers = position_soup.find_all("symbol")
-        quantity = position_soup.find_all("quantity")
-        price = position_soup.find_all("price")
-        change = position_soup.find_all("change")
-        change_percent = position_soup.find_all("changepercent")
-        vol = position_soup.find_all("vol")
-        for i, ticker in enumerate(tickers):
-            ticker = ticker.text
-            self.securities_held[ticker] = {
-                "quantity": quantity[i].text,
-                "price": price[i].text,
-                "change": change[i].text,
-                "change_percent": change_percent[i].text,
-                "vol": vol[i].text,
-            }
-        return self.securities_held
+        response = self.session.post(url=urls.cancel_order(), data=data)
+        return response.json()

firstrade/exceptions.py

@@ -0,0 +1,44 @@
+class QuoteError(Exception):
+    """Base class for exceptions in the Quote module."""
+    pass
+
+class QuoteRequestError(QuoteError):
+    """Exception raised for errors in the HTTP request during a Quote."""
+    def __init__(self, status_code, message="Error in HTTP request"):
+        self.status_code = status_code
+        self.message = f"{message}. HTTP status code: {status_code}"
+        super().__init__(self.message)
+
+class QuoteResponseError(QuoteError):
+    """Exception raised for errors in the API response."""
+    def __init__(self, symbol, error_message):
+        self.symbol = symbol
+        self.message = f"Failed to get data for {symbol}. API returned the following error: {error_message}"
+        super().__init__(self.message)
+        
+class LoginError(Exception):
+    """Exception raised for errors in the login process."""
+    pass
+
+class LoginRequestError(LoginError):
+    """Exception raised for errors in the HTTP request during login."""
+    def __init__(self, status_code, message="Error in HTTP request during login"):
+        self.status_code = status_code
+        self.message = f"{message}. HTTP status code: {status_code}"
+        super().__init__(self.message)
+
+class LoginResponseError(LoginError):
+    """Exception raised for errors in the API response during login."""
+    def __init__(self, error_message):
+        self.message = f"Failed to login. API returned the following error: {error_message}"
+        super().__init__(self.message)
+        
+class AccountError(Exception):
+    """Base class for exceptions in the Account module."""
+    pass
+
+class AccountResponseError(AccountError):
+    """Exception raised for errors in the API response when getting account data."""
+    def __init__(self, error_message):
+        self.message = f"Failed to get account data. API returned the following error: {error_message}"
+        super().__init__(self.message)