Skip to content

Commit

Permalink
Merge pull request #28 from luthermonson/add-lookup
Browse files Browse the repository at this point in the history 
add ip/host lookup hashmap
  • Loading branch information
@luthermonson
luthermonson authored Dec 2, 2021
2 parents d0115ab + 4106b03 commit bcac70b
Show file tree
Hide file tree
Showing 37 changed files with 3,923 additions and 279 deletions.
2 changes: 2 additions & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@ go 1.17
require (
github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d
github.com/dimchansky/utfbom v1.1.1
github.com/icrowley/fake v0.0.0-20180203215853-4178557ae428
github.com/stretchr/testify v1.7.0
)

require (
github.com/corpix/uarand v0.1.1 // indirect
github.com/davecgh/go-spew v1.1.0 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
Expand Down
12 changes: 12 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,27 @@
github.com/Masterminds/glide v0.13.2/go.mod h1:STyF5vcenH/rUqTEv+/hBXlSTo7KYwg2oc2f4tzPWic=
github.com/Masterminds/semver v1.4.2/go.mod h1:MB6lktGJrhw8PrUyiEoblNEGEQ+RzHPF078ddwwvV3Y=
github.com/Masterminds/vcs v1.13.0/go.mod h1:N09YCmOQr6RLxC6UNHzuVwAdodYbbnycGHSmwVJjcKA=
github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d h1:Byv0BzEl3/e6D5CLfI0j/7hiIEtvGVFPCZ7Ei2oq8iQ=
github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
github.com/codegangsta/cli v1.20.0/go.mod h1:/qJNoX69yVSKu5o4jLyXAENLRyk1uhi7zkbQ3slBdOA=
github.com/corpix/uarand v0.1.1 h1:RMr1TWc9F4n5jiPDzFHtmaUXLKLNUFK0SgCLo4BhX/U=
github.com/corpix/uarand v0.1.1/go.mod h1:SFKZvkcRoLqVRFZ4u25xPmp6m9ktANfbpXZ7SJ0/FNU=
github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dimchansky/utfbom v1.1.1 h1:vV6w1AhK4VMnhBno/TPVCoK9U/LP0PkLCS9tbxHdi/U=
github.com/dimchansky/utfbom v1.1.1/go.mod h1:SxdoEBH5qIqFocHMyGOXVAybYJdr71b1Q/j0mACtrfE=
github.com/icrowley/fake v0.0.0-20180203215853-4178557ae428 h1:Mo9W14pwbO9VfRe+ygqZ8dFbPpoIK1HFrG/zjTuQ+nc=
github.com/icrowley/fake v0.0.0-20180203215853-4178557ae428/go.mod h1:uhpZMVGznybq1itEKXj6RYw9I71qK4kH+OGMjRC4KEo=
github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
github.com/ngdinhtoan/glide-cleanup v0.2.0/go.mod h1:UQzsmiDOb8YV3nOsCxK/c9zPpCZVNoHScRE3EO9pVMM=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
200 changes: 130 additions & 70 deletions hosts.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,18 +8,26 @@ import (
"os"
"path/filepath"
"sort"
"sync"

"github.com/asaskevich/govalidator"
"github.com/dimchansky/utfbom"
)

type lookup struct {
sync.RWMutex
l map[string][]int
}

type Hosts struct {
Path string
Lines []HostsLine
ips lookup
hosts lookup
}

// NewHosts return a new instance of ``Hosts`` using the default hosts file path.
func NewHosts() (Hosts, error) {
func NewHosts() (*Hosts, error) {
osHostsFilePath := os.ExpandEnv(filepath.FromSlash(HostsFilePath))

if env, isset := os.LookupEnv("HOSTS_PATH"); isset && len(env) > 0 {
Expand All @@ -30,9 +38,11 @@ func NewHosts() (Hosts, error) {
}

// NewCustomHosts return a new instance of ``Hosts`` using a custom hosts file path.
func NewCustomHosts(osHostsFilePath string) (Hosts, error) {
hosts := Hosts{
Path: osHostsFilePath,
func NewCustomHosts(osHostsFilePath string) (*Hosts, error) {
hosts := &Hosts{
Path: osHostsFilePath,
ips: lookup{l: make(map[string][]int)},
hosts: lookup{l: make(map[string][]int)},
}

if err := hosts.Load(); err != nil {
Expand Down Expand Up @@ -64,18 +74,20 @@ func (h *Hosts) Load() error {

scanner := bufio.NewScanner(utfbom.SkipOnly(file))
for scanner.Scan() {
h.Lines = append(h.Lines, NewHostsLine(scanner.Text()))
}

if err := scanner.Err(); err != nil {
return err
hl := NewHostsLine(scanner.Text())
h.Lines = append(h.Lines, hl)
pos := len(h.Lines) - 1
h.addIpPosition(hl.IP, pos)
for _, host := range hl.Hosts {
h.addHostPositions(host, pos)
}
}

return nil
return scanner.Err()
}

// Flush any changes made to hosts file.
func (h Hosts) Flush() error {
func (h *Hosts) Flush() error {
h.preFlushClean()
file, err := os.Create(h.Path)
if err != nil {
Expand All @@ -85,7 +97,6 @@ func (h Hosts) Flush() error {
defer file.Close()

w := bufio.NewWriter(file)

for _, line := range h.Lines {
if _, err := fmt.Fprintf(w, "%s%s", line.ToRaw(), eol); err != nil {
return err
Expand All @@ -100,40 +111,80 @@ func (h Hosts) Flush() error {
return h.Load()
}

// AddRaw takes a line from a hosts file and parses/adds the HostsLine
func (h *Hosts) AddRaw(raw ...string) error {
for _, r := range raw {
nl := NewHostsLine(r)
if nl.IP != "" && net.ParseIP(nl.IP) == nil {
return fmt.Errorf("%q is an invalid IP address", nl.IP)
}

for _, host := range nl.Hosts {
if !govalidator.IsDNSName(host) {
return fmt.Errorf("hostname is not a valid dns name: %s", host)
}
}

h.Lines = append(h.Lines, nl)
pos := len(h.Lines) - 1
h.addIpPosition(nl.IP, pos)
for _, host := range nl.Hosts {
h.addHostPositions(host, pos)
}
}

return nil
}

// Add an entry to the hosts file.
func (h *Hosts) Add(ip string, hosts ...string) error {
if net.ParseIP(ip) == nil {
return fmt.Errorf("%q is an invalid IP address", ip)
}

position := h.getIpPosition(ip)
if position == -1 {
// ip not already in hostsfile
h.Lines = append(h.Lines, HostsLine{
position := h.getIpPositions(ip)
if len(position) == 0 {
nl := HostsLine{
Raw: buildRawLine(ip, hosts),
IP: ip,
Hosts: hosts,
})
}
h.Lines = append(h.Lines, nl)
pos := len(h.Lines) - 1
h.addIpPosition(ip, pos)
for _, host := range nl.Hosts {
h.addHostPositions(host, pos)
}
} else {
// add new hosts to the correct position for the ip
hostsCopy := h.Lines[position].Hosts
// add new host to the first one we find
hostsCopy := h.Lines[position[0]].Hosts
for _, addHost := range hosts {
if h.Has(ip, addHost) {
// this combo already exists
continue
}

if !govalidator.IsDNSName(addHost) {
return fmt.Errorf("hostname is not a valid dns name: %s", addHost)
}
if itemInSlice(addHost, hostsCopy) {
if itemInSliceString(addHost, hostsCopy) {
continue // host exists for ip already
}

hostsCopy = append(hostsCopy, addHost)
h.addHostPositions(addHost, position[0])
}
h.Lines[position].Hosts = hostsCopy
h.Lines[position].Raw = h.Lines[position].ToRaw() // reset raw
h.Lines[position[0]].Hosts = hostsCopy
h.Lines[position[0]].Raw = h.Lines[position[0]].ToRaw() // reset raw
}

return nil
}

func (h *Hosts) Clear() {
h.Lines = []HostsLine{}
}

// Clean merge duplicate ips and hosts per ip
func (h *Hosts) Clean() {
h.RemoveDuplicateIps()
Expand All @@ -144,17 +195,26 @@ func (h *Hosts) Clean() {
}

// Has return a bool if ip/host combo in hosts file.
func (h Hosts) Has(ip string, host string) bool {
return h.getHostPosition(ip, host) != -1
func (h *Hosts) Has(ip string, host string) bool {
ippos := h.getIpPositions(ip)
hostpos := h.getHostPositions(host)
for _, pos := range ippos {
if itemInSliceInt(pos, hostpos) {
// if ip and host have matching lookup positions we have a combo match
return true
}
}

return false
}

// HasHostname return a bool if hostname in hosts file.
func (h Hosts) HasHostname(host string) bool {
return h.getHostnamePosition(host) != -1
func (h *Hosts) HasHostname(host string) bool {
return len(h.getHostPositions(host)) > 0
}

func (h Hosts) HasIp(ip string) bool {
return h.getIpPosition(ip) != -1
func (h *Hosts) HasIp(ip string) bool {
return len(h.getIpPositions(ip)) > 0
}

// Remove an entry from the hosts file.
Expand All @@ -173,7 +233,7 @@ func (h *Hosts) Remove(ip string, hosts ...string) error {

var newHosts []string
for _, checkHost := range line.Hosts {
if !itemInSlice(checkHost, hosts) {
if !itemInSliceString(checkHost, hosts) {
newHosts = append(newHosts, checkHost)
}
}
Expand All @@ -196,27 +256,24 @@ func (h *Hosts) Remove(ip string, hosts ...string) error {

// RemoveByHostname remove entries by hostname from the hosts file.
func (h *Hosts) RemoveByHostname(host string) error {
newLines := []HostsLine{}
for _, line := range h.Lines {
for _, p := range h.getHostPositions(host) {
line := &h.Lines[p]
if len(line.Hosts) > 0 {
line.Hosts = removeFromSlice(host, line.Hosts)
line.Hosts = removeFromSliceString(host, line.Hosts)
line.RegenRaw()
}

if len(line.Hosts) > 0 {
newLines = append(newLines, line)
}
h.removeHostPositions(host, p)
}
h.Lines = newLines

return nil
}

func (h *Hosts) RemoveByIp(ip string) error {
pos := h.getIpPosition(ip)
for pos > -1 {
h.removeByPosition(pos)
pos = h.getIpPosition(ip)
pos := h.getIpPositions(ip)
for len(pos) > 0 {
for _, p := range pos {
h.removeByPosition(p)
}
}

return nil
Expand Down Expand Up @@ -315,8 +372,8 @@ func (h *Hosts) combineIp(ip string) {
}

func (h *Hosts) removeByPosition(pos int) {
if len(h.Lines) == 1 {
h.Lines = []HostsLine{}
if pos == 0 && len(h.Lines) == 1 {
h.Clear()
return
}
if pos == len(h.Lines) {
Expand All @@ -337,41 +394,44 @@ func (h *Hosts) removeIp(ip string) {
h.Lines = newLines
}

func (h Hosts) getHostPosition(ip string, host string) int {
for i := range h.Lines {
line := h.Lines[i]
if !line.IsComment() && line.Raw != "" {
if ip == line.IP && itemInSlice(host, line.Hosts) {
return i
}
}
func (h *Hosts) getHostPositions(host string) []int {
h.hosts.RLock()
defer h.hosts.RUnlock()
i, ok := h.hosts.l[host]
if ok {
return i
}

return -1
return []int{}
}

func (h Hosts) getHostnamePosition(host string) int {
for i := range h.Lines {
line := h.Lines[i]
if !line.IsComment() && line.Raw != "" {
if itemInSlice(host, line.Hosts) {
return i
}
}
}
func (h *Hosts) addHostPositions(host string, pos int) {
h.hosts.Lock()
defer h.hosts.Unlock()
h.hosts.l[host] = append(h.hosts.l[host], pos)
}

return -1
func (h *Hosts) removeHostPositions(host string, pos int) {
h.hosts.Lock()
defer h.hosts.Unlock()
positions := h.hosts.l[host]
h.hosts.l[host] = removeFromSliceInt(pos, positions)
}

func (h Hosts) getIpPosition(ip string) int {
for i := range h.Lines {
line := h.Lines[i]
if !line.IsComment() && line.Raw != "" && line.IP == ip {
return i
}
func (h *Hosts) getIpPositions(ip string) []int {
h.ips.RLock()
defer h.ips.RUnlock()
i, ok := h.ips.l[ip]
if ok {
return i
}

return -1
return []int{}
}

func (h *Hosts) addIpPosition(ip string, pos int) {
h.ips.Lock()
defer h.ips.Unlock()
h.ips.l[ip] = append(h.ips.l[ip], pos)
}

func buildRawLine(ip string, hosts []string) string {
Expand Down
Loading

0 comments on commit bcac70b

Please sign in to comment.