crossborder/utils/download_utils.py

import re
import threading
import time
import traceback
from pathlib import Path

from faker import Faker
from selenium.webdriver import FirefoxOptions, ActionChains
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait

from utils.log import log

DOWNLOAD_TIMEOUT = 60

download_lock = threading.Lock()

def configure_stealth_options(download_dir):
    """反检测浏览器配置"""
    opts = FirefoxOptions()
    opts.set_preference("dom.webdriver.enabled", False)
    opts.add_argument("--disable-blink-features=AutomationControlled")
    opts.add_argument("--headless")
    opts.add_argument("--window-size=1366,900")  # 固定视口
    opts.set_preference("general.useragent.override", Faker().user_agent())
    opts.set_preference("browser.download.dir",  str(download_dir))
    opts.set_preference("browser.download.folderList", 2)
    opts.set_preference("browser.helperApps.neverAsk.saveToDisk", "application/vnd.ms-excel")
    opts.set_preference("pdfjs.disabled", True)
    opts.set_preference("download.prompt_for_download", False)
    opts.set_preference("download.directory_upgrade", True)
    return opts


def generate_month_sequence(start_year, start_month, end_year=None, skip_january=False):
    """
    动态生成倒序月份序列

    Args:
        start_year (int): 检测到的最新数据年份
        start_month (int): 检测到的最新数据月份
        end_year (int, optional): 终止年份（默认 None）
        skip_january (bool, optional): 是否跳过所有1月数据（默认 False）

    Returns:
        List[Tuple[int, int]]: 月份序列列表，格式为 [(year, month), ...]
    """
    sequence = []
    current_year = start_year
    current_month = start_month

    # 当指定终止年份时
    if end_year:
        while not (current_year == end_year and current_month < 1):
            # 跳过1月判断
            if not (skip_january and current_month == 1):
                sequence.append((current_year, current_month))

            # 跨年处理
            if current_month == 1:
                current_year -= 1
                current_month = 12
            else:
                current_month -= 1

            # 终止条件：到达目标年份的1月
            if current_year < end_year:
                break
    else:
        # 未指定年份时取最近两个月
        if not (skip_january and current_month == 1):
            sequence.append((current_year, current_month))

        prev_year, prev_month = get_previous_month(current_year, current_month)
        if not (skip_january and prev_month == 1):
            sequence.append((prev_year, prev_month))

    return sequence


def get_previous_month(year, month):
    """跨年月份计算"""
    if month == 1:
        return year - 1, 12
    return year, month - 1

def download_excel(driver, url, year, month, title, download_dir):
    """文件下载模块"""

    download_dir.mkdir(parents=True, exist_ok=True)

    driver.execute_script(f"window.open('{url}')")
    driver.switch_to.window(driver.window_handles[-1])

    try:
        download_btn = WebDriverWait(driver, 20).until(
            EC.element_to_be_clickable(
                (By.XPATH,
                 '//a[substring(@href, string-length(@href)-3) = ".xls" or substring(@href, string-length(@href)-4) = ".xlsx"]')
            )
        )
        # download_btn.click()
        driver.execute_script("arguments[0].scrollIntoView({block: 'center'});", download_btn)
        # 使用 ActionChains 安全点击
        ActionChains(driver).move_to_element(download_btn).pause(0.3).click().perform()
        # log.info(f"√ 已点击下载按钮：{download_btn.get_attribute("href")}")

        downloaded_file  = wait_for_download(download_dir)
        final_path = Path(f'{download_dir}/{year}/{month:02d}/{title}{downloaded_file.suffix}')
        if final_path.exists():
            final_path.unlink()
        download_rel_dir = Path(f'{download_dir}/{year}/{month:02d}')
        download_rel_dir.mkdir(parents=True, exist_ok=True)
        downloaded_file.rename(final_path)
        log.info(f"√ 文件已保存至：{final_path}")
    finally:
        driver.close()
        driver.switch_to.window(driver.window_handles[0])


# def batch_download_excel(driver, url, year, month, base_title, download_dir):
#     """批量下载Excel文件模块"""
#     download_dir.mkdir(parents=True, exist_ok=True)
#
#     driver.execute_script(f"window.open('{url}')")
#     driver.switch_to.window(driver.window_handles[-1])
#
#     try:
#         # 获取所有Excel下载按钮
#         download_btns = WebDriverWait(driver, 20).until(
#             EC.presence_of_all_elements_located(
#                 (By.XPATH,
#                  '//a[substring(@href, string-length(@href)-3) = ".xls" or substring(@href, string-length(@href)-4) = ".xlsx"]')
#             )
#         )
#
#         for index, btn in enumerate(download_btns):
#             # 生成唯一标题（可自定义规则）
#             title = f"{base_title}_{index + 1}"
#
#             # 点击下载按钮
#             btn.click()
#
#             # 等待下载完成
#             downloaded_file = wait_for_download(download_dir)
#
#             # 处理文件路径
#             final_dir = download_dir / f'{year}' / f'{month:02d}'
#             final_dir.mkdir(parents=True, exist_ok=True)
#             final_path = final_dir / f'{title}{downloaded_file.suffix}'
#
#             # 重命名文件
#             if final_path.exists():
#                 final_path.unlink()
#             downloaded_file.rename(final_path)
#             log.info(f"√ 文件 {title} 已保存至：{final_path}")
#
#     finally:
#         driver.close()
#         driver.switch_to.window(driver.window_handles[0])

