From 6e445b91e8b12d9ab3ca2da05e5879a689a944cd Mon Sep 17 00:00:00 2001
From: Boshi Lian <farmer1992@gmail.com>
Date: Wed, 5 Feb 2025 03:30:12 -0800
Subject: [PATCH] Enhance failtoban plugin with log-only mode and SIGHUP cache
 reset

---
 plugin/failtoban/main.go | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/plugin/failtoban/main.go b/plugin/failtoban/main.go
index 2b71b9ac..924211c6 100644
--- a/plugin/failtoban/main.go
+++ b/plugin/failtoban/main.go
@@ -5,6 +5,9 @@ package main
 import (
 	"fmt"
 	"net"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	gocache "github.com/patrickmn/go-cache"
@@ -17,7 +20,7 @@ func main() {
 
 	libplugin.CreateAndRunPluginTemplate(&libplugin.PluginTemplate{
 		Name:  "failtoban",
-		Usage: "sshpiperd fixed plugin, only password auth is supported",
+		Usage: "failtoban plugin, block ip after too many auth failures",
 		Flags: []cli.Flag{
 			&cli.IntFlag{
 				Name:    "max-failures",
@@ -31,14 +34,33 @@ func main() {
 				EnvVars: []string{"SSHPIPERD_FAILTOBAN_BAN_DURATION"},
 				Value:   60 * time.Minute,
 			},
+			&cli.BoolFlag{
+				Name:    "log-only",
+				Usage:   "log only mode, no ban, useful for working with other tools like fail2ban",
+				EnvVars: []string{"SSHPIPERD_FAILTOBAN_LOG_ONLY"},
+				Value:   false,
+			},
 		},
 		CreateConfig: func(c *cli.Context) (*libplugin.SshPiperPluginConfig, error) {
 
 			maxFailures := c.Int("max-failures")
 			banDuration := c.Duration("ban-duration")
+			logOnly := c.Bool("log-only")
 
 			cache := gocache.New(banDuration, banDuration/2*3)
 
+			// register signal handler
+			go func() {
+				sigChan := make(chan os.Signal, 1)
+				signal.Notify(sigChan, syscall.SIGHUP)
+
+				for {
+					<-sigChan
+					cache.Flush()
+					log.Info("failtoban: cache reset due to SIGHUP")
+				}
+			}()
+
 			return &libplugin.SshPiperPluginConfig{
 				NoClientAuthCallback: func(conn libplugin.ConnMetadata) (*libplugin.Upstream, error) {
 					// in case someone put the failtoban plugin before other plugins
@@ -47,6 +69,10 @@ func main() {
 					}, nil
 				},
 				NewConnectionCallback: func(conn libplugin.ConnMetadata) error {
+					if logOnly {
+						return nil
+					}
+
 					ip, _, _ := net.SplitHostPort(conn.RemoteAddr())
 
 					failed, found := cache.Get(ip)
@@ -64,12 +90,12 @@ func main() {
 				UpstreamAuthFailureCallback: func(conn libplugin.ConnMetadata, method string, err error, allowmethods []string) {
 					ip, _, _ := net.SplitHostPort(conn.RemoteAddr())
 					failed, _ := cache.IncrementInt(ip, 1)
-					log.Debugf("failtoban: %v auth failed. current status: fail %v times, max allowed %v", ip, failed, maxFailures)
+					log.Warnf("failtoban: %v auth failed. current status: fail %v times, max allowed %v", ip, failed, maxFailures)
 				},
 				PipeCreateErrorCallback: func(remoteAddr string, err error) {
 					ip, _, _ := net.SplitHostPort(remoteAddr)
 					failed, _ := cache.IncrementInt(ip, 1)
-					log.Debugf("failtoban: %v pipe create failed, reason %v. current status: fail %v times, max allowed %v", ip, err, failed, maxFailures)
+					log.Warnf("failtoban: %v pipe create failed, reason %v. current status: fail %v times, max allowed %v", ip, err, failed, maxFailures)
 				},
 			}, nil
 		},