45 lines
830 B
Go
45 lines
830 B
Go
package netutil
|
|
|
|
import (
|
|
"net"
|
|
|
|
"github.com/yl2chen/cidranger"
|
|
)
|
|
|
|
type NetworkTree struct {
|
|
ranger cidranger.Ranger
|
|
}
|
|
|
|
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
|
|
}
|