spray/cmd/cmd.go

package cmd

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/chainreactors/logs"
	"github.com/chainreactors/spray/core"
	"github.com/chainreactors/spray/pkg"
	"github.com/jessevdk/go-flags"
)

var ver = "dev"
var DefaultConfig = "config.yaml"

func init() {
	logs.Log.SetColorMap(map[logs.Level]func(string) string{
		logs.InfoLevel:      logs.PurpleBold,
		logs.ImportantLevel: logs.GreenBold,
		pkg.LogVerbose:      logs.Green,
	})
}

func Spray() {
	var option core.Option

	parser := flags.NewParser(&option, flags.Default)
	parser.Usage = core.Usage()
	_, err := parser.ParseArgs(os.Args[1:])
	if err != nil {
		if err.(*flags.Error).Type != flags.ErrHelp {
			fmt.Println(err.Error())
		}
		return
	}

	defer time.Sleep(time.Second)
	ctx, canceler := context.WithTimeout(context.Background(), time.Duration(option.Deadline)*time.Second)
	go func() {
		select {
		case <-ctx.Done():
			time.Sleep(10 * time.Second)
			logs.Log.Errorf("deadline and timeout not work, hard exit!!!")
			os.Exit(0)
		}
	}()

	go func() {
		exitChan := make(chan os.Signal, 2)
		signal.Notify(exitChan, os.Interrupt, syscall.SIGTERM)

		go func() {
			sigCount := 0
			for {
				<-exitChan
				sigCount++
				if sigCount == 1 {
					logs.Log.Infof("Exit signal received, saving task and exiting...")
					canceler()
				} else if sigCount == 2 {
					logs.Log.Infof("forcing exit...")
					os.Exit(1)
				}
			}
		}()
	}()

	if err := core.RunWithArgs(ctx, os.Args[1:], core.RunOptions{
		DefaultConfig: DefaultConfig,
		Version:       ver,
	}); err != nil {
		logs.Log.Errorf("%s", err.Error())
		return
	}
}