mtbf/source.go

package main

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

type Source struct {
	name string // name of this source

	plain    []string // sources from commandline
	contents []string // sources from files

	files         []*os.File // file pointers
	fileNames     []string   // file names
	plainParmName string     // name of "plain" commandline parameter
	filesParmName string     // name of "files" commandline parameter

	transform func(item string) (string, error) // optional transformation function

	fetchMutex sync.Mutex // sync mutex
}

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

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

// Exhausted checks if a SourcePos can no longer produce any sources.
func (pos *SourcePos) Exhausted() bool {
	return pos.plainIdx == -1 && pos.contentIdx == -1
}

// Reset moves a SourcePos to its starting position.
func (pos *SourcePos) Reset() {
	pos.plainIdx = 0
	pos.contentIdx = 0
	log("src", 3, "resetting source pos")
}

// passwordTransform is a transformation function for a password.
func passwordTransform(item string) (res string, err error) {
	if getParamSwitch("no-password-trim") {
		return item, nil
	} else {
		return strings.TrimSpace(item), nil
	}
}

// SrcIP, SrcLogin and SrcPassword represent different sources.
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", transform: passwordTransform}

// String converts a Source to its string representation.
func (src *Source) String() string {
	return src.name
}

// ValidateAndTransformItem attempts to validate a source item
// and performs transformations, if any.
func (src *Source) ValidateAndTransformItem(item string) (res string, err error) {
	if src.transform != nil {
		res, err := src.transform(item)
		if err != nil {
			log("src", 1, "error validating %v \"%v\": %v", src, item, err.Error())
			res = ""
		}

		return res, err
	} else {
		return item, nil
	}
}

// FetchFromSlice retrieves an item from a string slice and optionally increments its current position.
func (src *Source) FetchFromSlice(name string, idx *int, slice []string, inc bool) (res string, empty bool) {
	if *idx == -1 { // exhausted
		log("src", 5, "fetch %v from %v: idx is -1, return empty", src, name)
		return "", true
	}

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

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

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

	return res, false
}

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

	if getParamSwitch("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, pos.String())
	return res, empty
}

// ParsePlain parses commandline parameters for a Source.
func (src *Source) ParsePlain() {
	for _, plain := range getParamStringSlice(src.plainParmName) {
		plain, err := src.ValidateAndTransformItem(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)
	}
}

// OpenFiles opens all files for a Source.
func (src *Source) OpenFiles() {
	fileNames := getParamStringSlice(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)
	}
}

// ParseFiles parses all files for a Source.
func (src *Source) ParseFiles() {
	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.ValidateAndTransformItem(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))
	}
}

// FailIfEmpty throws an exception if a source is empty.
func (src *Source) FailIfEmpty() {
	failIf(len(src.contents)+len(src.plain) == 0, "no %vs defined: check %v and %v parameters", src, src.plainParmName, src.filesParmName)
}

// ReportLoaded prints a console message about the number of loaded items for a Source.
func (src *Source) ReportLoaded() {
	log("src", 0, "loaded %vs: %v items from commandline and %v items from files", src, len(src.plain), len(src.contents))
}

// LoadSource fills a Source with data (from commandline and from files).
func (src *Source) LoadSource(wg *sync.WaitGroup) {
	if wg != nil {
		defer wg.Done()
	}

	if src.name == "password" && getParamSwitch("add-empty-password") {
		src.plain = append(src.plain, "")
	}

	src.ParsePlain()

	src.OpenFiles()
	defer src.CloseSource()

	src.ParseFiles()
	src.FailIfEmpty()
}

// CloseSource closes all files for a Source.
func (src *Source) CloseSource() {
	l := len(src.files)
	for _, file := range src.files {
		if file != nil {
			file.Close()
		}
	}

	src.files = nil
	src.fileNames = nil
	log("src", 1, "closed all %v %v files", l, src)
}

// ---
// ---
// ---

// LoadSources loads contents for all sources, both from commandline and from files.
func LoadSources() {
	log("src", 1, "loading sources")

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

	SrcIP.ReportLoaded()
	SrcLogin.ReportLoaded()
	SrcPassword.ReportLoaded()

	ParseEndpoints(SrcIP.plain)
	ParseEndpoints(SrcIP.contents)

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

// CloseSources closes all source files.
func CloseSources() {
	log("src", 1, "closing sources")

	SrcIP.CloseSource()
	SrcLogin.CloseSource()
	SrcPassword.CloseSource()

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

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

	registerSwitch("add-empty-password", "insert an empty password to the password list")
	registerSwitch("no-password-trim", "preserve leading and trailing spaces in passwords")
	registerSwitch("logins-first", "increment logins before passwords")
	registerSwitch("file-contents-first", "try to go through source files first, defer commandline args for later")
}