cookie/cookie_extractor.py

"""Cookie提取器 - 使用DrissionPage库"""
import tkinter as tk
from tkinter import ttk, messagebox
import json
import os
import logging
from DrissionPage import Chromium, ChromiumOptions

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('cookie_extractor.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class CookieExtractorApp:
    """Cookie提取器应用"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("Cookie提取器")
        self.root.geometry("600x250")
        self.root.resizable(True, True)
        
        self.browser = None
        self.tab = None
        
        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()
            
            co = ChromiumOptions()
            co.set_browser_path(r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe")
            co.set_argument('--no-sandbox')
            co.set_argument('--disable-dev-shm-usage')
            
            self.browser = Chromium(co)
            self.tab = self.browser.new_tab()
            
            self.status_var.set(f"正在访问: {url}")
            self.root.update()
            
            self.tab.get(url)
            
            self.status_var.set("等待页面加载...")
            self.root.update()
            
            self.tab.wait.load_start()
            self.tab.wait.doc_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 or not self.tab:
            messagebox.showerror("错误", "请先打开网站")
            return
        
        try:
            self.status_var.set("正在提取cookies...")
            self.root.update()
            
            cookies = self.browser.cookies()
            logger.info(f"成功提取 {len(cookies)} 个 cookies")
            
            data = {
                "cookies": cookies,
                "url": self.url_var.get()
            }
            
            output_file = "cookie_data.json"
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            self.status_var.set("正在关闭浏览器...")
            self.root.update()
            
            self._cleanup_browser()
            
            self.status_var.set(f"数据已保存到: {output_file}")
            messagebox.showinfo("成功", f"已成功保存 {len(cookies)} 个cookie到: {output_file}\n\n浏览器已关闭")
            
        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信息，然后打开指定网站"""
        try:
            output_file = "cookie_data.json"
            if not os.path.exists(output_file):
                messagebox.showerror("错误", f"文件不存在: {output_file}")
                return
            
            self.status_var.set("正在读取cookies...")
            self.root.update()
            
            with open(output_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            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()
            
            co = ChromiumOptions()
            co.set_browser_path(r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe")
            co.set_argument('--no-sandbox')
            co.set_argument('--disable-dev-shm-usage')
            
            self.browser = Chromium(co)
            self.tab = self.browser.new_tab()
            
            self.status_var.set("正在准备环境...")
            self.root.update()
            self.tab.get("about:blank")
            
            if cookies:
                self.status_var.set("正在导入cookies...")
                self.root.update()
                
                success_count = 0
                for cookie in cookies:
                    try:
                        if not cookie.get('name') or not cookie.get('value'):
                            continue
                        
                        cookie.setdefault('path', '/')
                        cookie.setdefault('httpOnly', False)
                        cookie.setdefault('secure', False)
                        
                        self.browser.set.cookies([cookie])
                        success_count += 1
                    except Exception as e:
                        logger.error(f"导入cookie失败: {cookie.get('name')} - {e}")
                
                logger.info(f"成功导入 {success_count} 个 cookies")
            
            self.status_var.set(f"正在访问: {url}")
            self.root.update()
            self.tab.get(url)
            
            self.status_var.set("等待页面加载...")
            self.root.update()
            self.tab.wait.load_start()
            self.tab.wait.doc_loaded()
            
            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):
        if self.browser:
            try:
                self.browser.quit()
            except Exception as e:
                logger.error(f"关闭浏览器失败: {e}")
            finally:
                self.browser = None
                self.tab = None
                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()