selenium_download.py

import time
import json
import os
import re
import sys
import requests
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.firefox.options import Options as FirefoxOptions
from selenium.common.exceptions import (
    WebDriverException, TimeoutException
)
from urllib.parse import urlparse, unquote

class FirefoxAutoBrowser:
    def __init__(self, command_executor='http://127.0.0.1:4444',
                 firefox_binary_path="C:\\Program Files\\Mozilla Firefox\\firefox.exe",
                 page_load_timeout=30, implicitly_wait=10):
        self.driver = None
        self.command_executor = command_executor
        self.firefox_binary_path = firefox_binary_path
        self.page_load_timeout = page_load_timeout
        self.implicitly_wait = implicitly_wait
        self.init_result = self._init_browser()
        self._download_session = requests.Session()
        self._default_img_suffix = (".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".pdf", ".zip", ".gz",".7z", ".doc", ".docx", ".xlsx", ".ppt")

    def _sync_browser_cookies(self):
        """同步浏览器Cookie到下载session"""
        if not self.driver:
            return
        try:
            cookies = self.driver.get_cookies()
            for ck in cookies:
                self._download_session.cookies.set(
                    ck["name"], ck["value"], domain=ck.get("domain"), path=ck.get("path")
                )
        except Exception:
            pass

    def _get_header(self, referer=None):
        """构造请求头"""
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98",
            "Accept": "*/*",
            "Connection": "keep-alive"
        }
        if referer:
            headers["Referer"] = referer
        return headers

    def _sanitize_filename(self, filename):
        """清理文件名非法字符"""
        if not filename:
            return "unnamed"
        return re.sub(r'[\\/:*?"<>|\r\n\t]+', "_", filename).strip(" .")

    def _get_unique_path(self, save_dir, filename):
        """文件存在自动重命名"""
        base, ext = os.path.splitext(filename)
        full_path = os.path.join(save_dir, filename)
        counter = 1
        while os.path.exists(full_path):
            full_path = os.path.join(save_dir, f"{base}_{counter}{ext}")
            counter += 1
        return full_path

    def download_single_file(self, url, save_dir="./download", referer=None, timeout=30, resume=True):
        """
        单文件/图片下载
        :param url: 资源链接
        :param save_dir: 保存目录
        :param referer: 来源页面地址
        :param timeout: 超时
        :param resume: 断点续传
        :return: (bool, 保存路径/错误信息)
        """
        try:
            # 同步Cookie
            self._sync_browser_cookies()
            # 创建目录
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            headers = self._get_header(referer)
            parsed = urlparse(unquote(url))
            raw_name = os.path.basename(parsed.path) or "file"
            safe_name = self._sanitize_filename(raw_name)
            save_path = os.path.join(save_dir, safe_name)

            # 断点续传
            exist_size = 0
            if resume and os.path.exists(save_path):
                exist_size = os.path.getsize(save_path)
                headers["Range"] = f"bytes={exist_size}-"

            resp = self._download_session.get(url, headers=headers, timeout=timeout, stream=True)
            resp.raise_for_status()
            mode = "ab" if resp.status_code == 206 else "wb"

            with open(save_path, mode) as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            return True, save_path
        except Exception as e:
            return False, f"下载失败:{str(url)} | 错误:{str(e)}"

    def batch_download(self, url_list, save_dir="./batch_download", referer=None, timeout=30, filter_img=False):
        """
        批量下载，filter_img=True只下载图片后缀资源
        :param url_list: 链接列表
        :param save_dir: 保存目录
        :param referer: 全局来源页
        :param timeout: 超时
        :param filter_img: 是否仅过滤图片
        :return: (成功列表, 失败列表)
        """
        success = []
        fail = []
        for link in url_list:
            if filter_img and not link.lower().endswith(self._default_img_suffix):
                continue
            ok, path = self.download_single_file(link, save_dir, referer, timeout)
            print("[Down] " + path)
            if ok:
                success.append({"url": link, "path": path})
            else:
                fail.append({"url": link, "msg": path})
        return success, fail

    def _init_browser(self):
        try:
            firefox_options = FirefoxOptions()
            firefox_options.binary_location = self.firefox_binary_path
            custom_user_agent = "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98"
            firefox_options.add_argument("--headless")
            firefox_options.add_argument(f"--user-agent={custom_user_agent}")
            firefox_options.set_preference("dom.webnotifications.enabled", False)
            firefox_options.set_preference("dom.popup_maximum", -1)
            firefox_options.set_preference("browser.popups.showPopupBlocker", False)
            firefox_options.set_preference("dom.disable_open_during_load", False)
            firefox_options.set_preference("browser.link.open_newwindow", 3)
            firefox_options.set_preference("browser.link.open_newwindow.restriction", 0)
            firefox_options.add_argument("--ignore-certificate-errors")
            self.driver = webdriver.Remote(
                command_executor=self.command_executor,
                options=firefox_options
            )
            self.driver.set_page_load_timeout(self.page_load_timeout)
            self.driver.implicitly_wait(self.implicitly_wait)
            success_json = {
                "code": 200,
                "status": "success",
                "message": "Browser initialized successfully",
                "detail": None,
                "data": None
            }
            return json.dumps(success_json, ensure_ascii=False, indent=2)
        except (WebDriverException, Exception) as e:
            fail_json = {
                "code": 500,
                "status": "failed",
                "message": f"Browser initialization failed: {str(e)}",
                "detail": None,
                "data": None
            }
            return json.dumps(fail_json, ensure_ascii=False, indent=2)

    def switch_to_specific_tab(self, target):
        try:
            if not self.driver:
                raise Exception("Browser driver not initialized")
            all_handles = self.driver.window_handles
            target_handle = None
            if isinstance(target, int):
                if 0 <= target < len(all_handles):
                    target_handle = all_handles[target]
                else:
                    raise Exception(f"Tab index {target} invalid, total tabs: {len(all_handles)}")
            elif isinstance(target, str):
                if target in all_handles:
                    target_handle = target
                else:
                    raise Exception(f"Tab handle {target} not exists")
            else:
                raise Exception("target support int(index) or str(handle)")
            self.driver.switch_to.window(target_handle)
            result_json = {
                "code": 200,
                "status": "success",
                "detail": {"target_param": target, "target_handle": target_handle},
                "data": None
            }
            return json.dumps(result_json, ensure_ascii=False, indent=2)
        except (WebDriverException, Exception) as e:
            fail_json = {
                "code": 500,
                "status": "failed",
                "message": str(e),
                "detail": {"input_target": target},
                "data": None
            }
            return json.dumps(fail_json, ensure_ascii=False, indent=2)

    def open_url_in_specific_tab(self, url, target=0, wait_second=0.8):
        try:
            if not self.driver:
                raise Exception("Browser driver not initialized")
            if not url or not isinstance(url, str) or url.strip() == "":
                raise Exception("empty url")
            switch_result = json.loads(self.switch_to_specific_tab(target))
            if switch_result["code"] != 200:
                raise Exception(switch_result["message"])
            self.driver.get(url.strip())
            time.sleep(wait_second)
            result_json = {
                "code": 200,
                "status": "success",
                "message": f"Opened {url}",
                "detail": {"target_tab": target, "url": url, "wait": wait_second},
                "data": None
            }
            return json.dumps(result_json, ensure_ascii=False, indent=2)
        except (TimeoutException, WebDriverException, Exception) as e:
            fail_json = {
                "code": 500,
                "status": "failed",
                "message": str(e),
                "detail": {"target_tab": target, "url": url},
                "data": None
            }
            return json.dumps(fail_json, ensure_ascii=False, indent=2)

    def get_page_html(self, target=0):
        try:
            if not self.driver:
                raise Exception("Browser driver not initialized")
            switch_ret = json.loads(self.switch_to_specific_tab(target))
            if switch_ret["code"] != 200:
                raise Exception(switch_ret["message"])
            html_source = self.driver.page_source
            result_json = {
                "code": 200,
                "status": "success",
                "message": f"Get html source success, length: {len(html_source)}",
                "detail": {"target_tab": target},
                "data": html_source
            }
            return json.dumps(result_json, ensure_ascii=False, indent=2)
        except Exception as e:
            fail_json = {
                "code": 500,
                "status": "failed",
                "message": f"Get html failed: {str(e)}",
                "detail": {"target_tab": target},
                "data": None
            }
            return json.dumps(fail_json, ensure_ascii=False, indent=2)

