Go1.10通过http创建任务计划

package main

import (
    "encoding/json"
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "net/http"
    "net/smtp"
    "net/url"
    "os"
    "os/exec"
    "strings"

    "github.com/jakecoffman/cron"
)

var (
    cronpath         = "/etc/cronconfig.json"
    cronclient       bool
    cronclientArgs   = new(flag.FlagSet)
    cronconfig       CronConfig
    cronclientconfig CronClientConfig

    crontab      = cron.New()
    croncallback = func(err error) {}
    cronbaseauth = func(w http.ResponseWriter, r *http.Request) bool { return true }
)

func init() {
    if len(os.Args) > 1 {
        switch os.Args[1] {
        case "-h", "--help":
            fmt.Printf(`%s useage:
    %s 启动cron服务,默认使用/etc/cronconfig.json做配置路径
    %s /etc/cronconfig.json 通过指定配置文件启动
    %s client -h 作为客户端启动,详情使用-h查看帮助,语法详情:https://github.com/jakecoffman/cron`, os.Args[0], os.Args[0], os.Args[0], os.Args[0])
            os.Exit(1)
        case "client":
            cronclient = true
            cronclientArgs.BoolVar(&cronclientconfig.Add, "a", false, "-a=true 添加一个任务计划")
            cronclientArgs.BoolVar(&cronclientconfig.Del, "d", false, "-d=true 删除一个任务计划")
            cronclientArgs.BoolVar(&cronclientconfig.DelAll, "D", false, "-D=true 删除所有任务计划")
            cronclientArgs.BoolVar(&cronclientconfig.List, "l", false, "-l=true 查看所有任务计划")
            cronclientArgs.StringVar(&cronclientconfig.Tag, "tag", "", `-tag="deletelogcron" 根据指定的tag操作任务计划,结合-d或者-a使用`)
            cronclientArgs.StringVar(&cronclientconfig.Cron, "cron", "", `-cron='{\"log\":\"cron_01.log\",\"spec\":\"0 5 * * * *\",\"cmd\":\"ping\",\"args\":[\"baidu.com\"]}',结合-a使用`)
            cronclientArgs.StringVar(&cronclientconfig.BaseUser, "user", "", `-user="" 指定访问的baseauth用户名,可以为空`)
            cronclientArgs.StringVar(&cronclientconfig.BasePass, "pass", "", `-pass="" 指定访问的baseauth密码,可以为空`)
            cronclientArgs.StringVar(&cronclientconfig.Address, "addrs", "127.0.0.1:1789", `-addrs="h1,h2,h3" 指定要操作的目标地址`)
            cronclientArgs.Parse(os.Args[2:])
        default:
            cronpath = os.Args[1]
        }
    }

}

func main() {
    if !cronclient {
        if err := initcfg(cronpath, &cronconfig); err != nil {
            log.Fatalf("init config error:%s\n", err.Error())
        }
        initcronDaemon()
    } else {
        initClient()
    }
}

type Cron struct {
    Log  string   `json:"log"`
    Spec string   `json:"spec"`
    Cmd  string   `json:"cmd"`
    Args []string `json:"args"`
}

type CronConfig struct {
    Listen   string          `json:"listen"`
    LogPath  string          `json:"logpath"`
    BaseUser string          `json:"baseuser`
    BasePass string          `json:"basepass"`
    MailHost string          `json:"mailhost"`
    MailUser string          `json:"mailuser"`
    MailPass string          `json:"mailpass"`
    MailAddr string          `json:"mailaddr"`
    Cron     map[string]Cron `json:"cron"`
}

type CronClientConfig struct {
    Add    bool
    Del    bool
    DelAll bool
    List   bool

    Tag  string
    Cron string

    BaseUser string
    BasePass string
    Address  string
}

