From 12dc19a6f355a384a59d8fd7e2164e5bc9042684 Mon Sep 17 00:00:00 2001
From: Aria Edmonds <8436007+ar1a@users.noreply.github.com>
Date: Sun, 12 Feb 2023 03:40:08 +0800
Subject: [PATCH] Add keys per second display (#8)

* feature: add key rate

* feature: print keys processed with commas

* switch to immediate keys/s display

* switch to moving average

* use time.Since
---
 go.mod       |  5 ++++-
 go.sum       |  2 ++
 vanityssh.go | 34 ++++++++++++++++++++++++++++++++--
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index 8bf81a3..b87f8fd 100644
--- a/go.mod
+++ b/go.mod
@@ -7,4 +7,7 @@ require (
 	golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e
 )
 
-require golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 // indirect
+require (
+	github.com/dustin/go-humanize v1.0.1 // indirect
+	golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 // indirect
+)
diff --git a/go.sum b/go.sum
index ef49e5f..14f8f61 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/mikesmitty/edkey v0.0.0-20170222072505-3356ea4e686a h1:eU8j/ClY2Ty3qdHnn0TyW3ivFoPC/0F1gQZz8yTxbbE=
 github.com/mikesmitty/edkey v0.0.0-20170222072505-3356ea4e686a/go.mod h1:v8eSC2SMp9/7FTKUncp7fH9IwPfw+ysMObcEz5FWheQ=
 golang.org/x/crypto v0.0.0-20220525230936-793ad666bf5e h1:T8NU3HyQ8ClP4SEE+KbFlg6n0NhuTsN4MyznaarGsZM=
diff --git a/vanityssh.go b/vanityssh.go
index ccca0c7..2cdf81e 100644
--- a/vanityssh.go
+++ b/vanityssh.go
@@ -8,6 +8,7 @@ import (
 	"flag"
 	"fmt"
 	"io/ioutil"
+	"math"
 	"os"
 	"os/signal"
 	"regexp"
@@ -16,6 +17,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	"github.com/mikesmitty/edkey"
 	"golang.org/x/crypto/ed25519"
 	"golang.org/x/crypto/ssh"
@@ -26,7 +28,7 @@ var global_user_insensitive bool
 var global_user_streaming bool
 var global_user_fingerprint bool
 
-//var flagvar int
+// var flagvar int
 var global_counter int64
 var start time.Time
 var re *regexp.Regexp
@@ -109,6 +111,11 @@ func getAuthorizedKey(key ssh.PublicKey) string {
 	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key)))
 }
 
+func expMovingAverage(value, oldValue, deltaTime, timeWindow float64) float64 {
+	alpha := 1.0 - math.Exp(-deltaTime/timeWindow)
+	return alpha*value + (1.0-alpha)*oldValue
+}
+
 func main() {
 	//	input threads, else numcpu
 	for i := 1; i <= runtime.NumCPU(); i++ {
@@ -116,9 +123,32 @@ func main() {
 	}
 
 	fmt.Printf("Press Ctrl+C to end\n")
+
+	deleteLine := "\033[2K\r"
+	cursorUp := "\033[A"
+	avgKeyRate := float64(global_counter)
+	oldCounter := global_counter
+	oldTime := time.Now()
+
 	for {
-		fmt.Printf("\033[2K\r%s%d", "SSH Keys Processed = ", global_counter)
 		time.Sleep(250 * time.Millisecond)
+		relTime := time.Since(oldTime).Seconds()
+
+		// on first run, initialize the moving average with the current rate
+		// instead of starting at 0 and taking many seconds to tend towards the
+		// actual key rate
+		if oldCounter == 0 {
+			avgKeyRate = float64(global_counter)
+		}
+
+		fmt.Printf("%s%s%s", deleteLine, cursorUp, deleteLine)
+		fmt.Printf("SSH Keys Processed = %s\n", humanize.Comma(global_counter))
+		fmt.Printf("kKeys/s = %.2f", avgKeyRate/relTime/1000)
+
+		avgKeyRate = expMovingAverage(
+			float64(global_counter-oldCounter), avgKeyRate, relTime, 5)
+		oldCounter = global_counter
+		oldTime = time.Now()
 	}
 
 	WaitForCtrlC()