129 lines
2.8 KiB
Go
129 lines
2.8 KiB
Go
package cmd
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"os"
|
|
|
|
logger "git.towk2.me/towk/makeshift/pkg/log"
|
|
"github.com/rs/zerolog/log"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var (
|
|
loglevel logger.LogLevel = logger.INFO
|
|
)
|
|
|
|
var rootCmd = cobra.Command{
|
|
Use: "makeshift",
|
|
Short: "Extensible file cobbler",
|
|
PersistentPreRun: func(cmd *cobra.Command, args []string) {
|
|
var (
|
|
logFile string
|
|
err error
|
|
)
|
|
|
|
// initialize the logger
|
|
logFile, _ = cmd.Flags().GetString("log-file")
|
|
err = logger.InitWithLogLevel(loglevel, logFile)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("failed to initialize logger")
|
|
os.Exit(1)
|
|
}
|
|
},
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
// try and set flags using env vars
|
|
setenv(cmd, "log-file", "MAKESHIFT_LOG_FILE")
|
|
setenv(cmd, "log-level", "MAKESHIFT_LOG_LEVEL")
|
|
if len(args) == 0 {
|
|
err := cmd.Help()
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("failed to print help")
|
|
}
|
|
os.Exit(0)
|
|
}
|
|
},
|
|
PostRun: func(cmd *cobra.Command, args []string) {
|
|
log.Debug().Msg("closing log file")
|
|
err := logger.LogFile.Close()
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("failed to close log file")
|
|
}
|
|
},
|
|
}
|
|
|
|
func Execute() {
|
|
// run the main program
|
|
if err := rootCmd.Execute(); err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func init() {
|
|
cobra.OnInitialize(
|
|
initLogger,
|
|
)
|
|
// initialize the config a single time
|
|
rootCmd.PersistentFlags().VarP(&loglevel, "log-level", "l", "Set the log level output")
|
|
rootCmd.PersistentFlags().String("log-file", "", "Set the log file path (can be set with MAKESHIFT_LOG_FILE)")
|
|
}
|
|
|
|
func setenv(cmd *cobra.Command, varname string, envvar string) {
|
|
if cmd.Flags().Changed(varname) {
|
|
return
|
|
}
|
|
val := os.Getenv(envvar)
|
|
if val != "" {
|
|
cmd.Flags().Set(varname, val)
|
|
}
|
|
}
|
|
|
|
func setenvp(cmd *cobra.Command, varname string, envvar string) {
|
|
if cmd.Flags().Changed(varname) {
|
|
return
|
|
}
|
|
val := os.Getenv(envvar)
|
|
if val != "" {
|
|
cmd.PersistentFlags().Set(varname, val)
|
|
}
|
|
}
|
|
|
|
func initLogger() {
|
|
// initialize the logger
|
|
logfile, _ := rootCmd.PersistentFlags().GetString("log-file")
|
|
err := logger.InitWithLogLevel(loglevel, logfile)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("failed to initialize logger")
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func handleResponseError(res *http.Response, host, query string, err error) {
|
|
if err != nil {
|
|
log.Error().Err(err).
|
|
Str("host", host).
|
|
Str("query", query).
|
|
Msg("failed to make request")
|
|
os.Exit(1)
|
|
}
|
|
if res.StatusCode != http.StatusOK {
|
|
log.Error().
|
|
Any("status", map[string]any{
|
|
"code": res.StatusCode,
|
|
"message": res.Status,
|
|
}).
|
|
Str("host", host).
|
|
Msg("response returned bad status")
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// helper to write downloaded files
|
|
func writeFiles(path string, body []byte) {
|
|
var err = os.WriteFile(path, body, 0o755)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("failed to write file(s) from download")
|
|
os.Exit(1)
|
|
}
|
|
}
|