mirror of
https://github.com/putyy/res-downloader.git
synced 2026-01-12 06:04:55 +08:00
feat: Added header cache, optional header settings...
This commit is contained in:
@@ -24,6 +24,7 @@ type Config struct {
|
||||
WxAction bool `json:"WxAction"`
|
||||
TaskNumber int `json:"TaskNumber"`
|
||||
UserAgent string `json:"UserAgent"`
|
||||
UseHeaders string `json:"UseHeaders"`
|
||||
}
|
||||
|
||||
func initConfig() *Config {
|
||||
@@ -43,7 +44,8 @@ func initConfig() *Config {
|
||||
"AutoProxy": true,
|
||||
"WxAction": true,
|
||||
"TaskNumber": __TaskNumber__,
|
||||
"UserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36"
|
||||
"UserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
|
||||
"UseHeaders": "User-Agent,Referer,Authorization,Cookie"
|
||||
}
|
||||
`
|
||||
def = strings.ReplaceAll(def, "__TaskNumber__", strconv.Itoa(runtime.NumCPU()*2))
|
||||
@@ -51,9 +53,23 @@ func initConfig() *Config {
|
||||
storage: NewStorage("config.json", []byte(def)),
|
||||
}
|
||||
|
||||
defaultMap := make(map[string]interface{})
|
||||
_ = json.Unmarshal([]byte(def), &defaultMap)
|
||||
|
||||
data, err := globalConfig.storage.Load()
|
||||
if err == nil {
|
||||
_ = json.Unmarshal(data, &globalConfig)
|
||||
var loadedMap map[string]interface{}
|
||||
_ = json.Unmarshal(data, &loadedMap)
|
||||
|
||||
for key, val := range defaultMap {
|
||||
if _, ok := loadedMap[key]; !ok {
|
||||
loadedMap[key] = val
|
||||
}
|
||||
}
|
||||
|
||||
finalBytes, _ := json.Marshal(loadedMap)
|
||||
_ = json.Unmarshal(finalBytes, &globalConfig)
|
||||
|
||||
} else {
|
||||
globalLogger.Esg(err, "load config err")
|
||||
}
|
||||
@@ -77,6 +93,7 @@ func (c *Config) setConfig(config Config) {
|
||||
c.AutoProxy = config.AutoProxy
|
||||
c.TaskNumber = config.TaskNumber
|
||||
c.WxAction = config.WxAction
|
||||
c.UseHeaders = config.UseHeaders
|
||||
if oldProxy != c.UpstreamProxy {
|
||||
proxyOnce.setTransport()
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/cookiejar"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -32,17 +31,19 @@ type FileDownloader struct {
|
||||
totalTasks int
|
||||
TotalSize int64
|
||||
IsMultiPart bool
|
||||
Headers map[string]string
|
||||
DownloadTaskList []*DownloadTask
|
||||
progressCallback ProgressCallback
|
||||
}
|
||||
|
||||
func NewFileDownloader(url, filename string, totalTasks int) *FileDownloader {
|
||||
func NewFileDownloader(url, filename string, totalTasks int, headers map[string]string) *FileDownloader {
|
||||
return &FileDownloader{
|
||||
Url: url,
|
||||
FileName: filename,
|
||||
totalTasks: totalTasks,
|
||||
IsMultiPart: false,
|
||||
TotalSize: 0,
|
||||
Headers: headers,
|
||||
DownloadTaskList: make([]*DownloadTask, 0),
|
||||
}
|
||||
}
|
||||
@@ -53,10 +54,16 @@ func (fd *FileDownloader) buildClient() *http.Client {
|
||||
transport.Proxy = http.ProxyURL(fd.ProxyUrl)
|
||||
}
|
||||
// Cookie handle
|
||||
jar, _ := cookiejar.New(nil)
|
||||
return &http.Client{
|
||||
Transport: transport,
|
||||
Jar: jar,
|
||||
}
|
||||
}
|
||||
|
||||
func (fd *FileDownloader) setHeaders(request *http.Request) {
|
||||
for key, values := range fd.Headers {
|
||||
if strings.Contains(globalConfig.UseHeaders, key) {
|
||||
request.Header.Set(key, values)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,21 +83,22 @@ func (fd *FileDownloader) init() error {
|
||||
}
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("HEAD", fd.Url, nil)
|
||||
request, err := http.NewRequest("HEAD", fd.Url, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create request failed")
|
||||
}
|
||||
|
||||
// 设置请求头
|
||||
if globalConfig.UserAgent != "" {
|
||||
req.Header.Set("User-Agent", globalConfig.UserAgent)
|
||||
if _, ok := fd.Headers["User-Agent"]; !ok {
|
||||
fd.Headers["User-Agent"] = globalConfig.UserAgent
|
||||
}
|
||||
|
||||
if fd.Referer != "" {
|
||||
req.Header.Set("Referer", fd.Referer)
|
||||
if _, ok := fd.Headers["Referer"]; !ok {
|
||||
fd.Headers["Referer"] = fd.Referer
|
||||
}
|
||||
|
||||
resp, err := fd.buildClient().Do(req)
|
||||
fd.setHeaders(request)
|
||||
|
||||
resp, err := fd.buildClient().Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("request failed" + err.Error())
|
||||
}
|
||||
@@ -182,8 +190,9 @@ func (fd *FileDownloader) startDownloadTask(waitGroup *sync.WaitGroup, progressC
|
||||
globalLogger.Error().Stack().Err(err).Msgf("任务%d创建请求出错", task.taskID)
|
||||
return
|
||||
}
|
||||
request.Header.Set("User-Agent", globalConfig.UserAgent)
|
||||
request.Header.Set("Referer", fd.Referer)
|
||||
|
||||
fd.setHeaders(request)
|
||||
|
||||
if fd.IsMultiPart {
|
||||
request.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", task.rangeStart, task.rangeEnd))
|
||||
}
|
||||
|
||||
@@ -336,6 +336,12 @@ func (p *Proxy) httpResponseEvent(resp *http.Response, ctx *goproxy.ProxyCtx) *h
|
||||
Description: "",
|
||||
ContentType: resp.Header.Get("Content-Type"),
|
||||
}
|
||||
|
||||
// Store entire request headers as JSON
|
||||
if headers, err := json.Marshal(resp.Request.Header); err == nil {
|
||||
res.OtherData["headers"] = string(headers)
|
||||
}
|
||||
|
||||
resourceOnce.markMedia(urlSign)
|
||||
httpServerOnce.send("newResources", res)
|
||||
}
|
||||
|
||||
@@ -2,6 +2,8 @@ package core
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"os"
|
||||
@@ -142,7 +144,9 @@ func (r *Resource) download(mediaInfo MediaInfo, decodeStr string) {
|
||||
}
|
||||
}
|
||||
|
||||
downloader := NewFileDownloader(rawUrl, mediaInfo.SavePath, globalConfig.TaskNumber)
|
||||
headers, _ := r.parseHeaders(mediaInfo)
|
||||
|
||||
downloader := NewFileDownloader(rawUrl, mediaInfo.SavePath, globalConfig.TaskNumber, headers)
|
||||
downloader.progressCallback = func(totalDownloaded float64) {
|
||||
r.progressEventsEmit(mediaInfo, strconv.Itoa(int(totalDownloaded))+"%", DownloadStatusRunning)
|
||||
}
|
||||
@@ -162,6 +166,27 @@ func (r *Resource) download(mediaInfo MediaInfo, decodeStr string) {
|
||||
}(mediaInfo)
|
||||
}
|
||||
|
||||
// 解析并组装 headers
|
||||
func (r *Resource) parseHeaders(mediaInfo MediaInfo) (map[string]string, error) {
|
||||
headers := make(map[string]string)
|
||||
|
||||
if hh, ok := mediaInfo.OtherData["headers"]; ok {
|
||||
var tempHeaders map[string][]string
|
||||
// 解析 JSON 字符串为 map[string][]string
|
||||
if err := json.Unmarshal([]byte(hh), &tempHeaders); err != nil {
|
||||
return headers, fmt.Errorf("parse headers JSON err: %v", err)
|
||||
}
|
||||
|
||||
for key, values := range tempHeaders {
|
||||
if len(values) > 0 {
|
||||
headers[key] = values[0]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
func (r *Resource) wxFileDecode(mediaInfo MediaInfo, fileName, decodeStr string) (string, error) {
|
||||
sourceFile, err := os.Open(fileName)
|
||||
if err != nil {
|
||||
|
||||
Reference in New Issue
Block a user