def batch_download_excel(driver, url, base_year, base_month, base_title, download_dir):
    """批量下载Excel文件模块 - 支持根据文件名自动识别并分类到对应月份"""
    download_dir.mkdir(parents=True, exist_ok=True)

    driver.execute_script(f"window.open('{url}')")
    driver.switch_to.window(driver.window_handles[-1])

    try:
        # 获取所有Excel下载按钮
        download_btns = WebDriverWait(driver, 20).until(
            EC.presence_of_all_elements_located(
                (By.XPATH,
                 '//a[substring(@href, string-length(@href)-3) = ".xls" or substring(@href, string-length(@href)-4) = ".xlsx"]')
            )
        )

        for index, btn in enumerate(download_btns):
            raw_title = btn.text.strip() or btn.get_attribute("textContent").strip()
            if not raw_title:
                log.warning("⚠️ 标题为空，跳过处理")
                continue

            # 提取年份
            year_match = re.search(r'(\d{4})\s*年', raw_title)
            file_year = int(year_match.group(1)) if year_match else base_year

            # 提取月份（支持多种形式）
            month_match = re.search(
                r'(?:1[-－]\s*)?(\d{1,2})月|前?(\d{1,2})个?月|第([一二三四])季度',
                raw_title
            )
            file_month = base_month

            if month_match:
                digit_groups = month_match.groups()
                if digit_groups[0]:
                    file_month = int(digit_groups[0])
                elif digit_groups[1]:
                    file_month = int(digit_groups[1])
                elif digit_groups[2]:
                    quarter_map = {"一": 3, "二": 6, "三": 9, "四": 12}
                    file_month = quarter_map.get(digit_groups[2], base_month)

            # 确保月份合法
            if not (1 <= file_month <= 12):
                file_month = base_month

            # 构建目录
            final_dir = download_dir / f'{file_year}' / f'{file_month:02d}'
            final_dir.mkdir(parents=True, exist_ok=True)

            # 下载文件
            btn.click()
            downloaded_file = wait_for_download(download_dir)

            # 构建最终路径（去除已有后缀，防止重复）
            clean_title = Path(raw_title).stem
            final_path = final_dir / f'{clean_title}{downloaded_file.suffix}'

            # 删除已存在文件并重命名
            if final_path.exists():
                final_path.unlink()
            downloaded_file.rename(final_path)

            log.info(f"√ 文件 {clean_title} 已保存至：{final_path}")

    finally:
        driver.close()
        driver.switch_to.window(driver.window_handles[0])




def download_excel2(driver, link, year, month, title, download_dir):
    download_dir = Path(download_dir)
    download_dir.mkdir(parents=True, exist_ok=True)

    try:
        log.info(f"正在点击链接：{title}")
        link.click()
        log.info("等待文件下载完成...")
        downloaded_file = wait_for_download(download_dir)

        if not downloaded_file.suffix:
            downloaded_file = downloaded_file.with_suffix('.xlsx')

        final_dir = download_dir / f'{year}' / f'{month:02d}'
        final_dir.mkdir(parents=True, exist_ok=True)
        final_path = final_dir / f'{title}{downloaded_file.suffix}'

        if final_path.exists():
            final_path.unlink()
        downloaded_file.rename(final_path)
        log.info(f"√ 文件已保存至：{final_path}")

    except TimeoutError as te:
        log.info(f"[错误] 文件下载超时：{te}")
        raise
    except Exception as e:
        log.info(f"[错误] 发生异常：{e}")
        log.info(traceback.format_exc())
        raise


# def wait_for_download(directory):
#     """文件下载监控（只读取文件，忽略文件夹）"""
#     start_time = time.time()
#     while (time.time() - start_time) < DOWNLOAD_TIMEOUT:
#         files = [
#             f for f in directory.glob('*')
#             if f.is_file() and not f.name.endswith(('.part', '.crdownload'))
#         ]
#         if files:
#             # 按照创建时间排序并返回最新文件
#             return max(files, key=lambda x: x.stat().st_ctime)
#         time.sleep(1)
#     raise TimeoutError("文件下载超时")


def wait_for_download(directory):
    """文件下载监控（只读取文件，忽略文件夹）"""
    start_time = time.time()
    while (time.time() - start_time) < DOWNLOAD_TIMEOUT:
        with download_lock:
            files = [
                f for f in directory.glob('*')
                if f.is_file() and not f.name.endswith(('.part', '.crdownload'))
            ]
            if files:
                # 按照创建时间排序并返回最新文件
                return max(files, key=lambda x: x.stat().st_ctime)
        time.sleep(1)
    raise TimeoutError("文件下载超时")