d2vm/cmd/d2vm/main.go

// Copyright 2022 Linka Cloud  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/color"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"go.linka.cloud/d2vm"
	"go.linka.cloud/d2vm/pkg/exec"
)

var (
	verbose    = false
	timeFormat = ""
	format     = "qcow2"

	rootCmd = &cobra.Command{
		Use:          "d2vm",
		SilenceUsage: true,
		Version:      d2vm.Version,
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			switch timeFormat {
			case "full", "f":
			case "relative", "rel", "r":
			case "none", "":
			default:
				logrus.Fatalf("invalid time format: %s. Valid format: 'relative', 'full'", timeFormat)
			}
			if verbose {
				logrus.SetLevel(logrus.TraceLevel)
			}
			exec.SetDebug(verbose)

			// make the zsh completion work when sourced with `source <(d2vm completion zsh)`
			if cmd.Name() == "zsh" && cmd.Parent() != nil && cmd.Parent().Name() == "completion" {
				zshHead := fmt.Sprintf("#compdef %[1]s\ncompdef _%[1]s %[1]s\n", cmd.Root().Name())
				cmd.OutOrStdout().Write([]byte(zshHead))
			}
		},
	}
)

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, os.Interrupt, os.Kill)
	go func() {
		<-sigs
		fmt.Println()
		cancel()
	}()
	rootCmd.ExecuteContext(ctx)
}

func init() {
	rootCmd.PersistentFlags().BoolVarP(&verbose, "debug", "d", false, "Enable Debug output")
	rootCmd.PersistentFlags().MarkDeprecated("debug", "use -v instead")
	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "Enable Verbose output")
	rootCmd.PersistentFlags().StringVar(&timeFormat, "time", "none", "Enable formated timed output, valide formats: 'relative (rel | r)', 'full (f)'")
	color.NoColor = false
	logrus.StandardLogger().Formatter = &logfmtFormatter{start: time.Now()}
}

const (
	red    = 31
	yellow = 33
	blue   = 36
	white  = 39
	gray   = 90
)

type logfmtFormatter struct {
	start time.Time
}

func (f *logfmtFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b bytes.Buffer
	var c *color.Color
	switch entry.Level {
	case logrus.DebugLevel, logrus.TraceLevel:
		c = color.New(gray)
	case logrus.WarnLevel:
		c = color.New(yellow)
	case logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel:
		c = color.New(red)
	default:
		c = color.New(white)
	}
	msg := entry.Message
	if len(entry.Message) > 0 && entry.Level < logrus.DebugLevel {
		msg = strings.ToTitle(string(msg[0])) + msg[1:]
	}

	var err error
	switch timeFormat {
	case "full", "f":
		_, err = c.Fprintf(&b, "[%s] %s\n", entry.Time.Format("2006-01-02 15:04:05"), entry.Message)
	case "relative", "rel", "r":
		_, err = c.Fprintf(&b, "[%5v] %s\n", entry.Time.Sub(f.start).Truncate(time.Second).String(), msg)
	default:
		_, err = c.Fprintln(&b, msg)
	}
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func isRoot() bool {
	return os.Geteuid() == 0
}

func sudoUser() (uid int, sudo bool) {
	// if we are not running on linux, docker handle files user's permissions,
	// so we don't need to check for sudo here
	if runtime.GOOS != "linux" {
		return
	}
	v := os.Getenv("SUDO_UID")
	if v == "" {
		return 0, false
	}
	uid, err := strconv.Atoi(v)
	if err != nil {
		logrus.Errorf("invalid SUDO_UID: %s", v)
		return 0, false
	}
	return uid, true
}