apply geoip2, config file.

.gitignore

@@ -243,3 +243,4 @@ $RECYCLE.BIN/
 # End of https://www.toptal.com/developers/gitignore/api/go,goland,git,dotenv,database,macos,linux,windows
 
 keys/
+*.mmdb

go.mod

@@ -7,6 +7,8 @@ require github.com/gliderlabs/ssh v0.3.8
 require (
 	github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
 	github.com/joho/godotenv v1.5.1 // indirect
+	github.com/oschwald/geoip2-golang v1.13.0 // indirect
+	github.com/oschwald/maxminddb-golang v1.13.0 // indirect
 	golang.org/x/crypto v0.43.0 // indirect
 	golang.org/x/sys v0.37.0 // indirect
 )

go.sum

@@ -4,6 +4,10 @@ github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c=
 github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU=
 github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
 github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
+github.com/oschwald/geoip2-golang v1.13.0 h1:Q44/Ldc703pasJeP5V9+aFSZFmBN7DKHbNsSFzQATJI=
+github.com/oschwald/geoip2-golang v1.13.0/go.mod h1:P9zG+54KPEFOliZ29i7SeYZ/GM6tfEL+rgSn03hYuUo=
+github.com/oschwald/maxminddb-golang v1.13.0 h1:R8xBorY71s84yO06NgTmQvqvTvlS/bnYZrrWX1MElnU=
+github.com/oschwald/maxminddb-golang v1.13.0/go.mod h1:BU0z8BfFVhi1LQaonTwwGQlsHUEu9pWNdMfmq4ztm0o=
 golang.org/x/crypto v0.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04=
 golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0=
 golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=

inc.env

@@ -1 +1,2 @@
 PORT=2222
+GEOIP_DB=./GeoLite2-City.mmdb

main.go

@@ -3,15 +3,18 @@ package main
 import (
 	"fmt"
 	"log"
-	"os"
+	"slices"
+	"strings"
 
 	"sshchat/utils"
 
 	"github.com/gliderlabs/ssh"
-	"github.com/joho/godotenv"
+	"github.com/oschwald/geoip2-golang"
 )
 
