From 44f611aa6b22ac3217e86a45a8ad423dc5fe1478 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Torjus=20H=C3=A5kestad?= <torjus.hakestad@usit.uio.no>
Date: Sun, 11 Apr 2021 02:24:39 +0200
Subject: [PATCH] Make connection throttle configurable

---
 apiary.toml        |  4 ++++
 config/config.go   |  7 ++++---
 honeypot/conn.go   | 21 ++++++++++++---------
 honeypot/server.go | 11 +++++++----
 4 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/apiary.toml b/apiary.toml
index 72f60b4..5124c5a 100644
--- a/apiary.toml
+++ b/apiary.toml
@@ -22,6 +22,10 @@ LogLevel = "INFO"
 # Address and port to listen to 
 # Default: ":2222"
 ListenAddr = ":2222"
+# Throttle incoming and outgoing data per connection
+# Values are in bytes per second. Empty means no unlimited
+# Default: ""
+ThrottleSpeed = 10240
 
 [Frontend]
 # Log level for SSH Honeypot
diff --git a/config/config.go b/config/config.go
index 68761f8..085a763 100644
--- a/config/config.go
+++ b/config/config.go
@@ -24,9 +24,10 @@ type PostgresStoreConfig struct {
 }
 
 type HoneypotConfig struct {
-	ListenAddr  string `toml:"ListenAddr"`
-	LogLevel    string `toml:"LogLevel"`
-	HostKeyPath string `toml:"HostKeyPath"`
+	ListenAddr    string  `toml:"ListenAddr"`
+	LogLevel      string  `toml:"LogLevel"`
+	HostKeyPath   string  `toml:"HostKeyPath"`
+	ThrottleSpeed float64 `toml:"ThrottleSpeed"`
 }
 
 type FrontendConfig struct {
diff --git a/honeypot/conn.go b/honeypot/conn.go
index f934d25..8f5ae4f 100644
--- a/honeypot/conn.go
+++ b/honeypot/conn.go
@@ -11,29 +11,32 @@ import (
 type throttledConn struct {
 	ID            uuid.UUID
 	conn          net.Conn
-	speed         float64
+	writer        *shapeio.Writer
+	reader        *shapeio.Reader
 	CloseCallback func(c *throttledConn)
 }
 
 func newThrottledConn(conn net.Conn) *throttledConn {
 	id := uuid.Must(uuid.NewRandom())
-	return &throttledConn{ID: id, conn: conn, speed: 1024 * 10}
+	return &throttledConn{
+		ID:     id,
+		conn:   conn,
+		writer: shapeio.NewWriter(conn),
+		reader: shapeio.NewReader(conn),
+	}
 }
 
 func (sc *throttledConn) SetSpeed(bytesPerSec float64) {
-	sc.speed = bytesPerSec
+	sc.writer.SetRateLimit(bytesPerSec)
+	sc.reader.SetRateLimit(bytesPerSec)
 }
 
 func (sc *throttledConn) Read(b []byte) (n int, err error) {
-	slowReader := shapeio.NewReader(sc.conn)
-	slowReader.SetRateLimit(sc.speed)
-	return slowReader.Read(b)
+	return sc.reader.Read(b)
 }
 
 func (sc *throttledConn) Write(b []byte) (n int, err error) {
-	slowWriter := shapeio.NewWriter(sc.conn)
-	slowWriter.SetRateLimit(sc.speed)
-	return slowWriter.Write(b)
+	return sc.writer.Write(b)
 }
 
 func (sc *throttledConn) Close() error {
diff --git a/honeypot/server.go b/honeypot/server.go
index 9a357bb..e20fe03 100644
--- a/honeypot/server.go
+++ b/honeypot/server.go
@@ -19,12 +19,14 @@ import (
 )
 
 type HoneypotServer struct {
-	attemptStore      store.LoginAttemptStore
-	attemptsCallbacks []func(l models.LoginAttempt)
+	Logger *zap.SugaredLogger
 
 	sshServer *ssh.Server
 
-	Logger *zap.SugaredLogger
+	attemptStore      store.LoginAttemptStore
+	attemptsCallbacks []func(l models.LoginAttempt)
+
+	throttleSpeed float64
 }
 
 func NewHoneypotServer(cfg config.HoneypotConfig, store store.LoginAttemptStore) (*HoneypotServer, error) {
@@ -107,8 +109,9 @@ func (hs *HoneypotServer) passwordHandler(ctx ssh.Context, password string) bool
 
 func (s *HoneypotServer) connCallback(ctx ssh.Context, conn net.Conn) net.Conn {
 	throttledConn := newThrottledConn(conn)
+	throttledConn.SetSpeed(s.throttleSpeed)
 	ctx.SetValue("uuid", throttledConn.ID)
-	throttledConn.SetSpeed(2048)
+	throttledConn.SetSpeed(s.throttleSpeed)
 	return throttledConn
 }