Add optional -g/--grimoire flag to choose source file other than Niet…

README.md

@@ -16,7 +16,7 @@ Clients (hopefully bots) that disregard `robots.txt` and connect to your instanc
 
 HellPot will send an infinite stream of data that is *just close enough* to being a real website that they might just stick around until their soul is ripped apart and they cease to exist.
 
-Under the hood of this eternal suffering is a markov engine that chucks bits and pieces of [The Birth of Tragedy (Hellenism and Pessimism)](https://www.gutenberg.org/files/51356/51356-h/51356-h.htm) by Friedrich Nietzsche at the client using [fasthttp](https://github.com/valyala/fasthttp).
+Under the hood of this eternal suffering is a markov engine that chucks bits and pieces of [The Birth of Tragedy (Hellenism and Pessimism)](https://www.gutenberg.org/files/51356/51356-h/51356-h.htm) by Friedrich Nietzsche at the client using [fasthttp](https://github.com/valyala/fasthttp), or optionally your (least?) favorite text using the `-b`/`--book` flag.
 
 ## Building From Source
 

cmd/HellPot/HellPot.go

@@ -38,7 +38,7 @@ func init() {
 
 	extra.Banner()
 
-	log.Info().Str("caller", "config").Str("file", config.Filename).Msg(config.Filename)
+	log.Info().Str("caller", "config").Str("file", config.ConfigFilename).Msg(config.ConfigFilename)
 	log.Info().Str("caller", "logger").Msg(config.CurrentLogFile)
 	log.Debug().Str("caller", "logger").Msg("debug enabled")
 	log.Trace().Str("caller", "logger").Msg("trace enabled")

heffalump/heffalump.go

@@ -14,8 +14,7 @@ import (
 
 var log = config.GetLogger()
 
-// DefaultHeffalump represents a Heffalump type
-var DefaultHeffalump *Heffalump
+const DefaultBuffSize = 100 * 1 << 10
 
 // Heffalump represents our buffer pool and markov map from Heffalump
 type Heffalump struct {
@@ -36,6 +35,12 @@ func NewHeffalump(mm MarkovMap, buffsize int) *Heffalump {
 	}
 }
 
+// NewDefaultHeffalump instantiates a new default Heffalump from a MarkovMap created using
+// using the default source text.
+func NewDefaultHeffalump() *Heffalump {
+	return NewHeffalump(NewDefaultMarkovMap(), DefaultBuffSize)
+}
+
 // WriteHell writes markov chain heffalump hell to the provided io.Writer
 func (h *Heffalump) WriteHell(bw *bufio.Writer) (int64, error) {
 	var n int64

heffalump/markov.go

@@ -11,19 +11,17 @@ import (
 	"git.tcp.direct/kayos/common/squish"
 )
 
-var DefaultMarkovMap MarkovMap
-
-func init() {
-	// DefaultMarkovMap is a Markov chain based on src.
+// NewDefaultMarkovMap creates a new MarkovMap from the default source text.
+func NewDefaultMarkovMap() MarkovMap {
 	src, err := squish.UnpackStr(srcGz)
 	if err != nil {
 		panic(err)
 	}
 	if len(src) < 1 {
 		panic("failed to unpack source")
 	}
-	DefaultMarkovMap = MakeMarkovMap(strings.NewReader(src))
-	DefaultHeffalump = NewHeffalump(DefaultMarkovMap, 100*1<<10)
+
+	return MakeMarkovMap(strings.NewReader(src))
 }
 
 // ScanHTML is a basic split function for a Scanner that returns each

internal/config/arguments.go

@@ -22,14 +22,21 @@ func argParse() {
 			continue
 		}
 		switch arg {
-		case "-h":
+		case "-h", "--help":
 			CLI.printUsage()
 		case "-c", "--config":
 			if len(os.Args) < i+2 {
 				println("missing config file after -c/--config")
 				os.Exit(1)
 			}
 			loadCustomConfig(os.Args[i+1])
+		case "-b", "--book":
+			if len(os.Args) < i+2 {
+				println("missing book file after -b/--book")
+				os.Exit(1)
+			}
+			BookFilename = os.Args[i+1]
+			UseCustomHeffalump = true
 		default:
 			continue
 		}

internal/config/config.go

@@ -29,8 +29,12 @@ var (
 	Trace bool
 	// Debug is the value of our debug (verbose) on/off toggle as per the current configuration.
 	Debug bool
-	// Filename returns the current location of our toml config file.
-	Filename string
+	// ConfigFilename returns the current location of our toml config file.
+	ConfigFilename string
+	// UseCustomHeffalump decides if a custom Heffalump is to be used
+	UseCustomHeffalump = false
+	// BookFilename returns the current location of a possible book file
+	BookFilename string
 )
 
 func writeConfig() {
@@ -40,9 +44,9 @@ func writeConfig() {
 			os.Exit(1)
 		}
 	}
-	Filename = prefConfigLocation + "/" + "config.toml"
-	if err := snek.SafeWriteConfigAs(Filename); err != nil {
-		fmt.Println("Failed to write new configuration file to '" + Filename + "': " + err.Error())
+	ConfigFilename = prefConfigLocation + "/" + "config.toml"
+	if err := snek.SafeWriteConfigAs(ConfigFilename); err != nil {
+		fmt.Println("Failed to write new configuration file to '" + ConfigFilename + "': " + err.Error())
 		os.Exit(1)
 	}
 }
@@ -71,8 +75,8 @@ func Init() {
 		writeConfig()
 	}
 
-	if len(Filename) < 1 {
-		Filename = snek.ConfigFileUsed()
+	if len(ConfigFilename) < 1 {
+		ConfigFilename = snek.ConfigFileUsed()
 	}
 
 	snek.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
@@ -102,9 +106,9 @@ func loadCustomConfig(path string) {
 		os.Exit(1)
 	}
 
-	Filename, err = filepath.Abs(path)
-	if len(Filename) < 1 || err != nil {
-		Filename = path
+	ConfigFilename, err = filepath.Abs(path)
+	if len(ConfigFilename) < 1 || err != nil {
+		ConfigFilename = path
 	}
 
 	defer func(f *os.File) {

internal/config/help.go

@@ -18,11 +18,12 @@ var CLI = help{
 	title:   Title,
 	version: Version,
 	usage: map[int][]string{
-		0: {0: "--config", 1: "<file>", 2: "Specify config file"},
+		0: {0: "-c, --config", 1: "<file>", 2: "Specify config file"},
 		1: {0: "--nocolor", 1: "disable color and banner"},
 		2: {0: "--banner", 1: "show banner + version and exit"},
 		3: {0: "--genconfig", 1: "write default config to " + Title + ".toml then exit"},
-		4: {0: "--help", 1: "show this help and exit"},
+		4: {0: "-b, --book", 1: "<file>", 2: "Specify a custom file used for text generation"},
+		5: {0: "-h, --help", 1: "show this help and exit"},
 	},
 	out: os.Stdout,
 }

internal/http/router.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"fmt"
 	"net/http"
+	"os"
 	"runtime"
 	"strings"
 	"time"
@@ -16,7 +17,10 @@ import (
 	"github.com/yunginnanet/HellPot/internal/config"
 )
 
-var log *zerolog.Logger
+var (
+	log              *zerolog.Logger
+	hellpotHeffalump *heffalump.Heffalump
+)
 
 func getRealRemote(ctx *fasthttp.RequestCtx) string {
 	xrealip := string(ctx.Request.Header.Peek(config.HeaderName))
@@ -61,7 +65,7 @@ func hellPot(ctx *fasthttp.RequestCtx) {
 		var wn int64
 
 		for {
-			wn, err = heffalump.DefaultHeffalump.WriteHell(bw)
+			wn, err = hellpotHeffalump.WriteHell(bw)
 			n += wn
 			if err != nil {
 				slog.Trace().Err(err).Msg("END_ON_ERR")
@@ -74,7 +78,6 @@ func hellPot(ctx *fasthttp.RequestCtx) {
 			Dur("DURATION", time.Since(s)).
 			Msg("FINISH")
 	})
-
 }
 
 func getSrv(r *router.Router) fasthttp.Server {
@@ -120,6 +123,27 @@ func getSrv(r *router.Router) fasthttp.Server {
 // Serve starts our HTTP server and request router
 func Serve() error {
 	log = config.GetLogger()
+
+	if config.UseCustomHeffalump {
+		content, err := os.ReadFile(config.BookFilename)
+		if err != nil {
+			panic(err)
+		}
+		// Wasteful, but only done once at startup
+		src := string(content)
+		log.Info().Msgf("Using custom book file '%s'", config.BookFilename)
+
+		if len(src) < 1 {
+			panic("book file was empty!")
+		}
+
+		markovMap := heffalump.MakeMarkovMap(strings.NewReader(src))
+		hellpotHeffalump = heffalump.NewHeffalump(markovMap, heffalump.DefaultBuffSize)
+	} else {
+		log.Info().Msg("Using default source text")
+		hellpotHeffalump = heffalump.NewDefaultHeffalump()
+	}
+
 	l := config.HTTPBind + ":" + config.HTTPPort
 
 	r := router.New()