-func sessionHandler(s ssh.Session) {
+var config = utils.GetConfig()
+
+func sessionHandler(s ssh.Session, geoip *geoip2.Reader) {
 	ptyReq, _, isPty := s.Pty()
 	if !isPty {
 		_, _ = fmt.Fprintln(s, "Err: PTY requires. Reconnect with -t option.")
@@ -22,25 +25,47 @@ func sessionHandler(s ssh.Session) {
 	remote := s.RemoteAddr().String()
 	username := s.User()
 
-	log.Printf("[sshchat] %s connected. %s", username, remote)
+	geoStatus := utils.GetIPInfo(remote, geoip)
+	if geoStatus == nil {
+		log.Printf("[sshchat] %s connected. %s / UNK [FORCE DISCONNECT]", username, remote)
+		_, _ = fmt.Fprintf(s, "[system] Your access country is blacklisted. UNK")
+		_ = s.Close()
+		return
+	} else {
+		log.Printf("[sshchat] %s connected. %s / %s", username, remote, geoStatus.Country)
+	}
+
+	if slices.Contains(config.CountryBlacklist, geoStatus.Country) {
+		log.Printf("[sshchat] %s country blacklisted. %s", username, remote)
+		_, _ = fmt.Fprintf(s, "[system] Your access country is blacklisted. %s\n", geoStatus.Country)
+		_ = s.Close()
+	}
+
+	if geoStatus.Country == "ZZ" && !(strings.HasPrefix(remote, "127") || strings.HasPrefix(remote, "[::1]")) {
+		log.Printf("[sshchat] unknown country blacklisted. %s", username)
+		_, _ = fmt.Fprintf(s, "[system] Unknown country is blacklisted. %s\n", geoStatus.Country)
+		_ = s.Close()
+	} else {
+		log.Printf("[sshchat] %s is localhost whitelisted.", username)
+	}
+
 	client := utils.NewClient(s, ptyReq.Window.Height, ptyReq.Window.Width, username, remote)
 
 	defer func() {
 		client.Close()
-		log.Printf("[sshchat] %s disconnected. %s", username, remote)
+		log.Printf("[sshchat] %s disconnected. %s / %s", username, remote, geoStatus.Country)
 	}()
 
 	client.EventLoop()
 }
 
 func main() {
-	err := godotenv.Load()
+	geoip, err := utils.GetDB(config.Geoip)
+	port := config.Port
 	if err != nil {
-		log.Fatal("Error loading .env file")
+		log.Fatalf("Geoip db is error: %v", err)
 	}
 
-	port := os.Getenv("PORT")
-
 	keys, err := utils.CheckHostKey()
 	if err != nil {
 		log.Print("Failed to check SSH keys: generate one.\n", err)
@@ -57,7 +82,9 @@ func main() {
 
 	s := &ssh.Server{
 		Addr: ":" + port,
-		Handler: sessionHandler,
+		Handler: func(s ssh.Session) {
+			sessionHandler(s, geoip)
+		},
 	}
 	for _, key := range keys {
 		s.AddHostKey(key)

utils/config.go

@@ -0,0 +1,32 @@
+package utils
+
+import (
+	"log"
+	"os"
+	"strings"
+
+	"github.com/joho/godotenv"
+)
+
+type Config struct {
+	Port             string
+	Geoip            string
+	CountryBlacklist []string
+}
+
+func GetConfig() *Config {
+	err := godotenv.Load()
+	if err != nil {
+		log.Fatal("Error loading .env file")
+	}
+
+	port := os.Getenv("PORT")
+	geoip_dbfile := os.Getenv("GEOIP_DB")
+	country_blacklist := os.Getenv("COUNTRY_BLACKLIST")
+
+	return &Config{
+		Port:             port,
+		Geoip:            geoip_dbfile,
+		CountryBlacklist: strings.Split(country_blacklist, ","),
+	}
+}

utils/geoip.go

@@ -0,0 +1,79 @@
+package utils
+
+import (
+	"log"
+	"net"
+
+	"github.com/oschwald/geoip2-golang"
+)
+
+type IpInfo struct {
+	Country       string
+	City          string
+	Timezone      string
+	Isp           string
+	IsAnonymousIP bool
+}
+
+func GetDB(db string) (*geoip2.Reader, error) {
+	geoip, err := geoip2.Open(db)
+	if err != nil {
+		log.Panicf("Failed to open database: %v", err)
+	}
+
+	return geoip, err
+}
+
+func GetIPInfo(ip string, db *geoip2.Reader) *IpInfo {
+	parsedIp := net.ParseIP(ip)
+
+	country := func(ip net.IP) string {
+		country, _ := db.Country(parsedIp)
+
+		println(country.Country.IsoCode)
+
+		if country != nil && country.Country.IsoCode != "" {
+			println(country.Country.IsoCode)
+			return country.Country.IsoCode
+		} else {
+			return "ZZ"
+		}
+	}(parsedIp)
+	city, timezone := func(ip net.IP) (string, string) {
+		city, _ := db.City(parsedIp)
+
+		if city != nil {
+			return city.City.Names["en"], city.Location.TimeZone
+		} else {
+			return "Unknown", "UTC+0"
+		}
+	}(parsedIp)
+	isp := func(ip net.IP) string {
+		isp, _ := db.ISP(parsedIp)
+
+		if isp != nil {
+			return isp.ISP
+		} else {
+			return "Unknown"
+		}
+	}(parsedIp)
+	isAnonymousIP := func(ip net.IP) bool {
+		is, _ := db.AnonymousIP(parsedIp)
+
+		if is != nil {
+			return is.IsAnonymousVPN ||
+				is.IsPublicProxy ||
+				is.IsAnonymous
+		} else {
+			return false
+		}
+	}(parsedIp)
+
+	return &IpInfo{
+		Country:       country,
+		City:          city,
+		Timezone:      timezone,
+		Isp:           isp,
+		IsAnonymousIP: isAnonymousIP,
+	}
+}