/
hostshelper.go
91 lines (85 loc) · 2.09 KB
/
hostshelper.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
package main
import (
"bufio"
"errors"
"fmt"
"io/ioutil"
"os"
"regexp"
"runtime"
"strings"
)
func checkForAdmin() bool {
if runtime.GOOS == "windows" {
_, err := os.Open("\\\\.\\PHYSICALDRIVE0")
if err != nil {
return false
}
return true
} else {
fp, err := os.OpenFile("/etc/hosts", os.O_RDWR, 0644)
if err != nil {
return false
}
defer fp.Close()
return true
}
}
func getHostsLocation() string {
if runtime.GOOS == "windows" {
return `%SystemDrive%\Windows\System32\drivers\etc\hosts`
}
return `/etc/hosts`
}
func addHosts(records []HostsRecord) error {
if !checkForAdmin() {
return errors.New("Must run with priviledged user to modify hosts")
}
hostslocation := getHostsLocation()
b, err := ioutil.ReadFile(hostslocation)
if err != nil {
return err
}
originalhosts := string(b)
LineBreak := detectLineBreakFromString(originalhosts)
for _, record := range records {
substituionexp := regexp.MustCompile(fmt.Sprintf(`[a-fA-F\d\.\:]+\s+%s(.*)`, strings.Replace(record.hostname, ".", "\\.", -1)))
if occurences := substituionexp.FindAllString(originalhosts, -1); len(occurences) != 0 {
replaced := false
for _, occurence := range occurences {
if !replaced {
originalhosts = strings.Replace(originalhosts, occurence, fmt.Sprintf("%s %s", record.ip, record.hostname), 1)
replaced = true
}
originalhosts = strings.Replace(originalhosts, occurence, "", -1)
}
} else {
if !strings.HasSuffix(originalhosts, "\n") && !strings.HasSuffix(originalhosts, "\r") {
originalhosts += LineBreak
}
originalhosts += fmt.Sprintf("%s %s", record.ip, record.hostname) + LineBreak
}
}
fmt.Println("\n\n\n==========Modified Hosts file ===========")
fmt.Print(originalhosts)
fmt.Print("Write?[y/n]")
reader := bufio.NewReader(os.Stdin)
char, _, err := reader.ReadRune()
if err != nil {
return err
}
switch char {
case 'y':
fallthrough
case 'Y':
err := ioutil.WriteFile(hostslocation, []byte(originalhosts), 0644)
if err != nil {
return err
}
fmt.Println("Success")
break
default:
fmt.Println("Giving up...")
}
return nil
}