123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243 |
- package commutil
- import (
- "archive/zip"
- "bytes"
- "errors"
- "fmt"
- "io"
- "io/ioutil"
- "mime/multipart"
- "net/http"
- "os"
- "path/filepath"
- sp "spiderutil"
- "strconv"
- "strings"
- )
- // 任务发布
- func publishModule(taskCode, scriptCode, params string, workStart, workEnd, workTimeout, interval int) (err error) {
- if workStart < 0 || workEnd < 0 || workStart > workEnd || workEnd > 23 {
- return errors.New("起始或结束时间异常!")
- } else {
- req, err := http.NewRequest("GET", fmt.Sprintf("%s/publish", sp.Config.TaskMgrSerever), nil)
- if err != nil {
- return errors.New("服务端链接创建失败!")
- }
- req.Header.Set("task_code", taskCode)
- req.Header.Set("script_code", scriptCode)
- req.Header.Set("work_start", strconv.Itoa(workStart))
- req.Header.Set("work_end", strconv.Itoa(workEnd))
- req.Header.Set("worktimeout", strconv.Itoa(workTimeout))
- req.Header.Set("interval", strconv.Itoa(interval))
- req.Header.Set("params", params)
- client := httpClient()
- resp, err := client.Do(req)
- if err != nil {
- return errors.New("服务端发布请求失败!")
- }
- defer resp.Body.Close()
- }
- return
- }
- // postModule 推送模块
- func postModule(path, module string) error {
- name := module
- bs, err := dir2Zip(path)
- if err != nil {
- return err
- }
- uploadUrl := fmt.Sprintf("%s/push", sp.Config.CodeDbSerever)
- // 创建一个 multipart.FormDataWriter
- body := bytes.NewBuffer(nil)
- writer := multipart.NewWriter(body)
- // 添加文件
- part, err := writer.CreateFormFile("file", "1.zip")
- if err != nil {
- return err
- }
- _, err = io.Copy(part, bytes.NewReader(bs))
- if err != nil {
- return err
- }
- // 添加表单数据
- err = writer.WriteField("name", name)
- if err != nil {
- return err
- }
- // 完成 multipart 插头
- err = writer.Close()
- if err != nil {
- return err
- }
- // 创建一个表示 multipart/form-data 的 http.Client
- client := httpClient()
- // 构建 POST 请求
- req, err := http.NewRequest("POST", uploadUrl, nil)
- if err != nil {
- return err
- }
- // 设置 Content-Type 为 multipart/form-data
- req.Header.Set("Content-Type", writer.FormDataContentType())
- // 设置请求体
- req.Body = ioutil.NopCloser(body)
- // 发送请求
- resp, err := client.Do(req)
- if err != nil {
- return err
- }
- defer resp.Body.Close()
- // 读取响应
- _, err = ioutil.ReadAll(resp.Body)
- return err
- }
- // unPubModule 撤销任务发布
- func unPubModule(taskCode string) error {
- req, err := http.NewRequest("GET", fmt.Sprintf("%s/unpublish", sp.Config.CodeDbSerever), nil)
- if err != nil {
- return err
- }
- req.Header.Add("task_code", taskCode)
- client := httpClient()
- resp, err := client.Do(req)
- if err != nil {
- return err
- }
- defer resp.Body.Close()
- return nil
- }
- // 压缩目录为zip文件
- func dir2Zip(srcPath string) (bs []byte, err error) {
- buf := bytes.NewBuffer(nil)
- zipWriter := zip.NewWriter(buf)
- // 遍历目录并添加文件到 ZIP 归档
- err = filepath.Walk(srcPath, func(path string, info os.FileInfo, err error) error {
- if err != nil {
- return err
- }
- // 跳过根目录
- if path == srcPath {
- return nil
- }
- // 计算相对于根目录的路径
- relPath, err := filepath.Rel(srcPath, path)
- if err != nil {
- return err
- }
- // 创建 ZIP 文件条目
- header, err := zip.FileInfoHeader(info)
- if err != nil {
- return err
- }
- // 设置文件条目的名称
- header.Name = relPath
- // 处理目录
- if info.IsDir() {
- header.Name += "/"
- } else {
- header.Method = zip.Deflate
- }
- // 写入 ZIP 文件条目
- writer, err := zipWriter.CreateHeader(header)
- if err != nil {
- return err
- }
- // 如果是文件,写入文件内容
- if !info.IsDir() {
- file, err := os.Open(path)
- if err != nil {
- return err
- }
- defer file.Close()
- _, err = io.Copy(writer, file)
- if err != nil {
- return err
- }
- zipWriter.Flush()
- }
- return nil
- })
- zipWriter.Close()
- bs = buf.Bytes()
- return
- }
- // pullModule 拉取模块
- func pullModule(path, module string) error {
- pullUrl := fmt.Sprintf("%s/pull", sp.Config.CodeDbSerever)
- client := httpClient()
- req, err := http.NewRequest("POST", pullUrl, strings.NewReader(fmt.Sprintf("%s=%s", "name", module)))
- if err != nil {
- return err
- }
- req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
- resp, err := client.Do(req)
- if err != nil {
- return err
- }
- version := resp.Header.Get("code-version")
- fmt.Println("version", version)
- bs, err := ioutil.ReadAll(resp.Body)
- if err != nil {
- return err
- }
- resp.Body.Close()
- //保存临时文件
- // os.WriteFile(fmt.Sprintf("%s_%s_%04d.zip", path,
- // time.Now().Format("20060102150405"),
- // rand.Intn(10000)), bs, 0777)
- //zip解压
- return unzip(bs, path)
- }
- // 解压缩zip文件到指定目录
- func unzip(zipData []byte, destDir string) error {
- // 打开 zip 文件
- r, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
- if err != nil {
- return errors.New("open zipfile " + err.Error())
- }
- // 确保目标目录存在
- err = os.MkdirAll(destDir, 0755)
- if err != nil {
- return errors.New("open zipfile " + err.Error())
- }
- // 遍历 zip 文件中的所有条目
- for _, f := range r.File {
- // 打开文件
- func() error {
- rc, err := f.Open()
- if err != nil {
- return err
- }
- defer rc.Close()
- // 处理文件
- path := f.Name
- if path == "/" {
- return errors.New("skip root dir") // 跳过根目录
- }
- // 创建目标文件
- w, err := os.Create(destDir + "/" + path)
- if err != nil {
- return err
- }
- defer w.Close()
- // 复制文件内容
- _, err = io.Copy(w, rc)
- return err
- }()
- }
- return nil
- }
|