Merge pull request #131 from ginger51011/feat/add-file-flag

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 you may synchronize HellPot with your nightmares by using the `-g`/`--grimoire` flag
 
 ## Building From Source
 

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 "-g", "--grimoire":
+			if len(os.Args) < i+2 {
+				println("missing source of suffering file after -g/--grimoire")
+				os.Exit(1)
+			}
+			Grimoire = os.Args[i+1]
+			UseCustomHeffalump = true
 		default:
 			continue
 		}

internal/config/config.go

@@ -31,6 +31,10 @@ var (
 	Debug bool
 	// Filename returns the current location of our toml config file.
 	Filename string
+	// UseCustomHeffalump decides if a custom Heffalump is to be used
+	UseCustomHeffalump = false
+	// Grimoire returns the current location of a possible source of suffering file
+	Grimoire string
 )
 
 func writeConfig() {

internal/config/defaults.go

@@ -88,7 +88,7 @@ func gen(memfs afero.Fs) {
 		println(err.Error())
 		os.Exit(1)
 	}
-	print("default configuration successfully written to " + target)
+	println("default configuration successfully written to " + target)
 	os.Exit(0)
 }
 

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: "-g, --grimoire", 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,28 @@ func getSrv(r *router.Router) fasthttp.Server {
 // Serve starts our HTTP server and request router
 func Serve() error {
 	log = config.GetLogger()
+
+	switch config.UseCustomHeffalump {
+	case true:
+		content, err := os.ReadFile(config.Grimoire)
+		if err != nil {
+			panic(err)
+		}
+		// Wasteful, but only done once at startup
+		src := string(content)
+		log.Info().Msgf("Using custom grimoire file '%s'", config.Grimoire)
+
+		if len(src) < 1 {
+			panic("grimoire file was empty!")
+		}
+
+		markovMap := heffalump.MakeMarkovMap(strings.NewReader(src))
+		hellpotHeffalump = heffalump.NewHeffalump(markovMap, heffalump.DefaultBuffSize)
+	default:
+		log.Info().Msg("Using default source text")
+		hellpotHeffalump = heffalump.NewDefaultHeffalump()
+	}
+
 	l := config.HTTPBind + ":" + config.HTTPPort
 
 	r := router.New()