tunasync/worker/cmd_provider.go
2016-04-30 16:38:38 +08:00

119 lines
2.3 KiB
Go

package worker
import (
"os"
"os/exec"
"strings"
"github.com/anmitsu/go-shlex"
"github.com/codeskyblue/go-sh"
)
type cmdConfig struct {
name string
upstreamURL, command string
workingDir, logDir, logFile string
interval int
env map[string]string
}
type cmdProvider struct {
baseProvider
cmdConfig
command []string
cmd *exec.Cmd
session *sh.Session
}
func newCmdProvider(c cmdConfig) (*cmdProvider, error) {
// TODO: check config options
provider := &cmdProvider{
baseProvider: baseProvider{
name: c.name,
ctx: NewContext(),
interval: c.interval,
},
cmdConfig: c,
}
provider.ctx.Set(_WorkingDirKey, c.workingDir)
provider.ctx.Set(_LogDirKey, c.logDir)
provider.ctx.Set(_LogFileKey, c.logFile)
cmd, err := shlex.Split(c.command, true)
if err != nil {
return nil, err
}
provider.command = cmd
return provider, nil
}
// Copied from go-sh
func newEnviron(env map[string]string, inherit bool) []string { //map[string]string {
environ := make([]string, 0, len(env))
if inherit {
for _, line := range os.Environ() {
// if os environment and env collapses,
// omit the os one
k := strings.Split(line, "=")[0]
if _, ok := env[k]; ok {
continue
}
environ = append(environ, line)
}
}
for k, v := range env {
environ = append(environ, k+"="+v)
}
return environ
}
// TODO: implement this
func (p *cmdProvider) Run() error {
if len(p.command) == 1 {
p.cmd = exec.Command(p.command[0])
} else if len(p.command) > 1 {
c := p.command[0]
args := p.command[1:]
p.cmd = exec.Command(c, args...)
} else if len(p.command) == 0 {
panic("Command length should be at least 1!")
}
p.cmd.Dir = p.WorkingDir()
env := map[string]string{
"TUNASYNC_MIRROR_NAME": p.Name(),
"TUNASYNC_WORKING_DIR": p.WorkingDir(),
"TUNASYNC_UPSTREAM_URL": p.upstreamURL,
"TUNASYNC_LOG_FILE": p.LogFile(),
}
for k, v := range p.env {
env[k] = v
}
p.cmd.Env = newEnviron(env, true)
logFile, err := os.OpenFile(p.LogFile(), os.O_WRONLY|os.O_CREATE, 0644)
if err != nil {
return err
}
p.cmd.Stdout = logFile
p.cmd.Stderr = logFile
return p.cmd.Start()
}
func (p *cmdProvider) Wait() error {
return p.cmd.Wait()
}
// TODO: implement this
func (p *cmdProvider) Terminate() {
}
// TODO: implement this
func (p *cmdProvider) Hooks() {
}