main.go

package main

import (
	"bufio"
	"fmt"
	"log"
	"math"
	"os"
	"sort"
	"strings"

	"github.com/alecthomas/kingpin/v2"
)

var (
	wordlistPath = kingpin.Arg("wordlist", "Path to wordlist file").Required().String()
	top          = kingpin.Flag("top", "Number of top words to display").Default("10").Int()
	outputFile   = kingpin.Flag("output", "Output file").String()
)

func main() {

	kingpin.Parse()

	// Load wordlist
	wordlist := ReadLines(*wordlistPath)

	if *top > len(wordlist) {
		*top = len(wordlist)
	}

	// Calculate entropy for each word in the wordlist and store it in a map.
	wordEntropy := make(map[string]float64)
	for _, word := range wordlist {
		wordEntropy[word] = ShannonEntropy(word)
	}

	// Sort the wordlist by entropy.
	sortedWordlist := SortByValue(wordEntropy)

	// Reverse the sorted wordlist to get the highest entropy words first.
	for i, j := 0, len(sortedWordlist)-1; i < j; i, j = i+1, j-1 {
		sortedWordlist[i], sortedWordlist[j] = sortedWordlist[j], sortedWordlist[i]
	}

	// Print the top words with their entropy.
	for _, word := range sortedWordlist[:*top] {
		fmt.Printf("%s\t%.2f\n", word, wordEntropy[word])
	}

	// Write the output to a file if specified.
	if *outputFile != "" {
		file, err := os.Create(*outputFile)
		if err != nil {
			log.Fatal(err)
		}
		defer file.Close()

		for _, word := range sortedWordlist {
			fmt.Fprintf(file, "%s\n", word)
		}
	}
}

// SortByValue sorts a map by its values.
func SortByValue(m map[string]float64) []string {
	var keys []string
	for key := range m {
		keys = append(keys, key)
	}
	// Sort keys by value
	sort.Slice(keys, func(i, j int) bool {
		return m[keys[i]] < m[keys[j]]
	})
	return keys
}

// ShannonEntropy calculates the entropy of a string using the Shannon entropy formula.
//
// The Shannon entropy formula is defined as:
//
// H(X) = - Σ (P(xi) * log2(P(xi)))
//
// Where:
// - H(X) is the entropy of the string.
// - Σ is the sum over all characters in the string.
// - xi is a character in the string at i'th position.
// - p(x) is the probability of the character x.
func ShannonEntropy(s string) float64 {
	entropy := 0.0
	probabilities := Probabilities(s)

	for _, p := range probabilities {
		entropy -= p * math.Log2(p)
	}
	return entropy
}

// Probabilities calculates the probability of each character in a string.
func Probabilities(s string) map[string]float64 {
	probabilities := make(map[string]float64)
	for _, char := range s {
		probabilities[string(char)]++
	}
	for key, value := range probabilities {
		probabilities[key] = value / float64(len(s))
	}
	return probabilities
}

// ReadLines reads a file and returns a slice of strings representing each line.
func ReadLines(filePath string) []string {
	var lines []string
	file, err := os.Open(filePath)
	if err != nil {
		log.Fatal(err)
	}
	scanner := bufio.NewScanner(file)
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		lines = append(lines, strings.TrimSpace(scanner.Text()))
	}
	file.Close()
	return lines
}

func TestShannonEntropy() {
	// Test cases for ShannonEntropy function.
	testCases := []struct {
		input    string
		expected float64
	}{
		{"hello", 2.584962500721156},
		{"world", 2.584962500721156},
		{"hello world", 3.169925001442312},
		{"entropy", 2.807354922057604},
		{"shannon", 2.807354922057604},
		{"shannon entropy", 3.169925001442312},
	}

	for _, tc := range testCases {
		actual := ShannonEntropy(tc.input)
		if actual != tc.expected {
			fmt.Printf("FAIL: ShannonEntropy(%q) = %f; expected %f\n", tc.input, actual, tc.expected)
		}
	}
}