package hosts import ( "bufio" "bytes" "fmt" "io" "net" "os" "runtime" "strings" "sync" ) var lineBreaker string func init() { if runtime.GOOS == "windows" { lineBreaker = "\r\n" } else { lineBreaker = "\n" } } func SystemHostpath() string { if runtime.GOOS == "windows" { return `C:\Windows\System32\drivers\etc\hosts` } return `/etc/hosts` } var defaultParser = NewHosts() func ParseHosts(hostPath string) error { return defaultParser.Parse(hostPath) } func Parse() error { return defaultParser.Parse(SystemHostpath()) } func List() []*HostNode { return defaultParser.List() } func ListHosts() map[string][]string { return defaultParser.ListHosts() } func ListIPs() map[string][]string { return defaultParser.ListIPs() } func ListByHost(host string) []*HostNode { return defaultParser.ListByHost(host) } func ListIPsByHost(host string) []string { return defaultParser.ListIPsByHost(host) } func ListFirstIPByHost(host string) string { return defaultParser.ListFirstIPByHost(host) } func ListByIP(ip string) []*HostNode { return defaultParser.ListByIP(ip) } func ListHostsByIP(ip string) []string { return defaultParser.ListHostsByIP(ip) } func ListFirstHostByIP(ip string) string { return defaultParser.ListFirstHostByIP(ip) } func AddHosts(ip string, hosts ...string) error { return defaultParser.AddHosts(ip, hosts...) } func AddHostsComment(comment string, ip string, hosts ...string) error { return defaultParser.AddHostsComment(comment, ip, hosts...) } func RemoveIPHosts(ip string, hosts ...string) error { return defaultParser.RemoveIPHosts(ip, hosts...) } func RemoveIPs(ips ...string) error { return defaultParser.RemoveIPs(ips...) } func RemoveHosts(hosts ...string) error { return defaultParser.RemoveHosts(hosts...) } func SetIPHosts(ip string, hosts ...string) error { return defaultParser.SetIPHosts(ip, hosts...) } func SetHostIPs(host string, ips ...string) error { return defaultParser.SetHostIPs(host, ips...) } func InsertNodes(node *HostNode, before bool, comment string, ip string, hosts ...string) error { return defaultParser.InsertNodeByData(node, before, comment, ip, hosts...) } func SaveAs(path string) error { return defaultParser.SaveAs(path) } func Save() error { return defaultParser.Save() } func Build() ([]byte, error) { return defaultParser.Build() } func GetFirstNode() (*HostNode, error) { return defaultParser.GetFirstNode() } func GetNode(nodeID uint64) (*HostNode, error) { return defaultParser.GetNode(nodeID) } func GetNextNode(nodeID uint64) (*HostNode, error) { return defaultParser.GetNextNode(nodeID) } func GetLastNode(nodeID uint64) (*HostNode, error) { return defaultParser.GetLastNode(nodeID) } func GetLatestNode() (*HostNode, error) { return defaultParser.GetLatestNode() } func DeleteNode(node *HostNode) error { return defaultParser.DeleteNode(node) } func DeleteNodeByUID(uid uint64) error { return defaultParser.DeleteNodeByUID(uid) } func AddNode(node *HostNode) error { return defaultParser.AddNode(node) } func InsertNode(node *HostNode) error { return defaultParser.InsertNode(node) } func UpdateNode(node *HostNode) error { return defaultParser.UpdateNode(node) } type HostNode struct { uid uint64 nextuid uint64 lastuid uint64 ip string host []string comment string original string onlyComment bool valid bool } type Host struct { idx uint64 hostPath string firstUid uint64 lastUid uint64 fulldata map[uint64]*HostNode hostData map[string][]*HostNode ipData map[string][]*HostNode sync.RWMutex } func NewHosts() *Host { return &Host{ hostData: make(map[string][]*HostNode), ipData: make(map[string][]*HostNode), fulldata: make(map[uint64]*HostNode), } } func (h *Host) Parse(hostPath string) error { h.Lock() defer h.Unlock() h.idx = 0 h.firstUid = 0 h.lastUid = 0 h.hostPath = hostPath h.fulldata = make(map[uint64]*HostNode) h.hostData = make(map[string][]*HostNode) h.ipData = make(map[string][]*HostNode) return h.parse() } func (h *Host) parse() error { f, err := os.OpenFile(h.hostPath, os.O_RDONLY, 0666) if err != nil { return fmt.Errorf("open hosts file %s error: %s", h.hostPath, err) } defer f.Close() buf := bufio.NewReader(f) for { line, err := buf.ReadString('\n') if err == io.EOF { if h.idx-1 >= 0 { h.fulldata[h.idx].nextuid = 0 h.lastUid = h.idx } break } if err != nil { return fmt.Errorf("read hosts file error: %s", err) } h.idx++ line = strings.TrimSpace(line) data, _ := h.parseLine(line) data.uid = h.idx data.lastuid = h.idx - 1 data.nextuid = h.idx + 1 h.fulldata[data.uid] = &data if h.firstUid == 0 { h.firstUid = h.idx } if data.valid { for _, v := range data.host { h.hostData[v] = append(h.hostData[v], &data) } h.ipData[data.ip] = append(h.ipData[data.ip], &data) } } return nil } func (h *Host) parseLine(data string) (HostNode, error) { var res = HostNode{ original: data, } if len(data) == 0 { return res, fmt.Errorf("empty line") } if data[0] == '#' { res.comment = data res.onlyComment = true res.valid = true return res, nil } var dataArr []string cache := "" for k, v := range data { if v == '#' { if len(cache) > 0 { dataArr = append(dataArr, cache) cache = "" } dataArr = append(dataArr, data[k:]) break } if v == ' ' || v == '\t' { if len(cache) > 0 { dataArr = append(dataArr, cache) cache = "" } continue } cache += string(v) } if len(cache) > 0 { dataArr = append(dataArr, cache) } if len(dataArr) < 2 { return res, fmt.Errorf("invalid line") } if strings.HasPrefix(dataArr[1], "#") { return res, fmt.Errorf("invalid line") } if net.ParseIP(dataArr[0]) == nil { return res, fmt.Errorf("invalid ip address") } for k, v := range dataArr { switch k { case 0: res.ip = v default: if strings.HasPrefix(v, "#") { res.comment = v } else { res.host = append(res.host, v) } } } res.valid = true return res, nil } func (h *Host) List() []*HostNode { h.RLock() defer h.RUnlock() var res []*HostNode nextUid := h.firstUid for { if nextUid == 0 { break } res = append(res, h.fulldata[nextUid]) nextUid = h.fulldata[nextUid].nextuid } return res } func (h *Host) ListHosts() map[string][]string { h.RLock() defer h.RUnlock() if h.hostData == nil { return nil } var res = make(map[string][]string) for k, v := range h.hostData { for _, vv := range v { res[k] = append(res[k], vv.ip) } } return res } func (h *Host) ListIPs() map[string][]string { h.RLock() defer h.RUnlock() if h.ipData == nil { return nil } var res = make(map[string][]string) for k, v := range h.ipData { for _, vv := range v { res[k] = append(res[k], vv.host...) } } return res } func (h *Host) ListByHost(host string) []*HostNode { h.RLock() defer h.RUnlock() if h.hostData == nil { return nil } return h.hostData[host] } func (h *Host) ListIPsByHost(host string) []string { h.RLock() defer h.RUnlock() if h.hostData == nil { return nil } var res []string for _, v := range h.hostData[host] { res = append(res, v.ip) } return res } func (h *Host) ListFirstIPByHost(host string) string { h.RLock() defer h.RUnlock() if h.hostData == nil { return "" } for _, v := range h.hostData[host] { return v.ip } return "" } func (h *Host) ListByIP(ip string) []*HostNode { h.RLock() defer h.RUnlock() if h.ipData == nil { return nil } return h.ipData[ip] } func (h *Host) ListHostsByIP(ip string) []string { h.RLock() defer h.RUnlock() return h.listHostsByIP(ip) } func (h *Host) listHostsByIP(ip string) []string { if h.ipData == nil { return nil } var res []string for _, v := range h.ipData[ip] { res = append(res, v.host...) } return res } func (h *Host) ListFirstHostByIP(ip string) string { h.RLock() defer h.RUnlock() if h.ipData == nil { return "" } for _, v := range h.ipData[ip] { if len(v.host) > 0 { return v.host[0] } } return "" } func (h *Host) AddHosts(ip string, hosts ...string) error { return h.addHosts("", ip, hosts...) } func (h *Host) AddHostsComment(comment string, ip string, hosts ...string) error { return h.addHosts(comment, ip, hosts...) } func (h *Host) RemoveIPHosts(ip string, hosts ...string) error { h.Lock() defer h.Unlock() if h.ipData == nil { return fmt.Errorf("hosts data not initialized") } ipInfo := h.ipData[ip] if len(ipInfo) == 0 { return nil } cntfor: for _, v := range ipInfo { for _, host := range hosts { if len(v.host) == 1 && v.host[0] == host { delete(h.ipData, ip) if v.lastuid != 0 { h.fulldata[v.lastuid].nextuid = v.nextuid } else { h.firstUid = v.nextuid } if v.nextuid != 0 { h.fulldata[v.nextuid].lastuid = v.lastuid } else { h.lastUid = v.lastuid } var newHostData []*HostNode for _, vv := range h.hostData[v.host[0]] { if vv.uid != v.uid { newHostData = append(newHostData, vv) } } h.hostData[host] = newHostData delete(h.fulldata, v.uid) v = nil continue cntfor } if len(v.host) > 1 { var newHosts []string for _, vv := range v.host { if vv != host { newHosts = append(newHosts, vv) } } v.host = newHosts var newHostData []*HostNode for _, vv := range h.hostData[host] { if vv.uid != v.uid { newHostData = append(newHostData, vv) } } h.hostData[host] = newHostData } } } return nil } func (h *Host) RemoveIPs(ips ...string) error { h.Lock() defer h.Unlock() if h.ipData == nil { return fmt.Errorf("hosts data not initialized") } for _, ip := range ips { ipInfo := h.ipData[ip] if len(ipInfo) == 0 { continue } for _, v := range ipInfo { delete(h.ipData, ip) delete(h.fulldata, v.uid) if v.lastuid != 0 { h.fulldata[v.lastuid].nextuid = v.nextuid } else { h.firstUid = v.nextuid } if v.nextuid != 0 { h.fulldata[v.nextuid].lastuid = v.lastuid } else { h.lastUid = v.lastuid } for _, host := range v.host { var newHostData []*HostNode for _, vv := range h.hostData[host] { if vv.uid != v.uid { newHostData = append(newHostData, vv) } } h.hostData[host] = newHostData } } } return nil } func (h *Host) RemoveHosts(hosts ...string) error { h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } for _, host := range hosts { hostInfo := h.hostData[host] if len(hostInfo) == 0 { continue } delete(h.hostData, host) for _, v := range hostInfo { var newHosts []string for _, vv := range v.host { if vv != host { newHosts = append(newHosts, vv) } } v.host = newHosts if len(v.host) == 0 { delete(h.ipData, v.ip) if v.lastuid != 0 { h.fulldata[v.lastuid].nextuid = v.nextuid } else { h.firstUid = v.nextuid } if v.nextuid != 0 { h.fulldata[v.nextuid].lastuid = v.lastuid } else { h.lastUid = v.lastuid } delete(h.fulldata, v.uid) } } } return nil } func (h *Host) SetIPHosts(ip string, hosts ...string) error { info := h.ListByIP(ip) if len(info) == 0 { return h.AddHosts(ip, hosts...) } else if len(info) == 1 { info[0].host = hosts info[0].comment = "" return nil } err := h.RemoveIPs(ip) if err != nil { return err } return h.AddHosts(ip, hosts...) } func (h *Host) SetHostIPs(host string, ips ...string) error { info := h.ListByHost(host) if len(info) == 0 { return h.AddHosts(host, ips...) } else if len(info) == 1 { info[0].ip = ips[0] info[0].comment = "" return nil } err := h.RemoveHosts(host) if err != nil { return err } return h.AddHosts(ips[0], host) } func (h *Host) addHosts(comment string, ip string, hosts ...string) error { h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } ipInfo := h.listHostsByIP(ip) var needAddHosts []string for _, v := range hosts { if !inArray(ipInfo, v) { needAddHosts = append(needAddHosts, v) } } if len(needAddHosts) == 0 { return nil } if comment != "" && !strings.HasPrefix(strings.TrimSpace(comment), "#") { comment = "#" + comment } hostNode := HostNode{ uid: h.idx + 1, nextuid: 0, lastuid: h.lastUid, ip: ip, host: needAddHosts, valid: true, comment: comment, } h.idx++ h.fulldata[h.lastUid].nextuid = h.idx h.lastUid = h.idx h.fulldata[h.idx] = &hostNode h.ipData[ip] = append(h.ipData[ip], &hostNode) for _, v := range needAddHosts { h.hostData[v] = append(h.hostData[v], &hostNode) } return nil } func (h *Host) InsertNodeByData(node *HostNode, before bool, comment string, ip string, hosts ...string) error { h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } if _, ok := h.fulldata[node.uid]; !ok { return fmt.Errorf("node not exists") } if comment != "" && !strings.HasPrefix(strings.TrimSpace(comment), "#") { comment = "#" + comment } hostNode := HostNode{ uid: h.idx + 1, ip: ip, host: hosts, valid: true, comment: comment, } if ip == "" && len(hosts) == 0 && comment != "" { hostNode.onlyComment = true } if before { hostNode.nextuid = node.uid hostNode.lastuid = node.lastuid if node.lastuid != 0 { h.fulldata[node.lastuid].nextuid = h.idx } else { h.firstUid = h.idx } node.lastuid = h.idx } else { hostNode.lastuid = node.uid hostNode.nextuid = node.nextuid if node.nextuid != 0 { h.fulldata[node.nextuid].lastuid = h.idx } else { h.lastUid = h.idx } node.nextuid = h.idx } return nil } func (h *Host) DeleteNode(node *HostNode) error { h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } if _, ok := h.fulldata[node.uid]; !ok { return fmt.Errorf("node not exists") } if node.lastuid != 0 { h.fulldata[node.lastuid].nextuid = node.nextuid } else { h.firstUid = node.nextuid } if node.nextuid != 0 { h.fulldata[node.nextuid].lastuid = node.lastuid } else { h.lastUid = node.lastuid } for _, v := range node.host { var newHostData []*HostNode for _, vv := range h.hostData[v] { if vv.uid != node.uid { newHostData = append(newHostData, vv) } } h.hostData[v] = newHostData } ipInfo := h.ipData[node.ip] var newIPData []*HostNode for _, v := range ipInfo { if v.uid != node.uid { newIPData = append(newIPData, v) } } h.ipData[node.ip] = newIPData delete(h.fulldata, node.uid) return nil } func (h *Host) DeleteNodeByUID(uid uint64) error { h.RLock() node, ok := h.fulldata[uid] if !ok { h.RUnlock() return fmt.Errorf("node not exists") } h.RUnlock() return h.DeleteNode(node) } func (h *Host) AddNode(node *HostNode) error { h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } if node.comment != "" && !strings.HasPrefix(strings.TrimSpace(node.comment), "#") { node.comment = "#" + node.comment } if node.ip == "" && len(node.host) == 0 && node.comment != "" { node.onlyComment = true } h.idx++ node.uid = h.idx node.lastuid = h.lastUid node.nextuid = 0 h.fulldata[h.lastUid].nextuid = node.uid h.lastUid = node.uid h.fulldata[node.uid] = node node.valid = node.CheckValid() if node.valid { h.ipData[node.ip] = append(h.ipData[node.ip], node) for _, v := range node.host { h.hostData[v] = append(h.hostData[v], node) } } return nil } func (h *Host) InsertNode(node *HostNode) error { if node.lastuid == 0 && node.nextuid == 0 { return h.AddNode(node) } h.Lock() defer h.Unlock() if h.hostData == nil { return fmt.Errorf("hosts data not initialized") } if node.comment != "" && !strings.HasPrefix(strings.TrimSpace(node.comment), "#") { node.comment = "#" + node.comment } if node.ip == "" && len(node.host) == 0 && node.comment != "" { node.onlyComment = true } node.uid = h.idx + 1 if h.fulldata[node.lastuid].nextuid != node.nextuid { return fmt.Errorf("node lastuid nextuid not match") } h.idx++ h.fulldata[node.lastuid].nextuid = node.uid h.fulldata[node.nextuid].lastuid = node.uid h.fulldata[node.uid] = node node.valid = node.CheckValid() if node.valid { h.ipData[node.ip] = append(h.ipData[node.ip], node) for _, v := range node.host { h.hostData[v] = append(h.hostData[v], node) } } return nil } func inArray(arr []string, v string) bool { for _, vv := range arr { if v == vv { return true } } return false } func (h *Host) SaveAs(path string) error { return h.save(path) } func (h *Host) Save() error { return h.save(h.hostPath) } func (h *Host) save(path string) error { h.Lock() defer h.Unlock() if h.firstUid == 0 || h.fulldata == nil { return fmt.Errorf("no data") } f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666) if err != nil { return fmt.Errorf("open hosts file %s error: %s", h.hostPath, err) } defer f.Close() node := h.fulldata[h.firstUid] if node == nil { return fmt.Errorf("no data") } for { if node.onlyComment { if _, err := f.WriteString(node.comment + "\n"); err != nil { return fmt.Errorf("write hosts file error: %s", err) } if node.nextuid == 0 { break } node = h.fulldata[node.nextuid] continue } if _, err := f.WriteString(node.ip + " "); err != nil { return fmt.Errorf("write hosts file error: %s", err) } if _, err := f.WriteString(strings.Join(node.host, " ")); err != nil { return fmt.Errorf("write hosts file error: %s", err) } if len(node.comment) > 0 { if _, err := f.WriteString(" " + node.comment); err != nil { return fmt.Errorf("write hosts file error: %s", err) } } if _, err := f.WriteString(lineBreaker); err != nil { return fmt.Errorf("write hosts file error: %s", err) } if node.nextuid == 0 { break } node = h.fulldata[node.nextuid] } return nil } func (h *Host) Build() ([]byte, error) { if h.firstUid == 0 || h.fulldata == nil { return nil, fmt.Errorf("no data") } h.Lock() defer h.Unlock() var f bytes.Buffer node := h.fulldata[h.firstUid] if node == nil { return nil, fmt.Errorf("no data") } for { if node.onlyComment { if _, err := f.WriteString(node.comment + "\n"); err != nil { return nil, fmt.Errorf("write hosts file error: %s", err) } if node.nextuid == 0 { break } node = h.fulldata[node.nextuid] continue } if _, err := f.WriteString(node.ip + " "); err != nil { return nil, fmt.Errorf("write hosts file error: %s", err) } if _, err := f.WriteString(strings.Join(node.host, " ")); err != nil { return nil, fmt.Errorf("write hosts file error: %s", err) } if len(node.comment) > 0 { if _, err := f.WriteString(" " + node.comment); err != nil { return nil, fmt.Errorf("write hosts file error: %s", err) } } if _, err := f.WriteString(lineBreaker); err != nil { return nil, fmt.Errorf("write hosts file error: %s", err) } if node.nextuid == 0 { break } node = h.fulldata[node.nextuid] } return f.Bytes(), nil } func (h *Host) GetFirstNode() (*HostNode, error) { h.RLock() defer h.RUnlock() if h.firstUid == 0 { return nil, fmt.Errorf("no data") } if h.fulldata == nil { return nil, fmt.Errorf("no data") } return h.fulldata[h.firstUid], nil } func (h *Host) GetNode(nodeID uint64) (*HostNode, error) { h.RLock() defer h.RUnlock() if h.fulldata == nil { return nil, fmt.Errorf("no data") } if node, ok := h.fulldata[nodeID]; ok { return node, nil } return nil, fmt.Errorf("node not exists") } func (h *Host) GetNextNode(nodeID uint64) (*HostNode, error) { h.RLock() defer h.RUnlock() if h.fulldata == nil { return nil, fmt.Errorf("no data") } if node, ok := h.fulldata[nodeID]; ok { if node.nextuid == 0 { return nil, fmt.Errorf("no next node") } return h.fulldata[node.nextuid], nil } return nil, fmt.Errorf("node not exists") } func (h *Host) GetLastNode(nodeID uint64) (*HostNode, error) { h.RLock() defer h.RUnlock() if h.fulldata == nil { return nil, fmt.Errorf("no data") } if node, ok := h.fulldata[nodeID]; ok { if node.lastuid == 0 { return nil, fmt.Errorf("no last node") } return h.fulldata[node.lastuid], nil } return nil, fmt.Errorf("node not exists") } func (h *Host) GetLatestNode() (*HostNode, error) { h.RLock() defer h.RUnlock() if h.lastUid == 0 { return nil, fmt.Errorf("no data") } if h.fulldata == nil { return nil, fmt.Errorf("no data") } return h.fulldata[h.lastUid], nil } func (h *Host) UpdateNode(node *HostNode) error { h.Lock() defer h.Unlock() if h.fulldata == nil { return fmt.Errorf("no data") } if _, ok := h.fulldata[node.uid]; !ok { return fmt.Errorf("node not exists") } if node.comment != "" && !strings.HasPrefix(strings.TrimSpace(node.comment), "#") { node.comment = "#" + node.comment } if node.ip == "" && len(node.host) == 0 && node.comment != "" { node.onlyComment = true } h.fulldata[node.uid] = node node.valid = node.CheckValid() for k, v := range h.ipData { var newIPData []*HostNode for _, vv := range v { if vv.uid != node.uid { newIPData = append(newIPData, vv) } } h.ipData[k] = newIPData } for k, v := range h.hostData { var newHostData []*HostNode for _, vv := range v { if vv.uid != node.uid { newHostData = append(newHostData, vv) } } h.hostData[k] = newHostData } if node.valid { h.ipData[node.ip] = append(h.ipData[node.ip], node) for _, v := range node.host { h.hostData[v] = append(h.hostData[v], node) } } return nil } func (n *HostNode) String() string { if n.onlyComment { return n.comment } return fmt.Sprintf("%s %s %s", n.ip, strings.Join(n.host, " "), n.comment) } func (n *HostNode) Bytes() []byte { return []byte(n.String()) } func (n *HostNode) Hosts() []string { return n.host } func (n *HostNode) IP() string { return n.ip } func (n *HostNode) Comment() string { return n.comment } func (n *HostNode) UID() uint64 { return n.uid } func (n *HostNode) NextUID() uint64 { return n.nextuid } func (n *HostNode) LastUID() uint64 { return n.lastuid } func (n *HostNode) Valid() bool { return n.valid } func (n *HostNode) OnlyComment() bool { return n.onlyComment } func (n *HostNode) Original() string { return n.original } func (n *HostNode) SetHosts(hosts ...string) { n.host = hosts } func (n *HostNode) SetIP(ip string) { n.ip = ip } func (n *HostNode) SetComment(comment string) { if comment == "" { return } if !strings.HasPrefix(strings.TrimSpace(comment), "#") { n.comment = "#" + comment return } n.comment = comment } func (n *HostNode) SetValid(valid bool) { n.valid = valid } func (n *HostNode) SetOnlyComment(onlyComment bool) { n.onlyComment = onlyComment } func (n *HostNode) SetOriginal(original string) { n.original = original } func (n *HostNode) AddHosts(hosts ...string) { n.host = append(n.host, hosts...) } func (n *HostNode) SetLastUID(uid uint64) error { if n.uid != 0 { return fmt.Errorf("uid already set,you cannot change the existing node") } n.lastuid = uid return nil } func (n *HostNode) SetNextUID(uid uint64) error { if n.uid != 0 { return fmt.Errorf("uid already set,you cannot change the existing node") } n.nextuid = uid return nil } func (n *HostNode) CheckValid() bool { if n.ip == "" && len(n.host) == 0 && n.comment == "" { return true } if n.ip == "" && len(n.host) == 0 && n.comment != "" { return true } if n.ip == "" && len(n.host) > 0 { return false } if net.ParseIP(n.ip) == nil { return false } return true }