HellPot/config/config.go

251 lines
5.6 KiB
Go
Raw Normal View History

package config
import (
"bytes"
"fmt"
2021-11-14 22:06:15 +00:00
"io"
"os"
2021-09-15 14:04:16 +00:00
"runtime"
"strconv"
"github.com/rs/zerolog"
"github.com/spf13/viper"
)
func init() {
if home, err = os.UserHomeDir(); err != nil {
panic(err)
}
2021-09-15 14:04:16 +00:00
prefConfigLocation = home + "/.config/" + Title
snek = viper.New()
}
2021-09-18 12:22:15 +00:00
func writeConfig() {
2021-10-17 17:09:28 +00:00
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)
2021-09-18 12:22:15 +00:00
}
}
return
}
2021-10-17 17:09:28 +00:00
if _, err := os.Stat(prefConfigLocation); os.IsNotExist(err) {
2021-11-14 22:06:15 +00:00
if err = os.MkdirAll(prefConfigLocation, 0o755); err != nil {
2021-10-17 17:09:28 +00:00
println("error writing new config: " + err.Error())
2021-09-18 12:22:15 +00:00
os.Exit(1)
}
}
2021-10-17 17:09:28 +00:00
newconfig := prefConfigLocation + "/" + "config.toml"
if err = snek.SafeWriteConfigAs(newconfig); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
2021-09-18 12:22:15 +00:00
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
func Init() {
snek.SetConfigType("toml")
snek.SetConfigName("config")
argParse()
if customconfig {
2021-09-21 13:02:59 +00:00
associateExportedVariables()
return
}
2021-09-21 13:02:59 +00:00
setConfigFileLocations()
setDefaults()
for _, loc := range configLocations {
snek.AddConfigPath(loc)
}
2021-09-18 12:22:15 +00:00
if err = snek.MergeInConfig(); err != nil {
writeConfig()
2021-09-15 14:04:16 +00:00
}
if len(Filename) < 1 {
Filename = snek.ConfigFileUsed()
}
2021-09-21 13:02:59 +00:00
associateExportedVariables()
}
func setDefaults() {
2021-09-15 14:04:16 +00:00
var (
2021-09-21 13:02:59 +00:00
configSections = []string{"logger", "http", "performance", "deception", "ssh"}
2021-09-15 14:04:16 +00:00
deflogdir = home + "/.config/" + Title + "/logs/"
defNoColor = false
)
if runtime.GOOS == "windows" {
deflogdir = "logs/"
defNoColor = true
}
2021-09-18 12:22:15 +00:00
Opt := make(map[string]map[string]interface{})
Opt["logger"] = map[string]interface{}{
2021-09-15 14:04:16 +00:00
"debug": true,
"directory": deflogdir,
"nocolor": defNoColor,
"use_date_filename": true,
}
Opt["http"] = map[string]interface{}{
"use_unix_socket": false,
"unix_socket_path": "/var/run/hellpot",
"unix_socket_permissions": "0666",
"bind_addr": "127.0.0.1",
"bind_port": "8080",
"paths": []string{
"wp-login.php",
"wp-login",
},
}
2021-09-15 16:13:16 +00:00
Opt["performance"] = map[string]interface{}{
"restrict_concurrency": false,
2021-09-15 19:17:32 +00:00
"max_workers": 256,
2021-09-15 16:13:16 +00:00
}
2021-09-15 19:55:48 +00:00
Opt["deception"] = map[string]interface{}{
2021-09-15 16:13:16 +00:00
"server_name": "nginx",
}
for _, def := range configSections {
snek.SetDefault(def, Opt[def])
}
2021-09-16 07:51:16 +00:00
2021-09-18 12:22:15 +00:00
if GenConfig {
2021-09-16 07:51:16 +00:00
if err = snek.SafeWriteConfigAs("./config.toml"); err != nil {
fmt.Println(err.Error())
os.Exit(1)
}
os.Exit(0)
}
}
2021-09-21 13:02:59 +00:00
func setConfigFileLocations() {
configLocations = append(configLocations, "./")
2021-09-15 14:04:16 +00:00
if runtime.GOOS != "windows" {
2021-11-14 22:06:15 +00:00
configLocations = append(configLocations,
prefConfigLocation, "/etc/"+Title+"/", "../", "../../")
2021-09-15 14:04:16 +00:00
}
}
func loadCustomConfig(path string) {
if f, err = os.Open(path); err != nil {
println("Error opening specified config file: " + path)
panic("config file open fatal error: " + err.Error())
}
2021-11-14 22:06:15 +00:00
buf, err := io.ReadAll(f)
err2 := snek.ReadConfig(bytes.NewBuffer(buf))
switch {
case err != nil:
fmt.Println("config file read fatal error: ", err.Error())
case err2 != nil:
fmt.Println("config file read fatal error: ", err2.Error())
default:
break
}
customconfig = true
}
2021-09-16 07:51:16 +00:00
func printUsage() {
2021-09-18 12:22:15 +00:00
println("\n" + Title + " v" + Version + " Usage\n")
2021-09-16 07:51:16 +00:00
println("-c <config.toml> - Specify config file")
println("--nocolor - disable color and banner ")
println("--banner - show banner + version and exit")
println("--genconfig - write default config to 'default.toml' then exit")
os.Exit(0)
}
// TODO: should probably just make a proper CLI with flags or something
func argParse() {
for i, arg := range os.Args {
switch arg {
2021-09-15 14:04:16 +00:00
case "-h":
2021-09-16 07:51:16 +00:00
printUsage()
case "--genconfig":
2021-09-18 12:22:15 +00:00
GenConfig = true
2021-09-15 19:33:42 +00:00
case "--nocolor":
2021-09-18 12:22:15 +00:00
noColorForce = true
2021-09-15 19:17:32 +00:00
case "--banner":
BannerOnly = true
2021-11-14 22:06:15 +00:00
case "-c", "--config":
if len(os.Args) <= i-1 {
panic("syntax error! expected file after -c")
}
loadCustomConfig(os.Args[i+1])
default:
continue
}
}
}
2021-09-21 13:02:59 +00:00
func processOpts() {
// string options and their exported variables
stringOpt := map[string]*string{
"http.bind_addr": &HTTPBind,
"http.bind_port": &HTTPPort,
"logger.directory": &logDir,
"deception.server_name": &FakeServerName,
}
// string slice options and their exported variables
strSliceOpt := map[string]*[]string{
"http.paths": &Paths,
}
// bool options and their exported variables
2021-10-17 17:09:28 +00:00
boolOpt := map[string]*bool{
"http.use_unix_socket": &UseUnixSocket,
"logger.debug": &Debug,
2021-09-21 13:02:59 +00:00
"performance.restrict_concurrency": &RestrictConcurrency,
2021-10-17 17:09:28 +00:00
"logger.nocolor": &NoColor,
2021-09-21 13:02:59 +00:00
}
// integer options and their exported variables
2021-10-17 17:09:28 +00:00
intOpt := map[string]*int{
2021-09-21 13:02:59 +00:00
"performance.max_workers": &MaxWorkers,
}
2021-09-21 13:02:59 +00:00
for key, opt := range stringOpt {
*opt = snek.GetString(key)
}
for key, opt := range strSliceOpt {
*opt = snek.GetStringSlice(key)
}
for key, opt := range boolOpt {
*opt = snek.GetBool(key)
}
for key, opt := range intOpt {
*opt = snek.GetInt(key)
}
}
func associateExportedVariables() {
processOpts()
2021-09-18 12:22:15 +00:00
if noColorForce {
2021-09-15 19:33:42 +00:00
NoColor = true
}
2021-09-18 12:22:15 +00:00
2021-09-15 16:13:16 +00:00
if UseUnixSocket {
UnixSocketPath = snek.GetString("http.unix_socket_path")
parsedPermissions, err := strconv.ParseUint(snek.GetString("http.unix_socket_permissions"), 8, 32)
if err == nil {
UnixSocketPermissions = uint32(parsedPermissions)
}
2021-09-15 16:13:16 +00:00
}
if Debug {
zerolog.SetGlobalLevel(zerolog.DebugLevel)
}
}