100 lines
1.6 KiB
Go
100 lines
1.6 KiB
Go
package netutil
|
|
|
|
import (
|
|
"strings"
|
|
|
|
"github.com/miekg/dns"
|
|
)
|
|
|
|
type DomainTree struct {
|
|
root *domainTreeNode
|
|
}
|
|
|
|
type domainTreeNode struct {
|
|
leaf map[string]*domainTreeNode
|
|
isEnd bool
|
|
}
|
|
|
|
func NewDomainList(domains ...string) *DomainTree {
|
|
tree := &DomainTree{
|
|
root: &domainTreeNode{leaf: make(map[string]*domainTreeNode)},
|
|
}
|
|
for _, domain := range domains {
|
|
tree.Add(domain)
|
|
}
|
|
return tree
|
|
}
|
|
|
|
func (tree *DomainTree) Add(domain string) {
|
|
domain = normalizeDomain(domain)
|
|
if domain == "" {
|
|
return
|
|
}
|
|
|
|
labels := dns.SplitDomainName(domain)
|
|
if len(labels) == 0 {
|
|
return
|
|
}
|
|
|
|
node := tree.root
|
|
for i := len(labels) - 1; i >= 0; i-- {
|
|
label := labels[i]
|
|
if label == "" {
|
|
continue
|
|
}
|
|
if node.leaf == nil {
|
|
node.leaf = make(map[string]*domainTreeNode)
|
|
}
|
|
if node.leaf[label] == nil {
|
|
node.leaf[label] = &domainTreeNode{}
|
|
}
|
|
node = node.leaf[label]
|
|
}
|
|
node.isEnd = true
|
|
}
|
|
|
|
func (tree *DomainTree) Contains(domain string) bool {
|
|
domain = normalizeDomain(domain)
|
|
if domain == "" {
|
|
return false
|
|
}
|
|
|
|
labels := dns.SplitDomainName(domain)
|
|
if len(labels) == 0 {
|
|
return false
|
|
}
|
|
|
|
node := tree.root
|
|
for i := len(labels) - 1; i >= 0; i-- {
|
|
if node.isEnd {
|
|
return true
|
|
}
|
|
|
|
if node.leaf == nil {
|
|
return false
|
|
}
|
|
|
|
label := labels[i]
|
|
if node = node.leaf[label]; node == nil {
|
|
return false
|
|
}
|
|
}
|
|
return node.isEnd
|
|
}
|
|
|
|
func normalizeDomain(domain string) string {
|
|
domain = strings.ToLower(strings.TrimSpace(domain))
|
|
if domain == "" {
|
|
return ""
|
|
}
|
|
|
|
// Remove trailing dot if present, dns.Fqdn will add it back properly
|
|
domain = strings.TrimSuffix(domain, ".")
|
|
|
|
if domain == "" {
|
|
return ""
|
|
}
|
|
|
|
return dns.Fqdn(domain)
|
|
}
|