staros/hosts/hosts.go

package hosts

import (
	"bytes"
	"io/ioutil"
	"runtime"
	"sync"

	"b612.me/staros/sysconf"
)

var hostsPath string
var hostsCfg *sysconf.SysConf
var haveError error
var lock sync.Mutex
var reverse bool = false

func init() {
	if runtime.GOOS == "windows" {
		hostsPath = `C:\Windows\System32\drivers\etc\hosts`
	} else {
		hostsPath = `/etc/hosts`
	}
}

func SetHostsPath(path string) {
	hostsPath = path
}

func GetHostsPath() string {
	return hostsPath
}

func Parse() error {
	hostsCfg = new(sysconf.SysConf)
	hostsCfg.EqualFlag = ` `
	hostsCfg.HaveSegMent = false
	hostsCfg.CommentCR = true
	hostsCfg.CommentFlag = []string{"#"}
	data, haveError := ioutil.ReadFile(hostsPath)
	if haveError != nil {
		return haveError
	}
	data = bytes.ReplaceAll(data, []byte(`	`), []byte(" "))
	haveError = hostsCfg.Parse(data)
	return haveError
}

func GetIpByDomain(domainName string) []string {
	if haveError != nil {
		return []string{}
	}
	lock.Lock()
	defer lock.Unlock()
	if !reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	return hostsCfg.Data[0].GetAll(domainName)
}

func GetDomainByIp(IpAddr string) []string {
	if haveError != nil {
		return []string{}
	}
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	return hostsCfg.Data[0].GetAll(IpAddr)
}

func AddHosts(ip, host string) {
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Data[0].AddValue(host, ip, "")
	} else {
		hostsCfg.Data[0].AddValue(ip, host, "")
	}
}

func RemoveHost(ip, host string) {
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Data[0].DeleteValue(host, ip)
	} else {
		hostsCfg.Data[0].DeleteValue(ip, host)
	}
}

func RemoveHostbyIp(ip string) {
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	hostsCfg.Data[0].Delete(ip)
}

func RemoveHostbyHost(host string) {
	lock.Lock()
	defer lock.Unlock()
	if !reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	hostsCfg.Data[0].Delete(host)
}

// SetHost 设定唯一的Host对应值，一个host对应一个ip，其余的删除
func SetHost(ip, host string) {
	lock.Lock()
	defer lock.Unlock()
	if !reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	hostsCfg.Data[0].Set(host, ip, "")
}

func SetHostbyIp(ip, host string) {
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	hostsCfg.Data[0].Set(ip, host, "")
}

func Build() string {
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	return string(hostsCfg.Build())
}

func Write() error {
	lock.Lock()
	defer lock.Unlock()
	data := []byte(Build())
	return ioutil.WriteFile(hostsPath, data, 0644)
}

func GetHostList() []string {
	var res []string
	lock.Lock()
	defer lock.Unlock()
	if !reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	for _, v := range hostsCfg.Data[0].NodeData {
		if v != nil {
			res = append(res, v.Key)
		}
	}
	return res
}

func GetIpList() []string {
	var res []string
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	for _, v := range hostsCfg.Data[0].NodeData {
		if v != nil {
			res = append(res, v.Key)
		}
	}
	return res
}

func GetAllListbyHost() map[string][]string {
	res := make(map[string][]string)
	lock.Lock()
	defer lock.Unlock()
	if !reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	for _, v := range hostsCfg.Data[0].NodeData {
		if v != nil {
			res[v.Key] = v.Value
		}
	}
	return res
}

func GetAllListbyIp() map[string][]string {
	res := make(map[string][]string)
	lock.Lock()
	defer lock.Unlock()
	if reverse {
		hostsCfg.Reverse()
		reverse = !reverse
	}
	for _, v := range hostsCfg.Data[0].NodeData {
		if v != nil {
			res[v.Key] = v.Value
		}
	}
	return res
}