Merge pull request #342 from matthew55/generic-cookie-functions

Add helper functions for handling cookies

Author: NelsonDane

fidelityAPI.py

@@ -27,6 +27,8 @@
     printHoldings,
     stockOrder,
     type_slowly,
+    save_cookies,
+    load_cookies,
 )
 
 
@@ -74,6 +76,7 @@ def fidelity_init(FIDELITY_EXTERNAL=None, DOCKER=False, botObj=None, loop=None):
             if driver is None:
                 raise Exception("Error: Unable to get driver")
             # Log in to Fidelity account
+            load_cookies(driver, filename=f"fidelity{index}.pkl", path="./creds/")
             driver.get(
                 "https://digital.fidelity.com/prgw/digital/login/full-page?AuthRedUrl=digital.fidelity.com/ftgw/digital/portfolio/summary"
             )
@@ -186,6 +189,8 @@ def fidelity_init(FIDELITY_EXTERNAL=None, DOCKER=False, botObj=None, loop=None):
 
                 code_field = driver.find_element(by=By.CSS_SELECTOR, value=code_field)
                 code_field.send_keys(str(sms_code))
+                remember_device_checkbox = "#dom-trust-device-checkbox + label"
+                driver.find_element(By.CSS_SELECTOR, remember_device_checkbox).click()
                 continue_btn_selector = "#dom-otp-code-submit-button"
                 driver.find_element(By.CSS_SELECTOR, continue_btn_selector).click()
             except TimeoutException:
@@ -229,6 +234,7 @@ def fidelity_init(FIDELITY_EXTERNAL=None, DOCKER=False, botObj=None, loop=None):
                 fidelity_obj.set_account_totals(
                     name, acct, account_dict[acct]["balance"]
                 )
+            save_cookies(driver, filename=f"fidelity{index}.pkl", path="./creds/")
             print(f"Logged in to {name}!")
         except Exception as e:
             fidelity_error(driver, e)

helperAPI.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import os
+import pickle
 import subprocess
 import sys
 import textwrap
@@ -760,3 +761,35 @@ def printHoldings(brokerObj: Brokerage, loop=None, mask=True):
             EMBED["fields"].append(field)
     printAndDiscord(EMBED, loop, True)
     print("==============================")
+
+
+def save_cookies(driver, filename, path=None):
+    if path is not None:
+        filename = os.path.join(path, filename)
+    if path is not None and not os.path.exists(path):
+        os.makedirs(path)
+    with open(filename, "wb") as f:
+        pickle.dump(driver.get_cookies(), f)
+
+
+def load_cookies(driver, filename, path=None):
+    if path is not None:
+        filename = os.path.join(path, filename)
+    if not os.path.exists(filename):
+        return False
+    with open(filename, "rb") as f:
+        cookies = pickle.load(f)
+    for cookie in cookies:
+        try:
+            driver.add_cookie(cookie)
+        except Exception:
+            continue
+    return True
+
+
+def clear_cookies(driver, filename, path=None):
+    if path is not None:
+        filename = os.path.join(path, filename)
+    if os.path.exists(filename):
+        os.remove(filename)
+    driver.delete_all_cookies()