fixing email check

substack/api.py

@@ -8,20 +8,13 @@ import base64
 import json
 import logging
 import os
-from datetime import datetime
 from urllib.parse import urljoin
-from pyvirtualdisplay import Display
 import requests
-from selenium import webdriver
-from selenium.webdriver.common.keys import Keys
-from selenium.webdriver.common.by import By
-from selenium.webdriver.chrome.options import Options
-from selenium.webdriver.support.ui import WebDriverWait
 import pickle
 import time
 from substack.exceptions import SubstackAPIException, SubstackRequestException
-from selenium.webdriver.support import expected_conditions as EC
-import datetime
+from checkemail import get_verification_link
+from datetime import datetime
 logger = logging.getLogger(__name__)
 
 __all__ = ["Api"]
@@ -69,25 +62,31 @@ class Api:
             logging.getLogger().setLevel(logging.DEBUG)
 
         self._session = requests.Session()
-
+        
         # Load cookies from file if provided
-        # Helps with Captcha errors by reusing cookies from "local" auth, then switching to running code in the cloud
-
         if cookies_path is not None:
-            with open(cookies_path) as f:
-                cookies = json.load(f)
-            self._session.cookies.update(cookies)
-
-        elif email is not None and password is not None:
-            self.send_magic_link(email)
-            magic_link = input("Enter magic link: ")
-            self.login_v2(email, password, magic_link)
-            self.export_cookies(cookies_path)
             
+            if os.path.exists(cookies_path):
+                with open(cookies_path) as f:
+                    cookies = json.load(f)
+                self._session.cookies.update(cookies)        
+
+            if not os.path.exists(cookies_path) or self.are_cookies_expired(cookies):
+                print("Cookies are expired. Sending magic link and waiting for verification.")
+                start_time = time.time()  # Record the time when the magic link is sent
+                self.send_magic_link(email)
+                verification_link = self.wait_for_verification_link(start_time)
+                if verification_link:
+                    self.login_v2(email, password, verification_link)
+                    self.export_cookies(cookies_path)
+
+                else:
+                    raise Exception("Failed to get the verification link.")
+        
+        elif email is not None and password is not None:
+            self.login(email, password)
         else:
-            raise ValueError(
-                "Must provide email and password or cookies_path to authenticate."
-            )
+            raise ValueError("Must provide email and password or cookies_path to authenticate.")
 
         user_publication = None
         # if the user provided a publication url, then use that
@@ -112,6 +111,22 @@ class Api:
         # set the current publication to the users primary publication
         self.change_publication(user_publication)
 
+    def are_cookies_expired(self, cookies):
+        for cookie in cookies:
+            if 'expiry' in cookie and cookie['expiry'] < time.time():
+                return True
+        return False
+        
+    def wait_for_verification_link(self, start_time):
+        sender_email = "no-reply@substack.com"
+        
+        while True:
+            verification_link = get_verification_link(self.email, sender_email, start_time)
+            if verification_link:
+                return verification_link
+            time.sleep(10)  # Wait for X seconds before checking again
+
+
     def send_magic_link(self, email):
         body = {
             "email": email,