Merge pull request #692 from kongfei605/feature_install

self upgrade

agent/install/service_linux.go

@@ -4,7 +4,9 @@ import (
 	"bytes"
 	"fmt"
 	"log"
+	"os"
 	"os/exec"
+	"path/filepath"
 	"strings"
 	"time"
 
@@ -203,8 +205,7 @@ func ServiceConfig() *service.Config {
 		option["SysvScript"] = sysvScript
 		option["LogOutput"] = true
 	}
-
-	return &service.Config{
+	cfg := &service.Config{
 		// 服务显示名称
 		Name: ServiceName,
 		// 服务名称
@@ -216,4 +217,15 @@ func ServiceConfig() *service.Config {
 		Dependencies: depends, //
 		Option:       option,
 	}
+
+	ov, err := os.Executable()
+	if err == nil {
+		if len(filepath.Dir(ov)) != 0 {
+			cfg.WorkingDirectory = filepath.Dir(ov)
+		}
+	} else {
+		log.Println("E! get exeutable path error:", err)
+	}
+	cfg.Arguments = []string{"-configs", filepath.Dir(ov) + "/conf"}
+	return cfg
 }

agent/update/update_darwin.go

@@ -0,0 +1,10 @@
+package update
+
+import (
+	"fmt"
+)
+
+func Update(tar string) error {
+	// binary
+	return fmt.Errorf("linux support only")
+}

agent/update/update_freebsd.go

@@ -0,0 +1,9 @@
+package update
+
+import (
+	"fmt"
+)
+
+func Update(tar string) error {
+	return fmt.Errorf("linux support only")
+}

agent/update/update_linux.go

@@ -0,0 +1,155 @@
+package update
+
+import (
+	"archive/tar"
+	"bufio"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+func download(file string) (string, error) {
+	fname := filepath.Base(file)
+	log.Println("downloading file:", file, "save to:", fname)
+	res, err := http.Get(file)
+	if err != nil {
+		return fname, fmt.Errorf("cannot download file from %s", file)
+	}
+	defer res.Body.Close()
+
+	if res.StatusCode != 200 {
+		return fname, fmt.Errorf("download  %s error response: %s", file, res.Status)
+	}
+	f, err := os.OpenFile(fname, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
+	if err != nil {
+		return fname, err
+	}
+	defer f.Close()
+	bufWriter := bufio.NewWriter(f)
+
+	_, err = io.Copy(bufWriter, res.Body)
+	if err != nil {
+		return fname, err
+	}
+	bufWriter.Flush()
+	return fname, nil
+}
+
+func Update(tar string) error {
+	// download
+	fname, err := download(tar)
+	if err != nil {
+		return err
+	}
+	nv, err := UnTar("./", fname)
+	if err != nil {
+		return err
+	}
+
+	// old version
+	ov, err := os.Executable()
+	if err != nil {
+		return err
+	}
+	fm, err := os.Stat(ov)
+	if err != nil {
+		return err
+	}
+	fi, err := os.Stat(ov)
+	if err != nil {
+		return err
+	}
+	if fi.Mode().IsDir() {
+		return fmt.Errorf("%s is directory", ov)
+	}
+	log.Printf("I! replace old version:%s with new version:%s", ov, "./"+nv)
+
+	// replace
+	err = os.Rename(nv, ov)
+	if err != nil {
+		return err
+	}
+	err = os.RemoveAll("./" + filepath.Dir(nv))
+	if err != nil {
+		log.Println("E! clean dir:", "./"+filepath.Dir(nv), "error:", err)
+	} else {
+		log.Println("I! clean dir:", "./"+filepath.Dir(nv), "success")
+	}
+	err = os.Remove("./" + fname)
+	if err != nil {
+		log.Println("E! clean file:", "./"+fname, "error:", err)
+	} else {
+		log.Println("I! clean file:", "./"+fname, "success")
+	}
+	return os.Chmod(ov, fm.Mode().Perm())
+}
+
+func UnTar(dst, src string) (target string, err error) {
+	fr, err := os.Open(src)
+	if err != nil {
+		return
+	}
+	defer fr.Close()
+
+	gr, err := gzip.NewReader(fr)
+	if err != nil {
+		return
+	}
+	defer gr.Close()
+
+	tr := tar.NewReader(gr)
+
+	for {
+		hdr, err := tr.Next()
+
+		switch {
+		case err == io.EOF:
+			return target, nil
+		case err != nil:
+			return target, err
+		case hdr == nil:
+			continue
+		}
+
+		dstFileDir := filepath.Join(dst, hdr.Name)
+
+		switch hdr.Typeflag {
+		case tar.TypeDir:
+			if b := ExistDir(dstFileDir); !b {
+				if err := os.MkdirAll(dstFileDir, 0775); err != nil {
+					return target, err
+				}
+			}
+		case tar.TypeReg:
+			err := os.MkdirAll(filepath.Dir(dstFileDir), 0755)
+			if err != nil {
+				log.Printf("mdkir:%s, error:%s", filepath.Base(dstFileDir), err)
+				return target, err
+			}
+			file, err := os.OpenFile(dstFileDir, os.O_CREATE|os.O_RDWR, os.FileMode(hdr.Mode))
+			if err != nil {
+				return target, err
+			}
+			if strings.HasSuffix(dstFileDir, "categraf") {
+				target = dstFileDir
+			}
+			_, err = io.Copy(file, tr)
+			if err != nil {
+				return target, err
+			}
+			file.Close()
+		}
+	}
+
+	return target, nil
+}
+
+func ExistDir(dirname string) bool {
+	fi, err := os.Stat(dirname)
+	return (err == nil || os.IsExist(err)) && fi.IsDir()
+}

agent/update/update_windows.go

@@ -0,0 +1,9 @@
+package update
+
+import (
+	"fmt"
+)
+
+func Update(tar string) error {
+	return fmt.Errorf("linux support only")
+}

main.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strings"
 	"syscall"
 
 	"github.com/chai2010/winsvc"
@@ -18,6 +19,7 @@ import (
 
 	"flashcat.cloud/categraf/agent"
 	agentInstall "flashcat.cloud/categraf/agent/install"
+	agentUpdate "flashcat.cloud/categraf/agent/update"
 	"flashcat.cloud/categraf/api"
 	"flashcat.cloud/categraf/config"
 	"flashcat.cloud/categraf/heartbeat"
@@ -38,6 +40,8 @@ var (
 	start        = flag.Bool("start", false, "Start categraf service")
 	stop         = flag.Bool("stop", false, "Stop categraf service")
 	status       = flag.Bool("status", false, "Show categraf service status")
+	update       = flag.Bool("update", false, "Update categraf binary")
+	updateFile   = flag.String("update_url", "", "new version for categraf to download")
 )
 
 func init() {
@@ -78,7 +82,7 @@ func main() {
 		fmt.Println(config.Version)
 		os.Exit(0)
 	}
-	if *install || *remove || *start || *stop || *status {
+	if *install || *remove || *start || *stop || *status || *update {
 		err := serviceProcess()
 		if err != nil {
 			log.Println("E!", err)
@@ -270,5 +274,36 @@ func serviceProcess() error {
 
 		return nil
 	}
+	if *update {
+		if *updateFile == "" {
+			return fmt.Errorf("please input update_url")
+		}
+		if sts, err := s.Status(); err != nil {
+			if strings.Contains(err.Error(), "not installed") {
+				log.Println("E! update only support mode that running in service mode")
+			}
+			return nil
+		} else {
+			switch sts {
+			case service.StatusRunning:
+				log.Println("I! categraf service status: running, version:", config.Version)
+			case service.StatusStopped:
+				log.Println("I! categraf service status: stopped, version:", config.Version)
+			default:
+				log.Println("I! categraf service status: unknown, version:", config.Version)
+			}
+		}
+		err := agentUpdate.Update(*updateFile)
+		if err != nil {
+			log.Println("E! update categraf failed:", err)
+			return nil
+		}
+		err = s.Restart()
+		if err != nil {
+			log.Println("E! restart categraf failed:", err)
+			return nil
+		}
+		log.Println("I! update categraf success")
+	}
 	return nil
 }