mtbf/source.go

package main

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

// SrcIP, SrcLogin and SrcPassword represent different source types.
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: func(item string) (res string, err error) {
		if getParamSwitch("no-password-trim") {
			return item, nil
		} else {
			return strings.TrimSpace(item), nil
		}
	}}

func init() {
	registerParam("ip", []string{}, "IPs or subnets in CIDR notation")
	registerParam("ip-file", []string{"ip.txt"}, "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{"login.txt"}, "paths to files with logins (one entry per line)")
	registerParam("password", []string{}, "one or more passwords")
	registerParam("password-file", []string{"password.txt"}, "paths to files with passwords (one entry per line)")

	registerSwitch("add-empty-password", "insert an empty password to the password list")
	registerSwitch("add-logins-as-passwords", "append all logins 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")
}

// LoadSources loads contents for all sources.
func LoadSources() {
	log("src", 1, "loading sources")

	var wg sync.WaitGroup
	if !getParamBool("add-logins-as-passwords") {
		wg.Add(3)
		go SrcIP.LoadSource(&wg)
		go SrcLogin.LoadSource(&wg)
		go SrcPassword.LoadSource(&wg)
	} else {
		wg.Add(1)
		go SrcLogin.LoadSource(&wg)
		wg.Wait()

		wg.Add(3)
		go SrcIP.LoadSource(&wg)
		go SrcPassword.LoadSource(&wg)
	}
	wg.Wait()

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

	ParseEndpoints(SrcIP.plain, "commandline")
	ParseEndpoints(SrcIP.contents, "files")

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

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

	SrcIP.CloseFiles()
	SrcLogin.CloseFiles()
	SrcPassword.CloseFiles()

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

// 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, "")
	}

	if src.name == "password" && getParamSwitch("add-logins-as-passwords") {
		src.plain = append(src.plain, SrcLogin.plain...)
		src.plain = append(src.plain, SrcLogin.contents...)
	}

	src.ParsePlain()

	src.files = make(map[string]*os.File)
	src.OpenFiles()
	defer src.CloseFiles()

	src.ParseFiles()
	failIf(len(src.contents)+len(src.plain) == 0, "no %vs defined: check %v and %v parameters", src, src.plainParmName, src.filesParmName)
}

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

	log("src", 1, "closed %v %v files", len(src.files), src)
	src.files = nil
}

// OpenFiles opens all files for a Source.
func (src *Source) OpenFiles() {
	fileNames := getParamStringSlice(src.filesParmName)

	for _, fileName := range fileNames {
		if src.files[fileName] != nil {
			log("src", 0, "ignoring duplicate %v file \"%v\"", src, fileName)
			continue
		}

		f, err := os.Open(fileName)
		failIf(err != nil, "error opening source file \"%v\": %v", fileName, err)
		src.files[fileName] = f
	}

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

// 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
	}
}

// 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)
	}
}

// ParseFiles parses files for a Source.
func (src *Source) ParseFiles() {
	for name, file := range src.files {
		log("src", 1, "parsing %v", name)
		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", name, scannerErr)
		log("src", 1, "ok: parsed \"%v\", got %v contents, %v total", name, thisTotal, len(src.contents))
	}
}

// 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))
}

type Source struct {
	name          string // name of this source
	plainParmName string // name of "plain" commandline parameter
	filesParmName string // name of "files" commandline parameter

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

	files map[string]*os.File // file pointers and names

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

	fetchMutex sync.Mutex // sync mutex
}

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

// 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", 4, "exhausted source %v for pos %v", src, pos.String())
	return res, empty
}

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

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

func (pos *SourcePos) Init() {
	pos.plainIdx = 0
	pos.contentIdx = 0
}

// 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")
}