151 lines
2.6 KiB
Go
151 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"io"
|
|
"io/fs"
|
|
"io/ioutil"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
type AliasMap map[string]string
|
|
|
|
const (
|
|
HostsFile string = "/etc/hosts"
|
|
)
|
|
|
|
func FixedSplit(s, sep string, parts int) []string {
|
|
n := make([]string, parts)
|
|
p := strings.SplitN(s, sep, parts)
|
|
copy(n, p)
|
|
return n
|
|
}
|
|
|
|
func IsLegalLine(line string) bool {
|
|
c := line[0]
|
|
return c != '[' && c != '#'
|
|
}
|
|
|
|
func BuildRegionString(regionName string, aliases AliasMap) string {
|
|
b := strings.Builder{}
|
|
b.WriteString("#region ")
|
|
b.WriteString(regionName)
|
|
b.WriteString("\n")
|
|
|
|
for ip, alias := range aliases {
|
|
b.WriteString(ip)
|
|
b.WriteString("\t")
|
|
b.WriteString(alias)
|
|
b.WriteString("\n")
|
|
}
|
|
|
|
b.WriteString("#endregion")
|
|
return b.String()
|
|
}
|
|
|
|
func BuildRegionRegexp(regionName string) *regexp.Regexp {
|
|
b := strings.Builder{}
|
|
b.WriteString("(?s)\n#region ")
|
|
b.WriteString(regexp.QuoteMeta(regionName))
|
|
b.WriteString(".*#endregion")
|
|
r := regexp.MustCompile(b.String())
|
|
return r
|
|
}
|
|
|
|
func ScanAliases(fileReader io.Reader) (AliasMap, error) {
|
|
aliasMap := AliasMap{}
|
|
scanner := bufio.NewScanner(fileReader)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if IsLegalLine(line) {
|
|
s := FixedSplit(line, "#", 2)
|
|
ip, alias := strings.TrimSpace(s[0]), strings.TrimSpace(s[1])
|
|
if _, ok := aliasMap[ip]; !ok && alias != "" {
|
|
aliasMap[ip] = alias
|
|
}
|
|
}
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return aliasMap, nil
|
|
}
|
|
|
|
func AddAliases(fileName string) {
|
|
file, err := os.Open(fileName)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
h, err := ScanAliases(file)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
file.Close()
|
|
|
|
r := BuildRegionRegexp(fileName)
|
|
s := BuildRegionString(fileName, h)
|
|
|
|
content, err := ioutil.ReadFile(HostsFile)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
c := r.ReplaceAllString(string(content), s)
|
|
if !r.MatchString(c) {
|
|
c += ("\n" + s)
|
|
}
|
|
|
|
err = os.WriteFile(HostsFile, []byte(c[:]), fs.FileMode(os.O_WRONLY|os.O_TRUNC))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func RemoveAliases(fileName string) {
|
|
regionReg := BuildRegionRegexp(fileName)
|
|
content, err := ioutil.ReadFile(HostsFile)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
c := regionReg.ReplaceAll(content, []byte(""))
|
|
err = os.WriteFile(HostsFile, c, fs.FileMode(os.O_WRONLY|os.O_TRUNC))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
u := fmt.Sprintf("Please use: %s <add|rm> <file:path>\n", os.Args[0])
|
|
|
|
if len(os.Args) < 3 {
|
|
fmt.Println(u)
|
|
os.Exit(1)
|
|
}
|
|
|
|
p, err := filepath.Abs(os.Args[2])
|
|
if err != nil {
|
|
fmt.Println(err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
switch os.Args[1] {
|
|
case "add":
|
|
AddAliases(p)
|
|
case "rm":
|
|
RemoveAliases(p)
|
|
default:
|
|
fmt.Println(u)
|
|
os.Exit(1)
|
|
}
|
|
}
|