mtbf/source.go

package main

import (
	"bufio"
	"os"
	"strconv"
	"strings"
	"sync"
)

type Source struct {
	name, plainParmName, filesParmName string

	validator func(item string) (string, error)

	plain     []string
	files     []*os.File
	fileNames []string

	contents []string

	fetchMutex sync.Mutex
}

// both -1: exhausted
// both 0: not started yet
type SourcePos struct {
	plainIdx   int
	contentIdx int
}

func (pos *SourcePos) String() string {
	return "P" + strconv.Itoa(pos.plainIdx) + "/C" + strconv.Itoa(pos.contentIdx)
}

func (pos *SourcePos) Exhausted() bool {
	return pos.plainIdx == -1 && pos.contentIdx == -1
}

func ipValidator(item string) (res string, err error) {
	return item, nil
}

func passwordValidator(item string) (res string, err error) {
	if CfgGetSwitch("no-password-trim") {
		return item, nil
	} else {
		return strings.TrimSpace(item), nil
	}
}

var SrcIP Source = Source{name: "ip", plainParmName: "ip", filesParmName: "ip-file"}
var SrcLogin Source = Source{name: "login", plainParmName: "login", filesParmName: "login-file"}
var SrcPassword Source = Source{name: "password", plainParmName: "password",
	filesParmName: "password-file", validator: passwordValidator}

func (src *Source) validate(item string) (res string, err error) {
	if src.validator != nil {
		res, err := src.validator(item)
		if err != nil {
			log("src", 1, "error validating %v \"%v\": %v", src.name, item, err.Error())
		}
		return res, err
	} else {
		return item, nil
	}
}

func (src *Source) parsePlain() {
	if src.plain == nil {
		src.plain = []string{}
	}

	for _, plain := range CfgGetStringSlice(src.plainParmName) {
		var err error
		plain, err = src.validate(plain)
		if err != nil {
			continue
		}

		src.plain = append(src.plain, plain)
	}

	if len(src.plain) > 0 {
		log("src", 1, "parsed %v %v items", len(src.plain), src.name)
	}
}

func (src *Source) openFiles() {
	if src.files == nil {
		src.files = []*os.File{}
	}

	if src.fileNames == nil {
		src.fileNames = []string{}
	}

	fileNames := CfgGetStringSlice(src.filesParmName)
	for _, fileName := range fileNames {
		f, err := os.Open(fileName)
		if err != nil {
			fail("error opening source file \"%v\": %v", fileName, err.Error())
		}

		src.files = append(src.files, f)
		src.fileNames = append(src.fileNames, fileName)
	}

	if len(src.files) > 0 {
		log("src", 1, "opened %v %v files", len(src.files), src.name)
	}
}

// this parses all source files
func (src *Source) parseFiles() {
	if src.contents == nil {
		src.contents = []string{}
	}

	for i, file := range src.files {
		fileName := src.fileNames[i]
		log("src", 1, "parsing %v", fileName)
		thisTotal := 0

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			text := scanner.Text()
			if text == "" {
				continue
			}

			value, err := src.validate(text)
			if err != nil {
				continue
			}

			src.contents = append(src.contents, value)
			thisTotal++
		}

		scannerErr := scanner.Err()
		failIf(scannerErr != nil, "error reading source file \"%v\": %v", fileName, scannerErr)
		log("src", 1, "ok: parsed %v, got %v contents, %v total", fileName, thisTotal, len(src.contents))
	}
}

func (src *Source) failIfEmpty() {
	failIf(len(src.contents)+len(src.plain) == 0, "no %vs defined: check %v and %v parameters", src.name, src.plainParmName, src.filesParmName)
}

func (src *Source) reportLoaded() {
	log("src", 0, "loaded %vs: %v items from commandline and %v items from files", src.name, len(src.plain), len(src.contents))
}

func (src *Source) load(wg *sync.WaitGroup) {
	if wg != nil {
		defer wg.Done()
	}

	if src.name == "password" && CfgGetSwitch("add-empty-password") {
		if src.plain == nil {
			src.plain = []string{}
		}
		src.plain = append(src.plain, "")
	}

	src.parsePlain()
	src.openFiles()
	defer src.closeFiles()

	src.parseFiles()
	src.failIfEmpty()
}

