cookie/cookie_extractor.py

"""Cookie提取器 - 使用DrissionPage库"""
import tkinter as tk
from tkinter import ttk, messagebox
from browser_manager import BrowserManager
from config_manager import CookieDataManager
from logger import get_logger

logger = get_logger(__name__)


class CookieExtractorApp:
    """Cookie提取器应用"""
    
    def __init__(self, root):
        """
        初始化应用
        
        Args:
            root: 根窗口
        """
        self.root = root
        self.root.title("Cookie提取器")
        self.root.geometry("600x250")
        self.root.resizable(True, True)
        
        self.browser_manager = BrowserManager()
        self.cookie_data_manager = CookieDataManager()
        
        self._create_widgets()
        
    def _create_widgets(self):
        """创建界面控件"""
        main_frame = ttk.Frame(self.root, padding="20")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        url_frame = ttk.LabelFrame(main_frame, text="网址输入", padding="10")
        url_frame.pack(fill=tk.X, pady=10)
        
        ttk.Label(url_frame, text="网址:").pack(side=tk.LEFT, padx=5)
        self.url_var = tk.StringVar(value="http://192.168.8.147:5666/login")
        self.url_entry = ttk.Entry(url_frame, textvariable=self.url_var, width=50)
        self.url_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=10)
        
        self.open_button = ttk.Button(button_frame, text="打开", command=self.open_website, width=15)
        self.open_button.pack(side=tk.LEFT, padx=5)
        
        self.save_button = ttk.Button(button_frame, text="保存", command=self.save_cookies, width=15, state=tk.DISABLED)
        self.save_button.pack(side=tk.LEFT, padx=5)
        
        self.import_button = ttk.Button(button_frame, text="导入", command=self.import_cookies, width=15)
        self.import_button.pack(side=tk.LEFT, padx=5)
        
        self.close_button = ttk.Button(button_frame, text="关闭", command=self.close_app, width=15)
        self.close_button.pack(side=tk.RIGHT, padx=5)
        
        self.status_var = tk.StringVar(value="就绪")
        status_frame = ttk.Frame(main_frame)
        status_frame.pack(fill=tk.X, pady=10)
        
        ttk.Label(status_frame, text="状态:").pack(side=tk.LEFT, padx=5)
        ttk.Label(status_frame, textvariable=self.status_var, foreground="blue").pack(side=tk.LEFT, padx=5)
        
    def open_website(self):
        """打开指定网站"""
        url = self.url_var.get().strip()
        
        if not url:
            messagebox.showerror("错误", "请输入网址")
            return
        
        if not (url.startswith("http://") or url.startswith("https://")):
            url = "https://" + url
            self.url_var.set(url)
        
        try:
            self._cleanup_browser()
            
            self.status_var.set("正在打开浏览器...")
            self.root.update()
            
            self.browser_manager.init_browser()
            
            self.status_var.set(f"正在访问: {url}")
            self.root.update()
            
            self.browser_manager.page.get(url)
            self.browser_manager.wait_page_loaded()
            
            self.status_var.set(f"已打开: {url}")
            self.save_button.config(state=tk.NORMAL)
            messagebox.showinfo("成功", f"已成功打开网址: {url}\n\n请在浏览器中完成登录操作后，点击'保存'按钮")
            
        except Exception as e:
            logger.error(f"打开网站失败: {e}")
            messagebox.showerror("错误", f"打开网站失败: {str(e)}")
            self.status_var.set("就绪")
            self._cleanup_browser()
    
    def save_cookies(self):
        """保存当前的cookie信息，并强制关闭浏览器"""
        if not self.browser_manager.page:
            messagebox.showerror("错误", "请先打开网站")
            return
        
        try:
            self.status_var.set("正在提取cookies...")
            self.root.update()
            
            cookies = self.browser_manager.get_cookies()
            logger.info(f"成功提取 {len(cookies)} 个 cookies")
            
            self.cookie_data_manager.save_simple_cookies(cookies, self.url_var.get())
            
            self.status_var.set("正在关闭浏览器...")
            self.root.update()
            
            self._cleanup_browser()
            
            self.status_var.set(f"Cookies已保存到 cookie_data.json")
            messagebox.showinfo('成功', 'Cookies保存成功，浏览器已关闭')
            
        except Exception as e:
            logger.error(f"保存cookies失败: {e}")
            messagebox.showerror("错误", f"保存cookies失败: {str(e)}")
            self.status_var.set("就绪")
            self._cleanup_browser()
    
    def import_cookies(self):
        """打开浏览器，导入cookie信息，然后打开指定网站"""
        import time
        try:
            data = self.cookie_data_manager.load_simple_cookies()
            
            url = data.get("url", "")
            cookies = data.get("cookies", [])
            
            if not url:
                url = self.url_var.get().strip()
            
            if not url:
                messagebox.showerror("错误", "没有有效的网址")
                return
            
            if not (url.startswith("http://") or url.startswith("https://")):
                url = "https://" + url
            
            self._cleanup_browser()
            
            self.status_var.set("正在打开浏览器...")
            self.root.update()
            
            self.browser_manager.init_browser()
            
            self.status_var.set("等待浏览器就绪...")
            self.root.update()
            self.browser_manager.page.get("about:blank")
            time.sleep(1)
            
            from urllib.parse import urlparse
            parsed = urlparse(url)
            domain = parsed.netloc
            
            if cookies:
                self.status_var.set("正在导入cookies...")
                self.root.update()
                success_count = self.browser_manager.import_cookies(cookies, domain)
                logger.info(f"成功导入 {success_count} 个 cookies")
            
            self.status_var.set(f"正在访问: {url}")
            self.root.update()
            self.browser_manager.page.get(url)
            self.browser_manager.wait_page_loaded()
            
            self.status_var.set("刷新页面验证登录状态...")
            self.root.update()
            for i in range(3):
                self.browser_manager.page.refresh()
                self.browser_manager.wait_page_loaded()
                time.sleep(1)
                self.status_var.set(f"刷新第 {i+1} 次...")
                self.root.update()
            
            self.url_var.set(url)
            self.status_var.set(f"已导入并打开: {url}")
            self.save_button.config(state=tk.NORMAL)
            messagebox.showinfo("成功", f"已成功导入 {success_count} 个cookie并打开网址: {url}")
            
        except Exception as e:
            logger.error(f"导入cookies失败: {e}")
            messagebox.showerror("错误", f"导入cookies失败: {str(e)}")
            self.status_var.set("就绪")
            self._cleanup_browser()
    
    def close_app(self):
        """关闭应用"""
        self._cleanup_browser()
        self.root.destroy()
    
    def _cleanup_browser(self):
        """清理浏览器资源"""
        self.browser_manager.close()
        self.save_button.config(state=tk.DISABLED)
        self.status_var.set("就绪")


def main():
    """主函数"""
    root = tk.Tk()
    app = CookieExtractorApp(root)
    
    def on_closing():
        app._cleanup_browser()
        root.destroy()
    
    root.protocol("WM_DELETE_WINDOW", on_closing)
    root.mainloop()


if __name__ == "__main__":
    main()