Implementation of routing feature based on a geolocation database

cmd/root.go

@@ -2,6 +2,8 @@ package cmd
 
 import (
 	"fmt"
+	"github.com/nodauf/Go-RouterSocks/router"
+	"log"
 	"os"
 
 	"github.com/spf13/cobra"
@@ -13,13 +15,29 @@ import (
 var cfgFile string
 var port int
 var ip string
+var geoIPDB string
 
 var rootCmd = &cobra.Command{
 	Use:   "rsocks",
 	Short: "Router socks",
 	Long: `Run a socks server and redirect the traffic to other socks server
 	according of the defined routes.`,
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		if geoIPDB != "" {
+			if _, err := os.Open(geoIPDB); err != nil {
+				return fmt.Errorf("cannot open the file %s: %s", geoIPDB, err.Error())
+			}
+		}
+		return nil
+	},
 	Run: func(cmd *cobra.Command, args []string) {
+		if geoIPDB != "" {
+			if err := router.LoadGeoIPDatabase(geoIPDB); err != nil {
+				log.Fatalf("error while loading the GeoIP database: %s", err.Error())
+			} else {
+				log.Printf("[*] GeoIP database %s loaded\n", geoIPDB)
+			}
+		}
 		socks.StartSocks(ip, port)
 		prompt.Prompt()
 	},
@@ -42,4 +60,8 @@ func init() {
 		&ip, "ip", "i", "0.0.0.0",
 		"IP for socks5 server",
 	)
+	rootCmd.Flags().StringVarP(
+		&geoIPDB, "geoip", "g", "",
+		"Path to the GeoIP database (GeoLite2-Country.mmdb)",
+	)
 }

go.mod

@@ -1,12 +1,22 @@
 module github.com/nodauf/Go-RouterSocks
 
-go 1.15
+go 1.22
 
 require (
 	github.com/c-bata/go-prompt v0.2.5
-	github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5
-	github.com/mitchellh/go-homedir v1.1.0
+	github.com/oschwald/geoip2-golang v1.9.0
+	github.com/oschwald/maxminddb-golang v1.11.0
 	github.com/spf13/cobra v1.2.1
-	github.com/spf13/viper v1.8.1
 	golang.org/x/net v0.0.0-20210902165921-8d991716f632
 )
+
+require (
+	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/mattn/go-colorable v0.1.7 // indirect
+	github.com/mattn/go-isatty v0.0.12 // indirect
+	github.com/mattn/go-runewidth v0.0.9 // indirect
+	github.com/mattn/go-tty v0.0.3 // indirect
+	github.com/pkg/term v1.1.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	golang.org/x/sys v0.9.0 // indirect
+)

prompt/completer.go

@@ -8,6 +8,7 @@ import (
 
 var commands = []prompt.Suggest{
 	{Text: "route", Description: "Manage routes. (default print them)"},
+	{Text: "geoip", Description: "List the supported GeoIP country"},
 	{Text: "chisel", Description: "List chisel port on this machine"},
 	{Text: "help", Description: "Help menu"},
 
@@ -23,9 +24,15 @@ var actionsRoute = []prompt.Suggest{
 	{Text: "import", Description: "Parse multiple route add command"},
 }
 
+var geoIPSubCommand = []prompt.Suggest{
+	{Text: "load", Description: "Load a GeoIP database"},
+	{Text: "print", Description: "Print all ISO codes with the associated country"},
+}
+
 var helpSubCommand = []prompt.Suggest{
 	{Text: "route", Description: "Add new route"},
 	{Text: "chisel", Description: "List chisel port on this machine"},
+	{Text: "geoip", Description: "List the supported GeoIP country"},
 }
 
 func complete(d prompt.Document) []prompt.Suggest {
@@ -45,6 +52,12 @@ func complete(d prompt.Document) []prompt.Suggest {
 			return prompt.FilterHasPrefix(actionsRoute, second, true)
 		}
 
+	case "geoip":
+		second := args[1]
+		if len(args) == 2 {
+			return prompt.FilterHasPrefix(geoIPSubCommand, second, true)
+		}
+
 	case "help":
 		second := args[1]
 		if len(args) == 2 {

prompt/help.go

@@ -10,14 +10,14 @@ func helpRoute() {
 }
 
 func helpRouteAdd() {
-	fmt.Println(`route add cidr socksServer:socksPort
-route add cidr chiselID
+	fmt.Println(`route add cidr/geoIPCountry socksServer:socksPort
+route add cidr/geoIPCountry chiselID
   route add 192.168.1.0/24 127.0.0.1:1081
   route add 192.168.1.0/24 0`)
 }
 
 func helpRouteDelete() {
-	fmt.Println(`route delete cidr
+	fmt.Println(`route delete cidr/geoIPCountry
   route delete 192.168.1.0/24`)
 }
 
@@ -27,8 +27,26 @@ func helpChisel() {
   [0] 127.0.0.1:1081`)
 }
 
+func helpGeoIP() {
+	fmt.Println("geoip usage: ")
+	helpGeoIPLoad()
+	fmt.Print("\n")
+	helpGeoPrint()
+}
+
+func helpGeoIPLoad() {
+	fmt.Println(`geoip load Path_to_the_database
+	geoip load /tmp/GeoLite2-Country.mmdb`)
+}
+func helpGeoPrint() {
+	fmt.Println(`geoip print
+  Output:
+  iso code => Country name`)
+}
+
 func help() {
 	fmt.Println("route: Manage route to socks servers")
 	fmt.Println("chisel: Liste chisel socks server on localhost")
+	fmt.Println("geoip: List the supported GeoIP country")
 	fmt.Println("help: help command")
 }

prompt/prompt.go

@@ -25,6 +25,8 @@ func executor(in string) {
 				helpRoute()
 			case "chisel":
 				helpChisel()
+			case "geoip":
+				helpGeoIP()
 			}
 		} else {
 			help()
@@ -41,15 +43,18 @@ func executor(in string) {
 				} else if serverSocks := utils.IsChiselIDValid(command[3]); serverSocks != "" {
 					remoteNetwork := command[2]
 					router.AddRoutes(remoteNetwork, serverSocks)
-				} else if !utils.IsCIDRValid(command[2]) {
-					fmt.Println("[-] CIDR is not valid")
+				} else if !utils.IsCIDRValid(command[2]) && !router.IsValidIsoCode(command[2]) {
+					fmt.Println("[-] CIDR or ISO code is not valid")
 					helpRouteAdd()
 				} else if !utils.IsRemoteSocksValid(command[3]) {
 					fmt.Println("[-] Socks server, socks port or chisel ID is not valid")
 					helpRouteAdd()
 				} else if !utils.CanResolvedHostname(strings.Split(command[3], ":")[0]) {
 					fmt.Println("[-] Server socks can be resolved")
 					helpRouteAdd()
+				} else if !utils.ServerReachable(command[3]) {
+					fmt.Println("[-] Server is not reachable")
+					helpRouteAdd()
 				} else {
 					remoteNetwork := command[2]
 					remoteSocks := command[3]
@@ -97,6 +102,28 @@ func executor(in string) {
 		}
 	case "chisel":
 		utils.PrintChiselProcess()
+	case "geoip":
+		if len(command) > 1 {
+			second := command[1]
+			switch strings.ToLower(second) {
+			case "load":
+
+				if len(command) != 3 {
+					helpGeoIPLoad()
+				} else {
+					err := router.LoadGeoIPDatabase(command[2])
+					if err != nil {
+						fmt.Printf("[-] Fail to load the database %s: %s\n", command[2], err.Error())
+					} else {
+						fmt.Printf("[*] GeoIP database %s loaded\n", command[2])
+					}
+				}
+			case "print":
+				router.PrintCountry()
+			}
+		} else {
+			helpGeoIP()
+		}
 	case "exit":
 		os.Exit(0)
 	case "":

router/geoip.go

@@ -0,0 +1,68 @@
+package router
+
+import (
+	"fmt"
+	"github.com/oschwald/geoip2-golang"
+	"github.com/oschwald/maxminddb-golang"
+	"log"
+	"slices"
+	"strings"
+)
+
+var geoIPDB *geoip2.Reader
+var geoIPDB_low *maxminddb.Reader
+var geoipEnable bool
+var isoCodes []string
+var isoWithCountries = make(map[string]string)
+
+//var isoWithNetworks = make(map[string][]string)
+
+func LoadGeoIPDatabase(geoipDB string) error {
+	var err error
+	geoIPDB, err = geoip2.Open(geoipDB)
+	if err != nil {
+		return err
+	}
+	geoIPDB_low, err = maxminddb.Open(geoipDB)
+	geoipEnable = true
+	// Load the countries and the iso code
+	var record struct {
+		Country struct {
+			ISOCode      string            `maxminddb:"iso_code"`
+			CountryNames map[string]string `maxminddb:"names"`
+		} `maxminddb:"country"`
+	} // Or any appropriate struct
+	networks := geoIPDB_low.Networks(maxminddb.SkipAliasedNetworks)
+	for networks.Next() {
+		_, err := networks.Network(&record)
+		if err != nil {
+			log.Fatal(err)
+		}
+		ISOCode := strings.ToLower(record.Country.ISOCode)
+		if !slices.Contains(isoCodes, ISOCode) {
+			isoCodes = append(isoCodes, ISOCode)
+		}
+		isoWithCountries[ISOCode] = record.Country.CountryNames["en"]
+		//isoWithNetworks[ISOCode] = append(isoWithNetworks[ISOCode], subnet.String())
+	}
+	if networks.Err() != nil {
+		log.Panic(networks.Err())
+	}
+	slices.Sort(isoCodes)
+	return err
+}
+
+func PrintCountry() {
+	if !geoipEnable {
+		fmt.Println("GeoIP database was not loaded")
+		return
+	}
+
+	for _, isoCode := range isoCodes {
+		fmt.Printf("%s => %s\n", isoCode, isoWithCountries[isoCode])
+	}
+}
+
+func IsValidIsoCode(ISOCode string) bool {
+	return slices.Contains(isoCodes, strings.ToLower(ISOCode))
+}

router/router.go

@@ -2,56 +2,79 @@ package router
 
 import (
 	"fmt"
-
 	utils "github.com/nodauf/Go-RouterSocks/utils"
+	"net"
+	"strings"
 )
 
 var Routes = map[string]string{}
 
-func AddRoutes(network string, remoteSocks string) {
-	if _, ok := Routes[network]; ok {
+func AddRoutes(destination string, remoteSocks string) {
+	destination = strings.ToLower(destination)
+	if _, ok := Routes[destination]; ok {
 		fmt.Println("[-] Route already present")
 	} else {
-		Routes[network] = remoteSocks
-		fmt.Println("[*] Successfull route added for network " + network)
+		Routes[destination] = remoteSocks
+		fmt.Println("[*] Successfull route added for network " + destination)
 	}
 }
 
-func DeleteRoutes(network string) {
-	if _, ok := Routes[network]; ok {
-		delete(Routes, network)
-		fmt.Println("[*] Successfull route " + network + " deleted")
+func DeleteRoutes(destination string) {
+	destination = strings.ToLower(destination)
+	if _, ok := Routes[destination]; ok {
+		delete(Routes, destination)
+		fmt.Println("[*] Successfull route " + destination + " deleted")
 	} else {
 		fmt.Println("[-] Route not found")
 	}
 }
 
 func PrintRoutes() {
-	for network, remoteSocks := range Routes {
-		fmt.Println(network + " => " + remoteSocks)
+	for destination, remoteSocks := range Routes {
+		fmt.Println(destination + " => " + remoteSocks)
 	}
 }
 
 func FlushRoutes() {
-	for network, _ := range Routes {
-		delete(Routes, network)
+	for destination, _ := range Routes {
+		delete(Routes, destination)
 	}
 
 	fmt.Println("[*] Successfull route flushed")
 }
 
 func DumpRoutes() {
-	for network, remoteSocks := range Routes {
-		fmt.Println("route add " + network + " " + remoteSocks)
+	for destination, remoteSocks := range Routes {
+		fmt.Println("route add " + destination + " " + remoteSocks)
 	}
 
 	fmt.Println("[*] Successfull route dumped")
 }
 
 func GetRoute(ip string) (string, string) {
-	for network, remoteSocks := range Routes {
-		if utils.CIDRContainsIP(network, ip) {
-			return network, remoteSocks
+	if geoipEnable {
+		ipNet := net.ParseIP(ip)
+		record, err := geoIPDB.Country(ipNet)
+		if err != nil {
+			fmt.Printf("[-] Fail to retrieve the country of %s\n", ip)
+		}
+		isoCode := strings.ToLower(record.Country.IsoCode)
+		if _, exist := Routes[isoCode]; exist {
+			return isoCode, Routes[isoCode]
+		}
+	}
+	for destination, remoteSocks := range Routes {
+		// if destination is iso code
+		if IsValidIsoCode(destination) {
+			continue
+			//for _, network := range isoWithNetworks[destination] {
+			//	if utils.CIDRContainsIP(network, ip) {
+			//		return destination, remoteSocks
+			//	}
+			//}
+
+		} else if utils.CIDRContainsIP(destination, ip) {
+			return destination, remoteSocks
 		}
 	}
 	return "", ""