Add: CLI option for generating config

This commit is contained in:
kayos@tcp.direct 2021-09-18 05:22:15 -07:00
parent f0eda75b35
commit eeaa181910
2 changed files with 76 additions and 68 deletions

View File

@ -15,14 +15,16 @@ const (
// Version roughly represents the applications current version. // Version roughly represents the applications current version.
Version = "0.3" Version = "0.3"
// Title is the name of the application used throughout the configuration process. // Title is the name of the application used throughout the configuration process.
Title = "HellPot" Title = "HellPot"
) )
var ( var (
// BannerOnly when toggled causes HellPot to only print the banner and version then exit. // BannerOnly when toggled causes HellPot to only print the banner and version then exit.
BannerOnly = false BannerOnly = false
// ConfigGen when toggled causes HellPot to write its default config to the cwd and then exit. // GenConfig when toggled causes HellPot to write its default config to the cwd and then exit.
ConfigGen = false GenConfig = false
// NoColor when true will disable the banner and any colored console output.
NoColor bool
) )
// "http" // "http"
@ -36,7 +38,7 @@ var (
Paths []string Paths []string
// UseUnixSocket when toggled disables the TCP listener and listens on the given UnixSocketPath. // UseUnixSocket when toggled disables the TCP listener and listens on the given UnixSocketPath.
UseUnixSocket bool UseUnixSocket bool
// UnixSocketPath is the path of the unix socket used when UseUnixSocket is toggled. // UnixSocketPath is the path of the unix socket used when UseUnixSocket is toggled.
UnixSocketPath = "" UnixSocketPath = ""
) )
@ -61,21 +63,15 @@ var (
var ( var (
f *os.File f *os.File
err error err error
)
NoColorForce = false var (
NoColor bool noColorForce = false
customconfig = false customconfig = false
home string home string
configLocations []string configLocations []string
) )
/*
Opt represents our program options.
nitially the values that are defined in Opt will be used to define details.
Beyond that, default values will be replaced by options from our config file.
*/
var Opt map[string]map[string]interface{}
var ( var (
// Debug is our global debug toggle // Debug is our global debug toggle
Debug bool Debug bool
@ -89,10 +85,37 @@ func init() {
panic(err) panic(err)
} }
prefConfigLocation = home + "/.config/" + Title prefConfigLocation = home + "/.config/" + Title
Opt = make(map[string]map[string]interface{})
snek = viper.New() snek = viper.New()
} }
func writeConfig() {
if runtime.GOOS != "windows" {
if _, err := os.Stat(prefConfigLocation); os.IsNotExist(err) {
if err = os.Mkdir(prefConfigLocation, 0755); err != nil {
println("error writing new config: " + err.Error())
}
}
newconfig := prefConfigLocation + "/" + "config.toml"
if err = snek.SafeWriteConfigAs(newconfig); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
Filename = newconfig
return
}
newconfig := "hellpot-config"
snek.SetConfigName(newconfig)
if err = snek.MergeInConfig(); err != nil {
if err = snek.SafeWriteConfigAs(newconfig + ".toml"); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
}
Filename = newconfig
}
// Init will initialize our toml configuration engine and define our default configuration values which can be written to a new configuration file if desired // Init will initialize our toml configuration engine and define our default configuration values which can be written to a new configuration file if desired
func Init() { func Init() {
snek.SetConfigType("toml") snek.SetConfigType("toml")
@ -113,34 +136,8 @@ func Init() {
snek.AddConfigPath(loc) snek.AddConfigPath(loc)
} }
if err = snek.MergeInConfig(); err != nil && runtime.GOOS != "windows" { if err = snek.MergeInConfig(); err != nil {
if _, err := os.Stat(prefConfigLocation); os.IsNotExist(err) { writeConfig()
if err = os.Mkdir(prefConfigLocation, 0755); err != nil {
println("error writing new config: " + err.Error())
}
}
newconfig := prefConfigLocation + "/" + "config.toml"
if err = snek.SafeWriteConfigAs(newconfig); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
Filename = newconfig
}
if runtime.GOOS == "windows" {
newconfig := "hellpot-config"
snek.SetConfigName(newconfig)
if err = snek.MergeInConfig(); err != nil {
if err = snek.SafeWriteConfigAs(newconfig + ".toml"); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
}
Filename = newconfig
} }
if len(Filename) < 1 { if len(Filename) < 1 {
@ -162,6 +159,8 @@ func setDefaults() {
defNoColor = true defNoColor = true
} }
Opt := make(map[string]map[string]interface{})
Opt["logger"] = map[string]interface{}{ Opt["logger"] = map[string]interface{}{
"debug": true, "debug": true,
"directory": deflogdir, "directory": deflogdir,
@ -190,7 +189,7 @@ func setDefaults() {
snek.SetDefault(def, Opt[def]) snek.SetDefault(def, Opt[def])
} }
if ConfigGen { if GenConfig {
if err = snek.SafeWriteConfigAs("./config.toml"); err != nil { if err = snek.SafeWriteConfigAs("./config.toml"); err != nil {
fmt.Println(err.Error()) fmt.Println(err.Error())
os.Exit(1) os.Exit(1)
@ -230,7 +229,7 @@ func loadCustomConfig(path string) {
} }
func printUsage() { func printUsage() {
println("\n"+Title+" v"+Version+" Usage\n") println("\n" + Title + " v" + Version + " Usage\n")
println("-c <config.toml> - Specify config file") println("-c <config.toml> - Specify config file")
println("--nocolor - disable color and banner ") println("--nocolor - disable color and banner ")
println("--banner - show banner + version and exit") println("--banner - show banner + version and exit")
@ -245,9 +244,9 @@ func argParse() {
case "-h": case "-h":
printUsage() printUsage()
case "--genconfig": case "--genconfig":
ConfigGen = true GenConfig = true
case "--nocolor": case "--nocolor":
NoColorForce = true noColorForce = true
case "--banner": case "--banner":
BannerOnly = true BannerOnly = true
case "--config": case "--config":
@ -263,31 +262,39 @@ func argParse() {
} }
} }
func associate() { func bl(key string) bool {
var newOpt map[string]map[string]interface{} return snek.GetBool(key)
newOpt = Opt }
for category, opt := range Opt { func st(key string) string {
for optname, value := range opt { return snek.GetString(key)
if snek.IsSet(category + "." + optname) { }
newOpt[category][optname] = value func sl(key string) []string {
} return snek.GetStringSlice(key)
} }
} func it(key string) int {
return snek.GetInt(key)
}
Opt = newOpt
Debug = snek.GetBool("logger.debug") func associate() {
logDir = snek.GetString("logger.directory") BindAddr = st("http.bind_addr")
BindAddr = snek.GetString("http.bind_addr") BindPort = st("http.bind_port")
BindPort = snek.GetString("http.bind_port") Paths = sl("http.paths")
Paths = snek.GetStringSlice("http.paths") UseUnixSocket = bl("http.use_unix_socket")
UseUnixSocket = snek.GetBool("http.use_unix_socket") //
FakeServerName = snek.GetString("deception.server_name") Debug = bl("logger.debug")
RestrictConcurrency = snek.GetBool("performance.restrict_concurrency") logDir = st("logger.directory")
MaxWorkers = snek.GetInt("performance.max_workers") NoColor = bl("logger.nocolor")
NoColor = snek.GetBool("logger.nocolor") //
if NoColorForce { FakeServerName = st("deception.server_name")
//
RestrictConcurrency = bl("performance.restrict_concurrency")
MaxWorkers = it("performance.max_workers")
if noColorForce {
NoColor = true NoColor = true
} }
if UseUnixSocket { if UseUnixSocket {
UnixSocketPath = snek.GetString("http.unix_socket_path") UnixSocketPath = snek.GetString("http.unix_socket_path")
} }

1
go.mod
View File

@ -13,6 +13,7 @@ require (
github.com/spf13/viper v1.7.1 github.com/spf13/viper v1.7.1
github.com/stretchr/testify v1.6.1 // indirect github.com/stretchr/testify v1.6.1 // indirect
github.com/valyala/fasthttp v1.30.0 github.com/valyala/fasthttp v1.30.0
golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a
gopkg.in/ini.v1 v1.53.0 // indirect gopkg.in/ini.v1 v1.53.0 // indirect
gopkg.in/yaml.v2 v2.3.0 // indirect gopkg.in/yaml.v2 v2.3.0 // indirect
) )