From 65c6d629e8101d3262ba1c38de7316bf42ac8476 Mon Sep 17 00:00:00 2001
From: Den <2119348+dzianisv@users.noreply.github.com>
Date: Sat, 11 Jan 2025 15:30:37 +0300
Subject: [PATCH] Add a login mode

---
 README.md      |  3 +--
 poept/poept.py | 63 +++++++++++++++++++++++++-------------------------
 2 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/README.md b/README.md
index b6c4257..4712447 100644
--- a/README.md
+++ b/README.md
@@ -80,8 +80,7 @@ bot.close()
 ## Env variables
 
 - POE_HEADLESS=true
-- POE_EMAIL=<your email>
-
+- POE_COOKIES=<base64 encoded cookies> otherwise cookies are read from "~/.cache/poept.cookies.json"
 ## Extra
 
 - status of client
diff --git a/poept/poept.py b/poept/poept.py
index 451a2c2..a7261fc 100644
--- a/poept/poept.py
+++ b/poept/poept.py
@@ -17,6 +17,8 @@
 import json
 import base64
 from urllib.parse import urlparse
+import signal
+import sys
 
 class BotNotFound(Exception):
     pass
@@ -32,7 +34,6 @@ class EmptyResponse(Exception):
 default_bot = "Assistant"
 website = "https://poe.com/"
 letters = ["ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz"]
-email_area='input[type="email"]'
 code_area="input[class*=CodeInput"
 go_key='//button[text()="Go"]'
 log_key=f"//button[contains(translate(., '{letters[0]}', '{letters[-1]}' ), 'log')]"
@@ -54,11 +55,11 @@ class PoePT:
 
     def __init__(self,
             cookies: Optional[str] = os.environ.get("POE_COOKIES"),
-            email: Optional[str] = os.environ.get("POE_EMAIL"),
             headless: bool = os.environ.get("POE_HEADLESS", "true") == "true",
+            login_mode: bool = False
         ):
         chrome_options = Options()
-        if headless:
+        if headless or login_mode:
             chrome_options.add_argument("--headless")
             chrome_options.add_argument("--no-sandbox")
             chrome_options.add_argument("--disable-gpu")
@@ -75,6 +76,16 @@ def __init__(self,
         self.driver = webdriver.Chrome(service=chrome_service, options=chrome_options)
         self.stat = "false"
 
+        if login_mode:
+            logger.info("Login mode enabled. Waiting for login completion...")
+            self.driver.get(website)
+            WebDriverWait(self.driver, 120).until(
+                EC.presence_of_element_located((By.CSS_SELECTOR, "textarea[class*='GrowingTextArea_textArea']"))
+            )
+            logger.info("Login detected. Saving cookies...")
+            self.save_cookies_on_exit()
+            return
+
         self.cookies = None
 
         if cookies is not None:
@@ -84,19 +95,11 @@ def __init__(self,
         if self.cookies is None:
             self.cookies = self.read_cookies()
 
-        if self.cookies is None:
-            if email is None:
-                raise AuthenticationFailure("no cookies and email are set")
-            else:
-                self.login(email)
+        for _ in range(3):
+            if self.apply_cookies():
+                break
         else:
-            # i noticed that sometimes cookies are not applied from first run :facepalm:
-            # workaround
-            for _ in range(3):
-                if self.apply_cookies():
-                    break
-            else:
-                raise AuthenticationFailure("cookies wasn't applied after a few attempts")
+            raise AuthenticationFailure("cookies wasn't applied after a few attempts")
 
     def read_cookies(self):
         for cookies_file in [self.cookies_file_path, self.alternative_cookies_file_path]:
@@ -205,22 +208,6 @@ def get_message(self):
     def clearchat(self):
         click(self.driver, By.CSS_SELECTOR, clear_key)
 
-    def login(self, email: str):
-        self.driver.get(website)
-        self.driver.execute_script('window.scrollBy(0, 5);')
-
-        enter(self.driver, By.CSS_SELECTOR, email_area, email)
-        click(self.driver, By.XPATH, go_key)
-
-        code = input("Enter code: ")
-        enter(self.driver, By.CSS_SELECTOR, code_area, code)
-        click(self.driver, By.XPATH, log_key)
-
-        self.cookies = self.driver.get_cookies()
-        with open(self.cookies_file_path, "wb") as f:
-            json.dump(self.cookies, f)
-
-
     def _typein(self, element, text):
         js_code = """
 const element = arguments[0];
@@ -328,4 +315,16 @@ def get_chat_id(self) -> Optional[str]:
             chat_part = parsed_url.path.split('chat/')[1]
             return f'chat/{chat_part}'
         else:
-            return None
\ No newline at end of file
+            return None
+
+    def save_cookies_on_exit(self):
+        logger.info("Saving cookies...")
+        cookies = self.driver.get_cookies()
+        with open(self.cookies_file_path, 'w') as f:
+            json.dump(cookies, f)
+        print(json.dumps(cookies))
+        self.close()
+        sys.exit(0)
+
+if __name__ == "__main__":
+    PoePT(login_mode=True)