func initcfg(path string, cfg *CronConfig) error {
    buf, err := ioutil.ReadFile(path)
    if err != nil {
        return err
    }
    err = json.Unmarshal(buf, cfg)

    if cfg.LogPath != "" {
        File, err := os.OpenFile(cfg.LogPath, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
        if err != nil {
            return err
        }
        log.SetOutput(File)
    }

    if cfg.Listen == "" {
        cfg.Listen = "127.0.0.1:1789"
    }

    return nil
}

func router(w http.ResponseWriter, r *http.Request) {
    println(1)
    log.Printf("From:%s Method:%s Request:%s\n", r.RemoteAddr, r.Method, r.URL.RequestURI())
    if cronbaseauth(w, r) {
        tag := r.FormValue("tag")
        switch r.Method {
        case "GET":
            if r.URL.Path == "/list" {
                m := cronconfig.Cron
                if tag != "" {
                    if _, ok := m[tag]; ok {
                        m = map[string]Cron{tag: m[tag]}
                    } else {
                        m = map[string]Cron{}
                    }
                }
                buf, err := json.MarshalIndent(m, "", "\t")
                if err != nil {
                    http.Error(w, "Internal Server Error", http.StatusInternalServerError)
                } else {
                    w.Write(buf)
                }
                return
            }
        case "POST":
            if r.URL.Path == "/add" {
                if tag == "" {
                    http.Error(w, "tag is empty", http.StatusBadRequest)
                } else {
                    if _, ok := cronconfig.Cron[tag]; !ok {
                        var c Cron
                        err := json.Unmarshal([]byte(r.FormValue("cron")), &c)
                        if err == nil {
                            cronconfig.Cron[tag] = c
                            crontab.AddFunc(c.Spec, toFunc(c, croncallback), tag)
                            w.Write([]byte("successful"))
                            crontodisk()
                        } else {
                            http.Error(w, err.Error(), http.StatusBadRequest)
                        }
                    } else {
                        http.Error(w, "already is exist", http.StatusBadRequest)
                    }
                }
                return
            }
        case "DELETE":
            if r.URL.Path == "/delete" {
                if tag != "" || r.FormValue("all") == "true" {
                    if tag == "" {
                        for tag, _ := range cronconfig.Cron {
                            crontab.RemoveJob(tag)
                            delete(cronconfig.Cron, tag)
                        }
                    } else {
                        crontab.RemoveJob(tag)
                        delete(cronconfig.Cron, tag)
                    }
                    w.Write([]byte("successful"))
                    crontodisk()
                    return
                }
            }
        default:
            http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
        }
    } else {
        http.Error(w, "Unauthorized", http.StatusUnauthorized)
    }
    http.Error(w, "Forbidden", http.StatusForbidden)
}

func toFunc(c Cron, callback func(err error)) func() {
    cmd := exec.Command(c.Cmd, c.Args...)
    if c.Log != "" {
        File, err := os.OpenFile(c.Log, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
        if err == nil {
            cmd.Stdout = File
            cmd.Stderr = File
        }
    }
    return func() {
        if callback != nil {
            callback(cmd.Run())
        }
    }
}

func initcronDaemon() {
    if cronconfig.BaseUser != "" && cronconfig.BasePass != "" {
        cronbaseauth = func(w http.ResponseWriter, r *http.Request) bool {
            u, p, ok := r.BasicAuth()
            if !ok {
                return ok
            }
            return u == cronconfig.BaseUser && p == cronconfig.BasePass
        }
    }

    if cronconfig.MailHost != "" && cronconfig.MailUser != "" && cronconfig.MailPass != "" {
        auth := smtp.PlainAuth("", cronconfig.MailUser, cronconfig.MailPass, strings.Split(cronconfig.MailPass, ":")[0])
        base := "To: " + cronconfig.MailAddr + "\r\nFrom: " + cronconfig.MailUser + ">\r\nSubject: " + "\r\n" + `Content-Type: text/plain; charset=UTF-8` + "\r\n\r\n"
        croncallback = func(err error) {
            if err != nil {
                msg := []byte(base + err.Error())
                err := smtp.SendMail(cronconfig.MailHost, auth, cronconfig.MailAddr, []string{cronconfig.MailAddr}, msg)
                if err != nil {
                    log.Printf("send mail error:%s\n", err.Error())
                }
            }
        }
    }

    for tag, c := range cronconfig.Cron {
        crontab.AddFunc(c.Spec, toFunc(c, croncallback), tag)
    }

    crontab.Start()

    http.HandleFunc("/", router)
    err := http.ListenAndServe(cronconfig.Listen, nil)
    if err != nil {
        log.Fatalf("listen error:%s\n", err.Error())
    }
}

func crontodisk() {
    buf, err := json.MarshalIndent(cronconfig, "", "\t")
    if err != nil {
        log.Printf("CronToDisk marshal error:%s\n", err.Error())
        return
    }
    File, err := os.Create(cronpath)
    if err != nil {
        log.Printf("CronToDisk error:%s\n", err.Error())
        return
    }
    File.Write(buf)
    File.Close()
}

func initClient() {
    var (
        err    error
        req    *http.Request
        method string
        format func(str string) string
    )

    list := strings.Split(cronclientconfig.Address, ",")
    if len(list) == 0 {
        log.Fatalf("必须指定要操作的主机地址")
    }

    switch {
    case cronclientconfig.Add:
        if cronclientconfig.Tag == "" || cronclientconfig.Cron == "" {
            log.Fatalf("使用-a参数,必须指定-tag和-spec参数")
        }
        method = "POST"
        var args = url.Values{
            "tag":  []string{cronclientconfig.Tag},
            "cron": []string{cronclientconfig.Cron},
        }

        format = func(str string) string {
            return fmt.Sprintf("http://%s/add?%s", str, args.Encode())
        }

    case cronclientconfig.Del || cronclientconfig.DelAll:
        var uri string
        if cronclientconfig.Del {
            if cronclientconfig.Tag == "" {
                log.Fatalf("如果使用-d参数,必须指定-tag参数\n")
            }
            uri = "/delete?tag=" + cronclientconfig.Tag
        } else {
            uri = "/delete?all=true"
        }
        method = "DELETE"
        format = func(str string) string {
            return fmt.Sprintf("http://%s%s", str, uri)
        }
    case cronclientconfig.List:
        method = "GET"
        if cronclientconfig.Tag == "" {
            format = func(str string) string {
                return fmt.Sprintf("http://%s/list", str)
            }
        } else {
            format = func(str string) string {
                return fmt.Sprintf("http://%s/list?tag=%s", str, cronclientconfig.Tag)
            }
        }
    default:
        os.Exit(1)
    }
    for _, addr := range list {
        req, err = http.NewRequest(method, format(addr), nil)
        if err != nil {
            log.Printf("%s create request error:%s\n", cronclientconfig.Tag, err.Error())
            continue
        }
        if cronclientconfig.BaseUser != "" && cronclientconfig.BasePass != "" {
            req.SetBasicAuth(cronclientconfig.BaseUser, cronclientconfig.BasePass)
        }
        resp, err := http.DefaultClient.Do(req)
        if err != nil {
            log.Printf("%s request error:%s\n", cronclientconfig.Tag, err.Error())
            continue
        }
        if resp.StatusCode == 200 {
            buf, err := ioutil.ReadAll(resp.Body)
            if err == nil {
                log.Printf("Output:\n%s\n", string(buf))
            }
        } else {
            buf, err := ioutil.ReadAll(resp.Body)
            if err != nil {
                log.Printf("%s return error:%s\n", cronclientconfig.Tag, http.StatusText(resp.StatusCode))
            } else {
                log.Printf("%s return error:%s\n", cronclientconfig.Tag, string(buf))
            }
        }
    }
}

猜你喜欢

转载自blog.csdn.net/fyxichen/article/details/79375337