From 87b6d1c93fc3b988aaa20588cdf51f4d647d858d Mon Sep 17 00:00:00 2001
From: dushixiang <helloworld1024@foxmail.com>
Date: Mon, 15 Nov 2021 20:45:43 +0800
Subject: [PATCH] =?UTF-8?q?sshd=20=E6=9C=8D=E5=8A=A1=E5=A2=9E=E5=8A=A0=20I?=
 =?UTF-8?q?P=20=E6=A0=A1=E9=AA=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/api/sshd.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/server/api/sshd.go b/server/api/sshd.go
index 565b5aa..24ca18f 100644
--- a/server/api/sshd.go
+++ b/server/api/sshd.go
@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
+	"next-terminal/server/global/security"
 	"path"
 	"strings"
 	"time"
@@ -376,6 +378,55 @@ func passwordAuth(ctx ssh.Context, pass string) bool {
 	return true
 }
 
+func connCallback(ctx ssh.Context, conn net.Conn) net.Conn {
+	securities := security.GlobalSecurityManager.Values()
+	if len(securities) == 0 {
+		return conn
+	}
+
+	ip := strings.Split(conn.RemoteAddr().String(), ":")[0]
+
+	for _, s := range securities {
+		if strings.Contains(s.IP, "/") {
+			// CIDR
+			_, ipNet, err := net.ParseCIDR(s.IP)
+			if err != nil {
+				continue
+			}
+			if !ipNet.Contains(net.ParseIP(ip)) {
+				continue
+			}
+		} else if strings.Contains(s.IP, "-") {
+			// 范围段
+			split := strings.Split(s.IP, "-")
+			if len(split) < 2 {
+				continue
+			}
+			start := split[0]
+			end := split[1]
+			intReqIP := utils.IpToInt(ip)
+			if intReqIP < utils.IpToInt(start) || intReqIP > utils.IpToInt(end) {
+				continue
+			}
+		} else {
+			// IP
+			if s.IP != ip {
+				continue
+			}
+		}
+
+		if s.Rule == constant.AccessRuleAllow {
+			return conn
+		}
+		if s.Rule == constant.AccessRuleReject {
+			_, _ = conn.Write([]byte("your access request was denied :(\n"))
+			return nil
+		}
+	}
+
+	return conn
+}
+
 func Setup() {
 	ssh.Handle(func(s ssh.Session) {
 		_, _ = io.WriteString(s, fmt.Sprintf(constant.Banner, constant.Version))
@@ -393,6 +444,7 @@ func Setup() {
 		nil,
 		ssh.PasswordAuth(passwordAuth),
 		ssh.HostKeyFile(config.GlobalCfg.Sshd.Key),
+		ssh.WrapConn(connCallback),
 	)
 	log.Fatal(fmt.Sprintf("启动sshd服务失败: %v", err.Error()))
 }