Performance improvement

main.go

@@ -3,7 +3,7 @@ package main
 import (
 	"bufio"
 	"encoding/hex"
-	"fmt"
+	"io"
 	"os"
 	"runtime"
 	"strings"
@@ -14,23 +14,32 @@ import (
 	"github.com/spf13/pflag"
 )
 
+type hashed struct {
+	password string
+	hash     string
+}
+
 func main() {
 	hash := pflag.String("hash", "ntlm", "hash type to generate (ntlm)")
 	parallel := pflag.Int("parallel", runtime.NumCPU(), "number of threads to use")
+	inputname := pflag.String("input", "", "file to read passwords from, blank for stdin")
+	outputname := pflag.String("output", "", "file to write hashed passwords to, blank for stdout")
 	pflag.Parse()
 
-	queue := make(chan string, *parallel*4)
-	var wg sync.WaitGroup
+	plaintextqueue := make(chan string, *parallel*4)
+	hashedqueue := make(chan hashed, *parallel*4)
+
+	var producerWait, consumerWait sync.WaitGroup
 
 	for i := 0; i < *parallel; i++ {
-		wg.Add(1)
+		producerWait.Add(1)
 		go func() {
-			defer wg.Done()
+			defer producerWait.Done()
 			switch *hash {
 			case "ntlm":
 				u16 := make([]byte, 16)
 				mdfour := md4.New()
-				for password := range queue {
+				for password := range plaintextqueue {
 					/* Add all bytes, as well as the 0x00 of UTF-16 */
 					utf16encoded := utf16.Encode([]rune(password))
 					if cap(u16) < len(utf16encoded)*2 {
@@ -48,17 +57,70 @@ func main() {
 					md4 := mdfour.Sum(nil)
 
 					/* Return the output */
-					fmt.Println(password + ":" + hex.EncodeToString(md4))
+					hashedqueue <- hashed{
+						password: password,
+						hash:     hex.EncodeToString(md4),
+					}
 				}
 			}
 		}()
 	}
 
+	var writeto io.WriteCloser
+	writeto = os.Stdout
+	if *outputname != "" {
+		outputfile, err := os.Create(*outputname)
+		if err != nil {
+			panic("could not write to" + *outputname + ": " + err.Error())
+		}
+		writeto = outputfile
+	}
+	defer writeto.Close()
+
+	output := bufio.NewWriter(writeto)
+	defer output.Flush()
+
+	go func() {
+		defer consumerWait.Done()
+		var sb strings.Builder
+		for hashed := range hashedqueue {
+			sb.Reset()
+			sb.Grow(len(hashed.password) + len(hashed.hash) + 2)
+			sb.WriteString(hashed.password)
+			sb.WriteString(":")
+			sb.WriteString(hashed.hash)
+			sb.WriteString("\n")
+			_, err := output.WriteString(sb.String())
+			if err != nil {
+				panic("could not write: " + err.Error())
+			}
+		}
+	}()
+
 	// read lines from stdin and put them in the queue
-	scanner := bufio.NewScanner(os.Stdin)
+	var readfrom io.ReadCloser
+	readfrom = os.Stdin
+
+	if *inputname != "" {
+		inputfile, err := os.Open(*inputname)
+		if err != nil {
+			panic("could not read from" + *inputname + ": " + err.Error())
+		}
+		readfrom = inputfile
+	}
+	defer readfrom.Close()
+
+	scanner := bufio.NewScanner(bufio.NewReaderSize(readfrom, 1024*1024))
 	for scanner.Scan() {
-		queue <- strings.Trim(scanner.Text(), "\r\n")
+		line := scanner.Bytes()
+		// for len(line) > 0 && (line[len(line)-1] == '\r' || line[len(line)-1] == '\n') {
+		// 	line = line[:len(line)-1]
+		// }
+		plaintextqueue <- string(line)
 	}
-	close(queue)
-	wg.Wait()
+	close(plaintextqueue)
+	producerWait.Wait()
+
+	close(hashedqueue)
+	consumerWait.Wait()
 }