80 lines
1.4 KiB
Go
80 lines
1.4 KiB
Go
package dataset
|
|
|
|
import (
|
|
"net"
|
|
"net/netip"
|
|
|
|
"git.maze.io/maze/styx/dataset/nettrie"
|
|
"github.com/yl2chen/cidranger"
|
|
)
|
|
|
|
type NetworkTrie struct {
|
|
*nettrie.Trie
|
|
}
|
|
|
|
func NewNetworkTrie(prefixes ...netip.Prefix) *NetworkTrie {
|
|
trie := &NetworkTrie{
|
|
Trie: nettrie.New(),
|
|
}
|
|
for _, prefix := range prefixes {
|
|
trie.Insert(prefix)
|
|
}
|
|
return trie
|
|
}
|
|
|
|
func (trie *NetworkTrie) ContainsIP(ip net.IP) bool {
|
|
if ip == nil {
|
|
return false
|
|
}
|
|
addr, ok := netip.AddrFromSlice(ip)
|
|
if !ok {
|
|
return false
|
|
}
|
|
return trie.Contains(addr)
|
|
}
|
|
|
|
type NetworkTree struct {
|
|
ranger cidranger.Ranger
|
|
}
|
|
|
|
func MustNetworkTree(networks ...string) *NetworkTree {
|
|
tree, err := NewNetworkTree(networks...)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return tree
|
|
}
|
|
|
|
func NewNetworkTree(networks ...string) (*NetworkTree, error) {
|
|
tree := &NetworkTree{
|
|
ranger: cidranger.NewPCTrieRanger(),
|
|
}
|
|
for _, cidr := range networks {
|
|
if err := tree.AddCIDR(cidr); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return tree, nil
|
|
}
|
|
|
|
func (tree *NetworkTree) Add(ipnet *net.IPNet) {
|
|
if ipnet == nil {
|
|
return
|
|
}
|
|
tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
|
|
}
|
|
|
|
func (tree *NetworkTree) AddCIDR(cidr string) error {
|
|
_, ipnet, err := net.ParseCIDR(cidr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
|
|
return nil
|
|
}
|
|
|
|
func (tree *NetworkTree) Contains(ip net.IP) bool {
|
|
contains, _ := tree.ranger.Contains(ip)
|
|
return contains
|
|
}
|