func (src *Source) closeFiles() {
	l := len(src.files)
	for _, file := range src.files {
		if file != nil {
			file.Close()
		}
	}

	src.files = []*os.File{}
	src.fileNames = []string{}
	log("src", 1, "closed all %v %v files", l, src.name)
}

func (src *Source) fetchFromSlice(name string, idx *int, slice []string, inc bool) (res string, empty bool) {
	if *idx == -1 {
		// exhausted
		log("src", 3, "fetch %v from %v: idx is -1, return empty", src.name, name)
		return "", true
	}

	if *idx >= len(slice) {
		log("src", 3, "fetch %v from %v: idx >= slice length (%v >= %v), marking as exhausted, return empty", src.name, name, *idx, len(slice))
		*idx = -1
		return "", true
	}

	res = slice[*idx]
	log("src", 3, "fetch %v from %v: ok, got %v at idx %v", src.name, name, res, *idx)

	if inc {
		*idx = *idx + 1
		log("src", 3, "fetch %v from %v: incrementing idx to %v", src.name, name, *idx)
	}

	return res, false
}

// retrieves an item from source
// increments pos
func (src *Source) FetchOne(pos *SourcePos, inc bool) (res string, empty bool) {
	src.fetchMutex.Lock()
	defer src.fetchMutex.Unlock()

	if CfgGetSwitch("file-contents-first") {
		res, empty = src.fetchFromSlice("contents", &pos.contentIdx, src.contents, inc)
		if empty {
			res, empty = src.fetchFromSlice("plain", &pos.plainIdx, src.plain, inc)
		}
	} else {
		res, empty = src.fetchFromSlice("plain", &pos.plainIdx, src.plain, inc)
		if empty {
			res, empty = src.fetchFromSlice("contents", &pos.contentIdx, src.contents, inc)
		}
	}

	logIf(empty, "src", 2, "exhausted source %v for pos %v", src.name, pos.String())
	return res, empty
}

func (pos *SourcePos) Reset() {
	pos.plainIdx = 0
	pos.contentIdx = 0
	log("src", 3, "resetting source pos")
}

func LoadSources() {
	log("src", 1, "loading sources")

	var wg sync.WaitGroup
	wg.Add(3)
	go SrcIP.load(&wg)
	go SrcLogin.load(&wg)
	go SrcPassword.load(&wg)
	wg.Wait()

	SrcIP.reportLoaded()
	SrcLogin.reportLoaded()
	SrcPassword.reportLoaded()

	ParseEndpoints(SrcIP.plain)

	// TODO: dynamic loading of contents
	// currently this just dumps everything to a slice
	ParseEndpoints(SrcIP.contents)

	log("src", 1, "ok: finished loading sources")
}

func CloseSources() {
	log("src", 1, "closing sources")
	SrcIP.closeFiles()
	SrcLogin.closeFiles()
	SrcPassword.closeFiles()
	log("src", 1, "ok: finished closing sources")
}

func init() {
	CfgRegister("ip", []string{}, "IPs or subnets in CIDR notation")
	CfgRegister("ip-file", []string{}, "paths to files with IPs or subnets in CIDR notation (one entry per line)")
	CfgRegister("login", []string{}, "one or more logins")
	CfgRegister("login-file", []string{}, "paths to files with logins (one entry per line)")
	CfgRegister("password", []string{}, "one or more passwords")
	CfgRegister("password-file", []string{}, "paths to files with passwords (one entry per line)")

	CfgRegisterSwitch("add-empty-password", "insert an empty password to the password list")
	//CfgRegisterSwitch("no-source-validation", "do not attempt to validate and count lines in source files")
	CfgRegisterSwitch("no-password-trim", "preserve leading and trailing spaces in passwords")
	CfgRegisterSwitch("logins-first", "increment logins before passwords")
	CfgRegisterSwitch("file-contents-first", "try to go through source files first, defer commandline args for later")

	//
	//CfgRegisterSwitch("ignore-network-names", "always skip non-IPv4/non-IPv6 entries, do not attempt to resolve them as domain names")
	//
}