From d82d2b2ea923826121819703786b38fc7cee7402 Mon Sep 17 00:00:00 2001
From: GGP1 <ggpalomeque@gmail.com>
Date: Tue, 16 Feb 2021 23:34:38 -0300
Subject: [PATCH] Fix naming and entropy

---
 example_test.go  | 11 +++++++++++
 passphrase.go    | 11 +++++------
 password.go      | 27 ++++++++++++++++++---------
 password_test.go |  2 +-
 4 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/example_test.go b/example_test.go
index 0045199..083aa9c 100644
--- a/example_test.go
+++ b/example_test.go
@@ -66,3 +66,14 @@ func ExampleNewPassphrase() {
 	// Example output:
 	// ynuafnezm hvoq asruso jvoe psiro
 }
+
+func ExampleKeyspace() {
+	p := &atoll.Password{
+		Length: 6,
+		Levels: []atoll.Level{atoll.Lowercase},
+		Repeat: false,
+	}
+
+	fmt.Println(atoll.Keyspace(p))
+	// Output: 3.089157759999998e+08
+}
diff --git a/passphrase.go b/passphrase.go
index 760b9dc..016e5c1 100644
--- a/passphrase.go
+++ b/passphrase.go
@@ -84,7 +84,7 @@ func (p *Passphrase) generate() (string, error) {
 		}
 	}
 
-	// Initialize secret slice (included words will be appended)
+	// Initialize secret slice
 	p.words = make([]string, int(p.Length))
 	// Defaults
 	if p.Separator == "" {
@@ -109,7 +109,7 @@ func (p *Passphrase) generate() (string, error) {
 
 // includeWords randomly inserts included words in the passphrase.
 func (p *Passphrase) includeWords() {
-	// Append included words
+	// Add included words at the end of the secret
 	for i, word := range p.Include {
 		p.words[int(p.Length)-i-1] = word
 	}
@@ -121,7 +121,7 @@ func (p *Passphrase) includeWords() {
 	}
 }
 
-// Check if any excluded word is within the secret and (if true) replace it with another random word.
+// excludeWords checks if any excluded word is within the secret and (if true) replace it with another random word.
 func (p *Passphrase) excludeWords() {
 	for i, word := range p.words {
 		for _, excl := range p.Exclude {
@@ -177,7 +177,7 @@ func SyllableList(p *Passphrase) {
 	}
 }
 
-// Entropy returns the bits of entropy of the passphrase.
+// Entropy returns the passphrase entropy in bits.
 //
 // If the list used is "NoList" the secret must be already generated.
 func (p *Passphrase) Entropy() float64 {
@@ -192,8 +192,7 @@ func (p *Passphrase) Entropy() float64 {
 		words := strings.Join(p.words, "")
 		// Take out the separators from the secret length
 		secretLength := len(words) - (len(p.Separator) * int(p.Length))
-		// -26- represents the dictionary length (vowels+constants)
-		return math.Log2(math.Pow(float64(26), float64(secretLength)))
+		return math.Log2(math.Pow(float64(len(vowels)+len(constants)), float64(secretLength)))
 	case "WordList":
 		poolLength = len(atollWords)
 	case "SyllableList":
diff --git a/password.go b/password.go
index c9b13b0..650da91 100644
--- a/password.go
+++ b/password.go
@@ -107,12 +107,12 @@ func (p *Password) generate() (string, error) {
 	remaining := int(p.Length) - len(password)
 
 	for i := 0; i < remaining; i++ {
-		c := p.pool[randInt(len(p.pool))]
-		password = randInsert(password, c)
+		char := p.pool[randInt(len(p.pool))]
+		password = randInsert(password, char)
 
 		if !p.Repeat {
 			// Remove element used
-			p.pool = strings.Replace(p.pool, string(c), "", 1)
+			p.pool = strings.Replace(p.pool, string(char), "", 1)
 		}
 	}
 
@@ -122,19 +122,19 @@ func (p *Password) generate() (string, error) {
 }
 
 func (p *Password) generatePool() {
-	var b strings.Builder
+	var sb strings.Builder
 	unique := make(map[Level]struct{})
 
 	for _, lvl := range p.Levels {
 		// Ensure that duplicated levels aren't added twice
 		if _, ok := unique[lvl]; !ok && len(lvl) > 0 {
 			unique[lvl] = struct{}{}
-			b.Grow(len(lvl))
-			b.WriteString(string(lvl))
+			sb.Grow(len(lvl))
+			sb.WriteString(string(lvl))
 		}
 	}
 
-	p.pool = b.String()
+	p.pool = sb.String()
 }
 
 // initPassword creates the password, adds any included word and makes sure that it contains
@@ -207,7 +207,7 @@ repeat:
 	return password
 }
 
-// validateLevels checks if Exclude contains all the characters of a level that is in Format.
+// validateLevels checks if Exclude contains all the characters of a level that is in Levels.
 func (p *Password) validateLevels() error {
 	for _, lvl := range p.Levels {
 		if len(lvl) < 1 {
@@ -250,13 +250,17 @@ func (p *Password) validateLevels() error {
 	return nil
 }
 
-// Entropy returns the bits of entropy of the password.
+// Entropy returns the password entropy in bits.
 func (p *Password) Entropy() float64 {
 	poolLength := len(p.pool)
 	if p.pool == "" {
 		p.generatePool()
 		poolLength = len(p.pool)
 
+		if !p.Repeat {
+			poolLength -= int(p.Length)
+		}
+
 		// Do not count 2/3 byte characters as they aren't in the pool
 		for _, excl := range p.Exclude {
 			if len(string(excl)) != 1 {
@@ -265,5 +269,10 @@ func (p *Password) Entropy() float64 {
 		}
 	}
 
+	// Add the characters that were removed from the pool
+	if !p.Repeat {
+		poolLength += int(p.Length)
+	}
+
 	return math.Log2(math.Pow(float64(poolLength), float64(p.Length)))
 }
diff --git a/password_test.go b/password_test.go
index f07ab60..56843a4 100644
--- a/password_test.go
+++ b/password_test.go
@@ -274,7 +274,7 @@ func TestPasswordEntropy(t *testing.T) {
 		Exclude: "a1r/ö",
 	}
 
-	expected := 131.09177703355275
+	expected := 136.65780028329485
 
 	got := p.Entropy()
 	if got != expected {