class PageParser:
    def __init__(self, browser: FirefoxAutoBrowser):
        self.browser = browser

    @staticmethod
    def _distinct_list(raw_list):
        """列表去重，保持原有顺序"""
        seen = set()
        new_list = []
        for item in raw_list:
            if item not in seen and item.strip() != "":
                seen.add(item)
                new_list.append(item)
        return new_list

    def get_page_attrs(self, url, regx, attrs, timeout, type, distinct=True):
        """
        :param url: 待爬页面地址
        :param regx: CSS选择器
        :param attrs: 提取属性，text模式填空字符串
        :param timeout: 页面加载等待时长（秒）
        :param type: attribute / text
        :param distinct: 是否自动去重，默认开启
        :return: 提取文本/属性列表，异常/无匹配返回空列表
        """
        # 1. 参数基础校验
        if not url.strip():
            print(f"[警告] 传入URL为空，直接返回空列表")
            return []
        if not regx.strip():
            print(f"[警告] CSS选择器不能为空，直接返回空列表")
            return []
        if type not in ("attribute", "text"):
            print(f"[警告] 不支持提取类型 {type}，仅支持 attribute / text")
            return []
        if type == "attribute" and not attrs.strip():
            print(f"[警告] 提取属性模式下，attrs参数不可为空")
            return []
        # 2. 打开页面，使用timeout作为页面渲染等待时间
        try:
            open_ret_json = self.browser.open_url_in_specific_tab(url, target=0, wait_second=timeout)
            open_ret = json.loads(open_ret_json)
            if open_ret["code"] != 200:
                print(f"[页面加载失败] 地址:{url} 错误信息:{open_ret['message']}")
                return []
        except Exception as e:
            print(f"[打开页面异常] {url} 异常:{str(e)}")
            return []
        # 3. 获取网页完整源码
        try:
            html_json = self.browser.get_page_html(0)
            html_data = json.loads(html_json)
            html_text = html_data["data"]
            if not html_text or len(html_text.strip()) == 0:
                print(f"[页面源码为空] 地址:{url}")
                return []
        except Exception as e:
            print(f"[获取源码异常] {url} 异常:{str(e)}")
            return []
        # 4. 调用本地解析方法提取数据
        raw_result = self.search_page(html_text, regx, attrs, type)
        if len(raw_result) == 0:
            print(f"[匹配结果为空] 页面{url} 选择器:{regx} 未找到任何元素")
            return []
        # 5. 可选去重
        if distinct:
            final_result = self._distinct_list(raw_result)
            print(f"[提取完成] 原始匹配{len(raw_result)}条，去重后{len(final_result)}条")
        else:
            final_result = raw_result
            print(f"[提取完成] 共匹配{len(raw_result)}条数据")
        return final_result

    def search_page(self, data, regx, attrs, type, strip_text=True):
        """
        本地HTML解析工具，增强容错
        :param data: 网页源码字符串
        :param regx: CSS选择器
        :param attrs: 属性名
        :param type: attribute / text
        :param strip_text: text模式自动去除首尾空白
        :return: 提取列表
        """
        respon_page = []
        if not data or not data.strip():
            return respon_page
        try:
            soup = BeautifulSoup(data, "html.parser")
            ret = soup.select(regx)
            for item in ret:
                if type == "attribute":
                    val = item.attrs.get(attrs, "").strip()
                    respon_page.append(val)
                elif type == "text":
                    if strip_text:
                        val = item.get_text(strip=True)
                    else:
                        val = item.get_text()
                    respon_page.append(val)
        except Exception as e:
            print(f"[解析HTML异常] 选择器{regx} 错误:{str(e)}")
        return respon_page

if __name__ == "__main__":

    referer_url = "https://www.baidu.com"

    browser = None
    try:
        browser = FirefoxAutoBrowser()
        init_res = json.loads(browser.init_result)
        if init_res["code"] != 200:
            exit(1)

        parser = PageParser(browser)

        # 提取链接
        ref_img = parser.get_page_attrs(
            url=sys.argv[1],
            regx="div[id='wrapper'] div[id='content'] ul[class='wookmark-initialised'] li[class='thumbwook'] a[class='rel-link']",
            attrs="href",
            timeout=5,
            type="attribute",
            distinct=True
        )

        # 过滤后缀
        suffix_list = (".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".pdf", ".zip", ".gz",".7z", ".doc", ".docx", ".xlsx", ".ppt")
        img_links = [link for link in ref_img if link.lower().endswith(suffix_list)]

        success_list, fail_list = browser.batch_download(
            url_list=img_links,
            save_dir="./download",
            referer=referer_url,
            timeout=15,
            filter_img=False
        )

        print(f"\n下载完成：成功 {len(success_list)} | 失败 {len(fail_list)}")

    except Exception as err:
        print(f"\n程序全局异常：{str(err)}")
    finally:
        if browser and browser.driver:
            browser.driver.quit()
            print("程序结束")