diff --git a/src/manager.py b/src/manager.py new file mode 100644 index 0000000..b4b3567 --- /dev/null +++ b/src/manager.py @@ -0,0 +1,309 @@ +""" +Progress and Cookie Management Module + +This module provides functionality for managing application progress, cookies, +and database operations. + +Classes: + ProgressManager: Handles progress tracking and persistence + CookieManager: Manages browser cookie operations + DatabaseManager: Handles database interactions +""" + +import logging +import os +import pickle +import sqlite3 +import sys +import time +from datetime import date + +from selenium.common.exceptions import UnableToSetCookieException +from selenium.webdriver.common.by import By + +LOGGER_NAME = "ChatGPT-API-Leakage" +LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" +logging.basicConfig(level=logging.INFO, format=LOG_FORMAT, datefmt="[%X]") +logger = logging.getLogger(LOGGER_NAME) + + +class ProgressManagerError(Exception): + """Custom exception for ProgressManager class errors""" + + def __init__(self, message): + super().__init__(message) + + +class ProgressManager: + """ + Manages and persists progress information for long-running operations. + + Attributes: + progress_file (Path): Path to the progress file + + Methods: + save: Saves current progress + load: Loads saved progress + """ + + def __init__(self, progress_file=".progress.txt"): + """ + Initialize the ProgressManager with a specified progress file. + + Args: + progress_file (str): The file where progress data is stored. + """ + self.progress_file = progress_file + + def save(self, from_iter: int, total: int): + """ + Saves the current progress to a file. + + Args: + from_iter (int): The current iteration progress. + total (int): The total number of iterations. + """ + with open(self.progress_file, "w", encoding="utf-8") as file: + file.write(f"{from_iter}/{total}/{time.time()}") + + def load(self, total: int) -> int: + """ + Loads the previously saved progress if available and valid. + + Args: + total (int): The total number of iterations for the current process. + + Returns: + int: The iteration number to continue from. + """ + if not os.path.exists(self.progress_file): + return 0 + + with open(self.progress_file, "r", encoding="utf-8") as file: + last_, totl_, tmst_ = file.read().strip().split("/") + last, totl = int(last_), int(totl_) + + if time.time() - float(tmst_) < 3600 and totl == total: + action = input(f"🔍 Progress found, do you want to continue from the last progress ({last}/{totl})? [yes] | no: ").lower() + if action in {"yes", "y", ""}: + return last + + return 0 + + +class CookieManager: + """ + Manages browser cookie operations. + + Methods: + save: Saves cookies to a file + load: Loads cookies from a file + verify_user_login: Checks if the user is currently logged in + """ + + def __init__(self, driver): + """ + Initialize the CookieManager with a Selenium WebDriver instance. + + Args: + driver (WebDriver): The Selenium WebDriver for cookie operations. + """ + self.driver = driver + + def save(self): + """ + Save cookies from the current browser session to a file. + """ + cookies = self.driver.get_cookies() + with open("cookies.pkl", "wb") as file: + pickle.dump(cookies, file) + logger.info("🍪 Cookies saved") + + def load(self): + """ + Load cookies from a file and attempt to add them to the current browser session. + """ + try: + with open("cookies.pkl", "rb") as file: + cookies = pickle.load(file) + for cookie in cookies: + try: + self.driver.add_cookie(cookie) + except UnableToSetCookieException: + logger.debug("🟡 Warning, unable to set a cookie %s", cookie) + except (EOFError, pickle.UnpicklingError): + if os.path.exists("cookies.pkl"): + os.remove("cookies.pkl") + logger.error("🔴 Error, unable to load cookies, invalid cookies has been removed, please restart.") + + def verify_user_login(self): + """ + Test if the user is really logged in by navigating to GitHub and checking login status. + """ + logger.info("🤗 Redirecting ...") + self.driver.get("https://github.com/") + + if self.driver.find_elements(by=By.XPATH, value="//*[contains(text(), 'Sign in')]"): + if os.path.exists("cookies.pkl"): + os.remove("cookies.pkl") + logger.error("🔴 Error, you are not logged in, please restart and try again.") + sys.exit(1) + return True + + +class DatabaseManager: + """ + This class is used to manage the database, including creating tables and handling data interactions. + """ + + def __init__(self, db_filename: str): + """ + Initialize the DatabaseManager with the specified database filename. + + Args: + db_filename (str): Path to the SQLite database file. + """ + self.db_filename = db_filename + self.con = None + self.cur = None + + def __enter__(self): + """ + Enter the runtime context related to this object, initializing the database if needed. + """ + if not os.path.exists(self.db_filename): + logging.info("Creating database github.db") + + self.con = sqlite3.connect(self.db_filename) + self.cur = self.con.cursor() + + self.cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='APIKeys'") + if self.cur.fetchone() is None: + logging.info("Creating table APIKeys") + self.cur.execute("CREATE TABLE APIKeys(apiKey, status, lastChecked)") + + self.cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='URLs'") + if self.cur.fetchone() is None: + logging.info("Creating table URLs") + self.cur.execute("CREATE TABLE URLs(url, key)") + + return self + + def __exit__(self, exc_type, exc_value, traceback): + """ + Exit the runtime context and close the database connection. + """ + if self.con: + self.con.close() + + def all_iq_keys(self) -> list: + """ + Get all keys with the status 'insufficient_quota'. + + Returns: + list: A list of tuples containing API keys. + """ + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("SELECT apiKey FROM APIKeys WHERE status='insufficient_quota'") + return self.cur.fetchall() + + def all_keys(self) -> list: + """ + Get all keys with the status 'yes'. + + Returns: + list: A list of tuples containing API keys. + """ + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("SELECT apiKey FROM APIKeys WHERE status='yes'") + return self.cur.fetchall() + + def deduplicate(self) -> None: + """ + Deduplicate the 'APIKeys' table by retaining only the latest record for each key. + """ + if self.con is None: + raise ValueError("Connection is not initialized") + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("CREATE TABLE temp_table as SELECT apiKey, status, MAX(lastChecked) as lastChecked FROM APIKeys GROUP BY apiKey;") + self.cur.execute("DROP TABLE APIKeys;") + self.cur.execute("ALTER TABLE temp_table RENAME TO APIKeys;") + self.con.commit() + + def delete(self, api_key: str) -> None: + """ + Delete a specific API key from the database. + + Args: + api_key (str): The unique API key to remove. + """ + if self.con is None: + raise ValueError("Connection is not initialized") + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("DELETE FROM APIKeys WHERE apiKey=?", (api_key,)) + self.con.commit() + + def insert(self, api_key: str, status: str): + """ + Insert a new API key and status into the database. + + Args: + api_key (str): The API key to insert. + status (str): The status of the API key. + """ + if self.con is None: + raise ValueError("Connection is not initialized") + if self.cur is None: + raise ValueError("Cursor is not initialized") + today = date.today() + self.cur.execute("INSERT INTO APIKeys(apiKey, status, lastChecked) VALUES(?, ?, ?)", (api_key, status, today)) + self.con.commit() + + def key_exists(self, api_key: str) -> bool: + """ + Check if a given API key exists in the database. + + Args: + api_key (str): The API key to search for. + + Returns: + bool: True if the API key exists, False otherwise. + """ + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("SELECT apiKey FROM APIKeys WHERE apiKey=?", (api_key,)) + return self.cur.fetchone() is not None + + def insert_url(self, url: str) -> None: + """ + Insert a new URL into the 'URLs' table. + + Args: + url (str): The URL to add. + """ + if self.con is None: + raise ValueError("Connection is not initialized") + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("INSERT INTO URLs(url, key) VALUES(?, ?)", (url, 1)) + self.con.commit() + + def get_url(self, url: str) -> str | None: + """ + Retrieve the 'key' associated with the given URL. + + Args: + url (str): The URL to look up. + + Returns: + str | None: The key if it exists, None if not. + """ + if self.cur is None: + raise ValueError("Cursor is not initialized") + self.cur.execute("SELECT key FROM URLs WHERE url=?", (url,)) + fetch = self.cur.fetchone() + return fetch[0] if fetch else None