GreyDGL · Af7eR9l0W · Jun 2, 2023 · Jun 2, 2023 · Jun 2, 2023
diff --git a/extract_cookie.py b/extract_cookie.py
@@ -1,21 +1,79 @@
+import os
 import platform
-from pycookiecheat import chrome_cookies
-import requests
-import pprint
+from pathlib import Path
 
-url = "https://chat.openai.com/public-api/conversation_limit"
+from rich.console import Console
+import browser_cookie3 # solves the utf encoding issue on linux and windows
+COOKIES_TO_EXCLUDE = ["__cf_bm", "_puid","_ga","_ga_9YTZJE58M9"]
+# redundant but kept just in case
+URL = "https://chat.openai.com/public-api/conversation_limit"
 
-# Determine the operating system
-os_name = platform.system()
 
-if os_name == "Darwin":  # macOS
-    cookie_file = "~/Library/Application Support/Google/Chrome/Profile 2/Cookies"
-elif os_name == "Linux":
-    cookie_file = "~/.config/google-chrome/Profile 2/Cookies"
-else:
-    raise Exception("Unsupported operating system: " + os_name)
+def get_cookie_file(os_name: str) -> Path:
+    home = str(Path.home())
+    if os_name == "Darwin":  
+        return Path(home, "Library/Application Support/Google/Chrome/Profile 2/Cookies")
+    elif os_name == "Linux":
+        return Path(home, ".config/google-chrome/Default/Cookies")
+    elif os_name == "Windows":
+        return Path(home, "AppData/Local/Google/Chrome/User Data/Default/Network/Cookies")
+    else:
+        raise Exception(f"Unsupported operating system: {os_name}")
 
-cookies = chrome_cookies(url, cookie_file=cookie_file)
 
