v2 init commit

src/mod/geodb/geodb.go

@@ -0,0 +1,250 @@
+package geodb
+
+import (
+	_ "embed"
+	"log"
+	"net"
+	"net/http"
+	"strings"
+
+	"imuslab.com/zoraxy/mod/database"
+)
+
+//go:embed geoipv4.csv
+var geoipv4 []byte //Original embedded csv file
+
+type Store struct {
+	Enabled bool
+	geodb   [][]string //Parsed geodb list
+	//geoipCache sync.Map
+	geotrie *trie
+	sysdb   *database.Database
+}
+
+type CountryInfo struct {
+	CountryIsoCode string
+	ContinetCode   string
+}
+
+func NewGeoDb(sysdb *database.Database) (*Store, error) {
+	parsedGeoData, err := parseCSV(geoipv4)
+	if err != nil {
+		return nil, err
+	}
+
+	blacklistEnabled := false
+	if sysdb != nil {
+		err = sysdb.NewTable("blacklist-cn")
+		if err != nil {
+			return nil, err
+		}
+
+		err = sysdb.NewTable("blacklist-ip")
+		if err != nil {
+			return nil, err
+		}
+
+		err = sysdb.NewTable("blacklist")
+		if err != nil {
+			return nil, err
+		}
+		sysdb.Read("blacklist", "enabled", &blacklistEnabled)
+	} else {
+		log.Println("Database pointer set to nil: Entering debug mode")
+	}
+
+	return &Store{
+		Enabled: blacklistEnabled,
+		geodb:   parsedGeoData,
+		//geoipCache: sync.Map{},
+		geotrie: constrctTrieTree(parsedGeoData),
+		sysdb:   sysdb,
+	}, nil
+}
+
+func (s *Store) ToggleBlacklist(enabled bool) {
+	s.sysdb.Write("blacklist", "enabled", enabled)
+	s.Enabled = enabled
+}
+
+func (s *Store) ResolveCountryCodeFromIP(ipstring string) (*CountryInfo, error) {
+	cc := s.search(ipstring)
+	return &CountryInfo{
+		CountryIsoCode: cc,
+		ContinetCode:   "",
+	}, nil
+}
+
+func (s *Store) Close() {
+
+}
+
+func (s *Store) AddCountryCodeToBlackList(countryCode string) {
+	countryCode = strings.ToLower(countryCode)
+	s.sysdb.Write("blacklist-cn", countryCode, true)
+}
+
+func (s *Store) RemoveCountryCodeFromBlackList(countryCode string) {
+	countryCode = strings.ToLower(countryCode)
+	s.sysdb.Delete("blacklist-cn", countryCode)
+}
+
+func (s *Store) IsCountryCodeBlacklisted(countryCode string) bool {
+	countryCode = strings.ToLower(countryCode)
+	var isBlacklisted bool = false
+	s.sysdb.Read("blacklist-cn", countryCode, &isBlacklisted)
+	return isBlacklisted
+}
+
+func (s *Store) GetAllBlacklistedCountryCode() []string {
+	bannedCountryCodes := []string{}
+	entries, err := s.sysdb.ListTable("blacklist-cn")
+	if err != nil {
+		return bannedCountryCodes
+	}
+	for _, keypairs := range entries {
+		ip := string(keypairs[0])
+		bannedCountryCodes = append(bannedCountryCodes, ip)
+	}
+
+	return bannedCountryCodes
+}
+
+func (s *Store) AddIPToBlackList(ipAddr string) {
+	s.sysdb.Write("blacklist-ip", ipAddr, true)
+}
+
+func (s *Store) RemoveIPFromBlackList(ipAddr string) {
+	s.sysdb.Delete("blacklist-ip", ipAddr)
+}
+
+func (s *Store) IsIPBlacklisted(ipAddr string) bool {
+	var isBlacklisted bool = false
+	s.sysdb.Read("blacklist-ip", ipAddr, &isBlacklisted)
+	if isBlacklisted {
+		return true
+	}
+
+	//Check for IP wildcard and CIRD rules
+	AllBlacklistedIps := s.GetAllBlacklistedIp()
+	for _, blacklistRule := range AllBlacklistedIps {
+		wildcardMatch := MatchIpWildcard(ipAddr, blacklistRule)
+		if wildcardMatch {
+			return true
+		}
+
+		cidrMatch := MatchIpCIDR(ipAddr, blacklistRule)
+		if cidrMatch {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (s *Store) GetAllBlacklistedIp() []string {
+	bannedIps := []string{}
+	entries, err := s.sysdb.ListTable("blacklist-ip")
+	if err != nil {
+		return bannedIps
+	}
+
+	for _, keypairs := range entries {
+		ip := string(keypairs[0])
+		bannedIps = append(bannedIps, ip)
+	}
+
+	return bannedIps
+}
+
+// Check if a IP address is blacklisted, in either country or IP blacklist
+func (s *Store) IsBlacklisted(ipAddr string) bool {
+	if !s.Enabled {
+		//Blacklist not enabled. Always return false
+		return false
+	}
+
+	if ipAddr == "" {
+		//Unable to get the target IP address
+		return false
+	}
+
+	countryCode, err := s.ResolveCountryCodeFromIP(ipAddr)
+	if err != nil {
+		return false
+	}
+
+	if s.IsCountryCodeBlacklisted(countryCode.CountryIsoCode) {
+		return true
+	}
+
+	if s.IsIPBlacklisted(ipAddr) {
+		return true
+	}
+
+	return false
+}
+
+func (s *Store) GetRequesterCountryISOCode(r *http.Request) string {
+	ipAddr := GetRequesterIP(r)
+	if ipAddr == "" {
+		return ""
+	}
+	countryCode, err := s.ResolveCountryCodeFromIP(ipAddr)
+	if err != nil {
+		return ""
+	}
+
+	return countryCode.CountryIsoCode
+}
+
+// Utilities function
+func GetRequesterIP(r *http.Request) string {
+	ip := r.Header.Get("X-Forwarded-For")
+	if ip == "" {
+		ip = r.Header.Get("X-Real-IP")
+		if ip == "" {
+			ip = strings.Split(r.RemoteAddr, ":")[0]
+		}
+	}
+	return ip
+}
+
+// Match the IP address with a wildcard string
+func MatchIpWildcard(ipAddress, wildcard string) bool {
+	// Split IP address and wildcard into octets
+	ipOctets := strings.Split(ipAddress, ".")
+	wildcardOctets := strings.Split(wildcard, ".")
+
+	// Check that both have 4 octets
+	if len(ipOctets) != 4 || len(wildcardOctets) != 4 {
+		return false
+	}
+
+	// Check each octet to see if it matches the wildcard or is an exact match
+	for i := 0; i < 4; i++ {
+		if wildcardOctets[i] == "*" {
+			continue
+		}
+		if ipOctets[i] != wildcardOctets[i] {
+			return false
+		}
+	}
+
+	return true
+}
+
+// Match ip address with CIDR
+func MatchIpCIDR(ip string, cidr string) bool {
+	// parse the CIDR string
+	_, cidrnet, err := net.ParseCIDR(cidr)
+	if err != nil {
+		return false
+	}
+
+	// parse the IP address
+	ipAddr := net.ParseIP(ip)
+
+	// check if the IP address is within the CIDR range
+	return cidrnet.Contains(ipAddr)
+}

src/mod/geodb/geodb_test.go

@@ -0,0 +1,73 @@
+package geodb_test
+
+import (
+	"testing"
+
+	"imuslab.com/zoraxy/mod/geodb"
+)
+
+/*
+func TestTrieConstruct(t *testing.T) {
+	tt := geodb.NewTrie()
+	data := [][]string{
+		{"1.0.16.0", "1.0.31.255", "JP"},
+		{"1.0.32.0", "1.0.63.255", "CN"},
+		{"1.0.64.0", "1.0.127.255", "JP"},
+		{"1.0.128.0", "1.0.255.255", "TH"},
+		{"1.1.0.0", "1.1.0.255", "CN"},
+		{"1.1.1.0", "1.1.1.255", "AU"},
+		{"1.1.2.0", "1.1.63.255", "CN"},
+		{"1.1.64.0", "1.1.127.255", "JP"},
+		{"1.1.128.0", "1.1.255.255", "TH"},
+		{"1.2.0.0", "1.2.2.255", "CN"},
+		{"1.2.3.0", "1.2.3.255", "AU"},
+	}
+
+	for _, entry := range data {
+		startIp := entry[0]
+		endIp := entry[1]
+		cc := entry[2]
+		tt.Insert(startIp, cc)
+		tt.Insert(endIp, cc)
+	}
+
+	t.Log(tt.Search("1.0.16.20"), "== JP")  //JP
+	t.Log(tt.Search("1.2.0.122"), "== CN")  //CN
+	t.Log(tt.Search("1.2.1.0"), "== CN")    //CN
+	t.Log(tt.Search("1.0.65.243"), "== JP") //JP
+	t.Log(tt.Search("1.0.62.243"), "== CN") //CN
+}
+*/
+
+func TestResolveCountryCodeFromIP(t *testing.T) {
+	// Create a new store
+	store, err := geodb.NewGeoDb(nil)
+	if err != nil {
+		t.Errorf("error creating store: %v", err)
+		return
+	}
+
+	// Test an IP address that should return a valid country code
+	ip := "8.8.8.8"
+	expected := "US"
+	info, err := store.ResolveCountryCodeFromIP(ip)
+	if err != nil {
+		t.Errorf("error resolving country code for IP %s: %v", ip, err)
+		return
+	}
+	if info.CountryIsoCode != expected {
+		t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip)
+	}
+
+	// Test an IP address that should return an empty country code
+	ip = "127.0.0.1"
+	expected = ""
+	info, err = store.ResolveCountryCodeFromIP(ip)
+	if err != nil {
+		t.Errorf("error resolving country code for IP %s: %v", ip, err)
+		return
+	}
+	if info.CountryIsoCode != expected {
+		t.Errorf("expected country code %s, but got %s for IP %s", expected, info.CountryIsoCode, ip)
+	}
+}

src/mod/geodb/geoloader.go

@@ -0,0 +1,89 @@
+package geodb
+
+import (
+	"bytes"
+	"encoding/csv"
+	"io"
+	"net"
+	"strings"
+)
+
+func (s *Store) search(ip string) string {
+	if strings.Contains(ip, ",") {
+		//This is a CF proxied request. We only need the front part
+		//Example 219.71.102.145, 172.71.139.178
+		ip = strings.Split(ip, ",")[0]
+		ip = strings.TrimSpace(ip)
+	}
+	//See if there are cached country code for this ip
+	/*
+		ccc, ok := s.geoipCache.Load(ip)
+		if ok {
+			return ccc.(string)
+		}
+	*/
+
+	//Search in geotrie tree
+	cc := s.geotrie.search(ip)
+	/*
+		if cc != "" {
+			s.geoipCache.Store(ip, cc)
+		}
+	*/
+	return cc
+}
+
+// Construct the trie data structure for quick lookup
+func constrctTrieTree(data [][]string) *trie {
+	tt := newTrie()
+	for _, entry := range data {
+		startIp := entry[0]
+		endIp := entry[1]
+		cc := entry[2]
+		tt.insert(startIp, cc)
+		tt.insert(endIp, cc)
+	}
+
+	return tt
+}
+
+// Parse the embedded csv as ipstart, ipend and country code entries
+func parseCSV(content []byte) ([][]string, error) {
+	var records [][]string
+	r := csv.NewReader(bytes.NewReader(content))
+	for {
+		record, err := r.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		records = append(records, record)
+	}
+	return records, nil
+}
+
+// Check if a ip string is within the range of two others
+func isIPInRange(ip, start, end string) bool {
+	ipAddr := net.ParseIP(ip)
+	if ipAddr == nil {
+		return false
+	}
+
+	startAddr := net.ParseIP(start)
+	if startAddr == nil {
+		return false
+	}
+
+	endAddr := net.ParseIP(end)
+	if endAddr == nil {
+		return false
+	}
+
+	if ipAddr.To4() == nil || startAddr.To4() == nil || endAddr.To4() == nil {
+		return false
+	}
+
+	return bytes.Compare(ipAddr.To4(), startAddr.To4()) >= 0 && bytes.Compare(ipAddr.To4(), endAddr.To4()) <= 0
+}

src/mod/geodb/trie.go

@@ -0,0 +1,131 @@
+package geodb
+
+import (
+	"fmt"
+	"net"
+	"strconv"
+	"strings"
+)
+
+type trie_Node struct {
+	childrens [2]*trie_Node
+	ends      bool
+	cc        string
+}
+
+// Initializing the root of the trie
+type trie struct {
+	root *trie_Node
+}
+
+func ipToBitString(ip string) string {
+	// Parse the IP address string into a net.IP object
+	parsedIP := net.ParseIP(ip)
+
+	// Convert the IP address to a 4-byte slice
+	ipBytes := parsedIP.To4()
+
+	// Convert each byte in the IP address to its 8-bit binary representation
+	var result []string
+	for _, b := range ipBytes {
+		result = append(result, fmt.Sprintf("%08b", b))
+	}
+
+	// Join the binary representation of each byte with dots to form the final bit string
+	return strings.Join(result, "")
+}
+
+func bitStringToIp(bitString string) string {
+	// Split the bit string into four 8-bit segments
+	segments := []string{
+		bitString[:8],
+		bitString[8:16],
+		bitString[16:24],
+		bitString[24:32],
+	}
+
+	// Convert each segment to its decimal equivalent
+	var decimalSegments []int
+	for _, s := range segments {
+		i, _ := strconv.ParseInt(s, 2, 64)
+		decimalSegments = append(decimalSegments, int(i))
+	}
+
+	// Join the decimal segments with dots to form the IP address string
+	return fmt.Sprintf("%d.%d.%d.%d", decimalSegments[0], decimalSegments[1], decimalSegments[2], decimalSegments[3])
+}
+
+// inititlaizing a new trie
+func newTrie() *trie {
+	t := new(trie)
+	t.root = new(trie_Node)
+	return t
+}
+
+// Passing words to trie
+func (t *trie) insert(ipAddr string, cc string) {
+	word := ipToBitString(ipAddr)
+	current := t.root
+	for _, wr := range word {
+		index := wr - '0'
+		if current.childrens[index] == nil {
+			current.childrens[index] = &trie_Node{
+				childrens: [2]*trie_Node{},
+				ends:      false,
+				cc:        cc,
+			}
+		}
+		current = current.childrens[index]
+	}
+	current.ends = true
+}
+
+func isReservedIP(ip string) bool {
+	parsedIP := net.ParseIP(ip)
+	if parsedIP == nil {
+		return false
+	}
+	// Check if the IP address is a loopback address
+	if parsedIP.IsLoopback() {
+		return true
+	}
+	// Check if the IP address is in the link-local address range
+	if parsedIP.IsLinkLocalUnicast() || parsedIP.IsLinkLocalMulticast() {
+		return true
+	}
+	// Check if the IP address is in the private address ranges
+	privateRanges := []*net.IPNet{
+		{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(8, 32)},
+		{IP: net.ParseIP("172.16.0.0"), Mask: net.CIDRMask(12, 32)},
+		{IP: net.ParseIP("192.168.0.0"), Mask: net.CIDRMask(16, 32)},
+	}
+	for _, r := range privateRanges {
+		if r.Contains(parsedIP) {
+			return true
+		}
+	}
+	// If the IP address is not a reserved address, return false
+	return false
+}
+
+// Initializing the search for word in node
+func (t *trie) search(ipAddr string) string {
+	if isReservedIP(ipAddr) {
+		return ""
+	}
+	word := ipToBitString(ipAddr)
+	current := t.root
+	for _, wr := range word {
+		index := wr - '0'
+		if current.childrens[index] == nil {
+			return current.cc
+		}
+		current = current.childrens[index]
+	}
+	if current.ends {
+		return current.cc
+	}
+
+	//Not found
+	return ""
+}