// GoPass
// Author: Maximilian Patterson
package main

import (
	"crypto/rand"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"strings"
)

var allowedCharacters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890`~!@#$%^&*()_+[]\\{}|;':,./<>?")

const (
	Version   = "1.3.5"
	symbols   = "`~!@#$%^&*()_+[]\\{}|;':,./<>?"
	chunkSize = 16 // The size of each chunk of the password to be generated by the worker goroutines
)

func matchArguments(args []string) string {
	// If there are no arguments
	if len(args) == 0 {
		return "No password length specified! (ex: gopass 16)"
	}

	// First argument is special, must be an integer, -v, or -h
	var size = 0 // Password length
	switch args[0] {
	case "-v":
		return "GoPass version " + Version
	case "-h":
		return "GoPass - A simple password generator written in Go\n" +
			"Usage: gopass [length] [disallowed characters] [optional remove symbols -s]\n" +
			"     Example: gopass 16\n" +
			"     Example: gopass 16 -r=abc123!@#\n" +
			"     Example: gopass 16 -s\n" +
			"\nFor help (this output): gopass -h\n" +
			"For version: gopass -v\n"
	default:
		var err error
		size, err = strconv.Atoi(args[0])
		if err != nil {
			return "Invalid first argument (\"" + args[0] + "\") supplied! (Type gopass -h for help)"
		}
	}

	for i := 1; i < len(args); i++ {
		v := args[i]
		switch {
		case v == "-s":
			removeDisallowed([]rune(symbols))
		case strings.HasPrefix(v, "-r="):
			// Remove all characters after the = until next whitespace
			removeDisallowed([]rune(v[2:]))
		default:
			return "Invalid argument (\"" + v + "\") supplied! (Type gopass -h for help)"
		}
	}

	if size <= 0 {
		return "No/invalid password length specified! (ex: gopass 16)"
	} else {
		return generatePassword(size)
	}
}

// Remove all disallowed characters from the allowedCharacters slice
func removeDisallowed(disallowed []rune) {
	disallowedMap := make(map[rune]bool, len(disallowed))
	for _, r := range disallowed {
		disallowedMap[r] = true
	}

	i := 0
	for _, v := range allowedCharacters {
		if !disallowedMap[v] {
			allowedCharacters[i] = v
			i++
		}
	}
	allowedCharacters = allowedCharacters[:i]
}

func generatePassword(size int) string {
	// Make empty array of runes with size of size
	pass := make([]rune, size)

	// Create a channel to receive chunks of the password
	passChan := make(chan []rune)

	// Determine the number of worker goroutines to use
	numWorkers := runtime.NumCPU()

	// Launch the worker goroutines
	for range numWorkers {
		go func() {
			allowedLen := len(allowedCharacters)

			// Calculate the rejection limit to avoid modulo bias.
			// Any random byte value > maxByte must be discarded.
			maxByte := 255 - (256 % allowedLen)

			// Create a buffer for random bytes.
			// We read enough bytes for the whole chunk at once (plus some extra
			// in case of rejections) to reduce system calls.
			randBuf := make([]byte, chunkSize*2)

			for {
				// Fill the buffer initially
				if _, err := rand.Read(randBuf); err != nil {
					println("Error securely generating random character chunk!")
					return
				}
				bufIdx := 0

				chunk := make([]rune, chunkSize)
				for i := range chunk {
					for {
						// If we exhausted the buffer due to rejections, refill it
						if bufIdx >= len(randBuf) {
							if _, err := rand.Read(randBuf); err != nil {
								println("Error securely generating random character chunk!")
								return
							}
							bufIdx = 0
						}

						b := randBuf[bufIdx]
						bufIdx++

						// If the byte causes bias, discard and retry
						if int(b) > maxByte {
							continue
						}

						// Safe to map to character set
						chunk[i] = allowedCharacters[int(b)%allowedLen]
						break
					}
				}
				// Send the chunk of the password to the main goroutine
				passChan <- chunk
			}
		}()
	}

	// Collect the chunks of the password from the channel
	for i := 0; i < size; i += chunkSize {
		chunk := <-passChan
		copy(pass[i:], chunk)
	}

	return string(pass)
}

func main() {
	// Process arguments
	fmt.Println(matchArguments(os.Args[1:]))
}