-cookies_string = "; ".join(f"{k}={v}" for k, v in cookies.items())
-print(cookies_string)
+def filter_cookies(cookies):
+    cookies_dict = {cookie.name: cookie.value for cookie in cookies}
+    return {k: v for k, v in cookies_dict.items() if k not in COOKIES_TO_EXCLUDE}
+
+
+def format_cookies(cookies):
+    return "; ".join(f"{k}={v}" for k, v in cookies.items())
+
+
+def get_console_command(os_name, cookies_string):
+    if os_name == "Windows":
+        # Setting env in Windows is different depending on if you're using cmd prompt or powershell
+        if os.environ.get('PROMPT'):
+            return "set CHATGPT_COOKIE='" + cookies_string + "'"
+        else:
+            return "$env:CHATGPT_COOKIE='" + cookies_string + "'"
+    else:
+        return "export CHATGPT_COOKIE='" + cookies_string + "'"
+
+
+def main():
+    console = Console()
+    os_name = platform.system()
+    cookie_file = os.getenv("BROWSER_COOKIE_DB") or get_cookie_file(os_name)
+
+    if not os.path.isfile(cookie_file):
+        console.print(
+            "Please run this script on the same machine with Chrome installed and/or "
+            "set BROWSER_COOKIE_DB to point to your cookies db",
+            style="bold yellow",
+        )
+
+        if os_name == "Windows":
+            if os.environ.get('PROMPT'):
+                console.print(f"e.g. set BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+            else:
+                console.print(f"e.g. $env:BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+        else:
+            console.print(f"e.g. export BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+        return
+
+    # This started as a Windows only line, but I've found it works more consistently 
+    # Chrome on windows doesn't store cookie urls with domain wildcards, so all cookie entries are openai.com
+    # adding *.openai.com worked as well
+    cookies = browser_cookie3.chrome(domain_name="openai.com", cookie_file=str(cookie_file))
+    filtered_cookies = filter_cookies(cookies)
+    cookies_string = format_cookies(filtered_cookies)
+    console_command = get_console_command(os_name, cookies_string)
+
+    console.print("Run the following command to set the cookie:\n")
+    console.print(console_command)
+
+
+if __name__ == "__main__":
+    main()
+
diff --git a/pentestgpt/extract_cookie.py b/pentestgpt/extract_cookie.py
@@ -1,43 +1,79 @@
-import platform
-from pycookiecheat import chrome_cookies
-import requests
-from rich.console import Console
-import pprint
 import os
+import platform
 from pathlib import Path
 
+from rich.console import Console
+import browser_cookie3 # solves the utf encoding issue on linux and windows
+COOKIES_TO_EXCLUDE = ["__cf_bm", "_puid","_ga","_ga_9YTZJE58M9"]
+# redundant but kept just in case
+URL = "https://chat.openai.com/public-api/conversation_limit"
+
+
+def get_cookie_file(os_name: str) -> Path:
+    home = str(Path.home())
+    if os_name == "Darwin":  
+        return Path(home, "Library/Application Support/Google/Chrome/Profile 2/Cookies")
+    elif os_name == "Linux":
+        return Path(home, ".config/google-chrome/Default/Cookies")
+    elif os_name == "Windows":
+        return Path(home, "AppData/Local/Google/Chrome/User Data/Default/Network/Cookies")
+    else:
+        raise Exception(f"Unsupported operating system: {os_name}")
+
+
+def filter_cookies(cookies):
+    cookies_dict = {cookie.name: cookie.value for cookie in cookies}
+    return {k: v for k, v in cookies_dict.items() if k not in COOKIES_TO_EXCLUDE}
+
+
+def format_cookies(cookies):
+    return "; ".join(f"{k}={v}" for k, v in cookies.items())
+
+
+def get_console_command(os_name, cookies_string):
+    if os_name == "Windows":
+        # Setting env in Windows is different depending on if you're using cmd prompt or powershell
+        if os.environ.get('PROMPT'):
+            return "set CHATGPT_COOKIE='" + cookies_string + "'"
+        else:
+            return "$env:CHATGPT_COOKIE='" + cookies_string + "'"
+    else:
+        return "export CHATGPT_COOKIE='" + cookies_string + "'"
+
 
 def main():
     console = Console()
-    url = "https://chat.openai.com/public-api/conversation_limit"
-
-    # Determine the operating system
     os_name = platform.system()
-    cookie_file = os.getenv("BROWSER_COOKIE_DB")
-    if not cookie_file:
-        home = str(Path.home())
-        if os_name == "Darwin":  # macOS
-            cookie_file = Path(
-                home, "Library/Application Support/Google/Chrome/Profile 2/Cookies"
-            )
-        elif os_name == "Linux":
-            cookie_file = Path(home, ".config/google-chrome/Default/Cookies")
-        else:
-            raise Exception("Unsupported operating system: " + os_name)
+    cookie_file = os.getenv("BROWSER_COOKIE_DB") or get_cookie_file(os_name)
 
-    if os.path.isfile(cookie_file):
-        cookies = chrome_cookies(url, cookie_file=cookie_file)
-        cookies_string = "; ".join(f"{k}={v}" for k, v in cookies.items())
-        console.print("Run the following command to set the cookie:\n")
-        console.print("export CHATGPT_COOKIE='" + cookies_string + "'")
-    else:
+    if not os.path.isfile(cookie_file):
         console.print(
-            "Please run this script on the same machine with Chrome installed and/or"
-            "set BROWSER_COOKIE_DB to point to your cookies db "
-            "(e.g. 'export BROWSER_COOKIE_DB=~/.config/google-chrome/Profile 2/Cookies').",
+            "Please run this script on the same machine with Chrome installed and/or "
+            "set BROWSER_COOKIE_DB to point to your cookies db",
             style="bold yellow",
         )
 
+        if os_name == "Windows":
+            if os.environ.get('PROMPT'):
+                console.print(f"e.g. set BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+            else:
+                console.print(f"e.g. $env:BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+        else:
+            console.print(f"e.g. export BROWSER_COOKIE_DB='{cookie_file}'", style="bold yellow")
+        return
+
+    # This started as a Windows only line, but I've found it works more consistently 
+    # Chrome on windows doesn't store cookie urls with domain wildcards, so all cookie entries are openai.com
+    # adding *.openai.com worked as well
+    cookies = browser_cookie3.chrome(domain_name="openai.com", cookie_file=str(cookie_file))
+    filtered_cookies = filter_cookies(cookies)
+    cookies_string = format_cookies(filtered_cookies)
+    console_command = get_console_command(os_name, cookies_string)
+
+    console.print("Run the following command to set the cookie:\n")
+    console.print(console_command)
+
 
 if __name__ == "__main__":
     main()
+
diff --git a/requirements.txt b/requirements.txt
@@ -14,3 +14,4 @@ openai
 langchain
 tiktoken
 pycookiecheat
+browser_cookie3