Zoraxy/mod/access/whitelist.go

package access

import (
	"encoding/json"
	"strings"

	"imuslab.com/zoraxy/mod/netutils"
)

/*
	Whitelist.go

	This script handles whitelist related functions
*/

const (
	EntryType_CountryCode int = 0
	EntryType_IP          int = 1
)

type WhitelistEntry struct {
	EntryType int    //Entry type of whitelist, Country Code or IP
	CC        string //ISO Country Code
	IP        string //IP address or range
	Comment   string //Comment for this entry
}

//Geo Whitelist

func (s *AccessRule) AddCountryCodeToWhitelist(countryCode string, comment string) {
	countryCode = strings.ToLower(countryCode)
	entry := WhitelistEntry{
		EntryType: EntryType_CountryCode,
		CC:        countryCode,
		Comment:   comment,
	}

	s.GetDatabase().Write(s.GetFullTableName("whitelist-cn"), countryCode, entry)
}

func (s *AccessRule) RemoveCountryCodeFromWhitelist(countryCode string) {
	countryCode = strings.ToLower(countryCode)
	s.GetDatabase().Delete(s.GetFullTableName("whitelist-cn"), countryCode)
}

func (s *AccessRule) IsCountryCodeWhitelisted(countryCode string) bool {
	countryCode = strings.ToLower(countryCode)
	return s.GetDatabase().KeyExists(s.GetFullTableName("whitelist-cn"), countryCode)
}

func (s *AccessRule) GetAllWhitelistedCountryCode() []*WhitelistEntry {
	whitelistedCountryCode := []*WhitelistEntry{}
	entries, err := s.GetDatabase().ListTable(s.GetFullTableName("whitelist-cn"))
	if err != nil {
		return whitelistedCountryCode
	}
	for _, keypairs := range entries {
		thisWhitelistEntry := WhitelistEntry{}
		json.Unmarshal(keypairs[1], &thisWhitelistEntry)
		whitelistedCountryCode = append(whitelistedCountryCode, &thisWhitelistEntry)
	}

	return whitelistedCountryCode
}

//IP Whitelist

func (s *AccessRule) AddIPToWhiteList(ipAddr string, comment string) {
	thisIpEntry := WhitelistEntry{
		EntryType: EntryType_IP,
		IP:        ipAddr,
		Comment:   comment,
	}

	s.GetDatabase().Write(s.GetFullTableName("whitelist-ip"), ipAddr, thisIpEntry)
}

func (s *AccessRule) RemoveIPFromWhiteList(ipAddr string) {
	s.GetDatabase().Delete(s.GetFullTableName("whitelist-ip"), ipAddr)
}

func (s *AccessRule) IsIPWhitelisted(ipAddr string) bool {
	isWhitelisted := s.GetDatabase().KeyExists(s.GetFullTableName("whitelist-ip"), ipAddr)
	if isWhitelisted {
		//single IP whitelist entry
		return true
	}

	//Check for IP wildcard and CIRD rules
	AllWhitelistedIps := s.GetAllWhitelistedIpAsStringSlice()
	for _, whitelistRules := range AllWhitelistedIps {
		wildcardMatch := netutils.MatchIpWildcard(ipAddr, whitelistRules)
		if wildcardMatch {
			return true
		}

		cidrMatch := netutils.MatchIpCIDR(ipAddr, whitelistRules)
		if cidrMatch {
			return true
		}
	}

	return false
}

func (s *AccessRule) GetAllWhitelistedIp() []*WhitelistEntry {
	whitelistedIp := []*WhitelistEntry{}
	entries, err := s.GetDatabase().ListTable(s.GetFullTableName("whitelist-ip"))
	if err != nil {
		return whitelistedIp
	}

	for _, keypairs := range entries {
		//ip := string(keypairs[0])
		thisEntry := WhitelistEntry{}
		json.Unmarshal(keypairs[1], &thisEntry)
		whitelistedIp = append(whitelistedIp, &thisEntry)
	}

	return whitelistedIp
}

func (s *AccessRule) GetAllWhitelistedIpAsStringSlice() []string {
	allWhitelistedIPs := []string{}
	entries := s.GetAllWhitelistedIp()
	for _, entry := range entries {
		allWhitelistedIPs = append(allWhitelistedIPs, entry.IP)
	}

	return allWhitelistedIPs
}