eta_data_analysis/services/ruizide/data_processor.go

package main

import (
	"context"
	"eta/eta_data_analysis/utils"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/chromedp/chromedp"
)

const downloadDir = "./downloads"

// 定义选择器
var (
	rzdLoginPath = "https://clients.rystadenergy.com/clients/"

	clientSearchLink               = `a[href="/clients/search/"]`
	supplyRevisionAnalysisSelector = `div.ais-Hits li[contains(., 'Supply Revision Analysis')]`
	oilDemandAnalysisSelector      = `div.ais-Hits li[contains(., 'Oil Demand Analysis')]`
	oilSupplyAnalysisSelector      = `div.ais-Hits li[contains(., 'Oil Supply Analysis')]`
	dateSlicerInputSelector        = `div.visual.customPadding.allow-deferred-rendering.visual-slicer input.date-slicer-input.enable-hover:first-of-type`
	downloadButtonSelector         = `div.btn.btn-link.btn-sm.dashboard-action.dashboard-action--download-data`
)

// 函数用于设置查询时间范围
func setQueryTime(ctx context.Context, year string) error {
	return chromedp.Run(ctx,
		chromedp.WaitVisible(dateSlicerInputSelector, chromedp.ByQuery),
		chromedp.SetValue(dateSlicerInputSelector, year, chromedp.ByQuery),
		chromedp.SendKeys(dateSlicerInputSelector, "\n"), // 回车查询
	)
}

// 函数用于点击下载按钮
func clickDownload(ctx context.Context) error {
	return chromedp.Run(ctx, chromedp.Click(downloadButtonSelector, chromedp.ByQuery))
}

// 处理数据下载的步骤
func downloadData(ctx context.Context) error {
	// Analytics Library
	if err := chromedp.Run(ctx,
		chromedp.WaitVisible(clientSearchLink, chromedp.ByQuery),
		chromedp.Click(clientSearchLink, chromedp.ByQuery),
		chromedp.SetValue(`input[class="ais-SearchBox-input rounded border py-2 px-3 shadow-sm font-size-14 w-100"]`, "oil demand signals weekly report", chromedp.ByQuery),
		chromedp.Click(`div.ais-InfiniteHits li:first-child img[src="/Static/img/icons/xls.png"]`, chromedp.ByQuery),
	); err != nil {
		return err
	}

	// Cube Dashboards: Supply Revision Analysis
	if err := chromedp.Run(ctx,
		chromedp.WaitVisible(`a[href="/clients/subscription/"]`, chromedp.ByQuery),
		chromedp.Click(`a[href="/clients/subscription/"]`, chromedp.ByQuery),
		chromedp.Click(supplyRevisionAnalysisSelector, chromedp.ByQuery),
	); err != nil {
		return err
	}
	if err := setQueryTime(ctx, "2020"); err != nil {
		return err
	}
	if err := clickDownload(ctx); err != nil {
		return err
	}
	if err := waitAndRenameDownloadedFile("Supply_Revision_Analysis_2020.xlsx"); err != nil {
		return err
	}

	// Oil Demand Analysis
	if err := chromedp.Run(ctx,
		chromedp.Click(`a[href="/clients/subscription/"]`, chromedp.ByQuery),
		chromedp.Click(oilDemandAnalysisSelector, chromedp.ByQuery),
	); err != nil {
		return err
	}
	if err := setQueryTime(ctx, "2015"); err != nil {
		return err
	}
	if err := clickDownload(ctx); err != nil {
		return err
	}
	if err := waitAndRenameDownloadedFile("Oil_Demand_Analysis_2015.xlsx"); err != nil {
		return err
	}

	// Oil Supply Analysis
	if err := chromedp.Run(ctx,
		chromedp.Click(`a[href="/clients/subscription/"]`, chromedp.ByQuery),
		chromedp.Click(oilSupplyAnalysisSelector, chromedp.ByQuery),
	); err != nil {
		return err
	}
	if err := setQueryTime(ctx, "2010"); err != nil {
		return err
	}
	if err := clickDownload(ctx); err != nil {
		return err
	}
	if err := waitAndRenameDownloadedFile("Oil_Supply_Analysis_2010.xlsx"); err != nil {
		return err
	}

	return nil
}

// 等待下载文件并重命名
func waitAndRenameDownloadedFile(newFileName string) error {
	// 等待一段时间以确保文件下载完成
	time.Sleep(5 * time.Second) // 可能需要根据实际情况调整

	// 查找下载目录中的文件
	files, err := filepath.Glob(filepath.Join(downloadDir, "*.xlsx"))
	if err != nil {
		return fmt.Errorf("查找文件时出错: %v", err)
	}

	// 重命名最新的文件
	for _, file := range files {
		if err := os.Rename(file, filepath.Join(downloadDir, newFileName)); err != nil {
			return fmt.Errorf("重命名文件时出错: %v", err)
		}
		// 打印重命名后的文件名
		fmt.Printf("文件重命名为: %s\n", newFileName)
		break // 只重命名第一个找到的文件
	}

	return nil
}

func main() {
	// 创建下载目录
	if err := os.MkdirAll(downloadDir, os.ModePerm); err != nil {
		fmt.Printf("创建下载目录时出错: %v\n", err)
		return
	}

	// 创建 chromedp 执行上下文
	options := []chromedp.ExecAllocatorOption{
		chromedp.Flag("headless", false),
		chromedp.Flag("disable-blink-features", "AutomationControlled"),
		chromedp.UserAgent(`Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/98.0.4758.80 Safari/537.36`),
		chromedp.Flag("download.default_directory", downloadDir),
		chromedp.Flag("download.prompt_for_download", false), // 不弹出下载对话框
		chromedp.Flag("safebrowsing.enabled", true),          // 启用安全浏览
	}

	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), options...)
	defer cancel()
	ctx, cancel := chromedp.NewContext(allocCtx)
	defer cancel()

	// 启动 Chrome 实例
	if err := chromedp.Run(ctx); err != nil {
		fmt.Printf("启动 Chrome 实例时出错: %v\n", err)
		return
	}

	// 登录操作
	if err := login(ctx); err != nil {
		fmt.Printf("登录错误: %v\n", err)
		return
	}

	// 下载数据
	if err := downloadData(ctx); err != nil {
		fmt.Printf("数据下载错误: %v\n", err)
		return
	}

	fmt.Println("数据下载完成")
}

func login(ctx context.Context) error {

	return chromedp.Run(ctx,
		chromedp.Navigate(rzdLoginPath),
		chromedp.SetValue(`input[id="Username"]`, utils.LY_USERNAME, chromedp.ByQuery),
		chromedp.SetValue(`input[id="Username"]`, utils.LY_USERNAME, chromedp.ByQuery),
		chromedp.SetValue(`input[id="Password"]`, utils.LY_PASSWORD, chromedp.ByQuery),
		chromedp.WaitEnabled(`//button[text()='Login']`, chromedp.BySearch),
		chromedp.Click(`//button[text()='Login']`, chromedp.BySearch),
		chromedp.Sleep(5*time.Second),

		// 等待并点击登录后页面的链接
		chromedp.WaitVisible(`a[href="/clients/"]`, chromedp.ByQuery), // 等待 Analytics Library 链接可见
		chromedp.Sleep(5*time.Second),                                 // 等待页面加载完成
	)
}