-
-
Notifications
You must be signed in to change notification settings - Fork 45
/
host_file.go
219 lines (185 loc) · 4.39 KB
/
host_file.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
package host
import (
"bufio"
"fmt"
"io"
"net"
"os"
"regexp"
"strings"
)
var (
profileStart = regexp.MustCompile(`(?i)# profile\s+([a-z0-9-_]+)\s*`)
profileEnd = regexp.MustCompile(`(?i)# end\s*`)
disableRe = regexp.MustCompile(`^#\s*`)
spaceRemover = regexp.MustCompile(`\s+`)
tabReplacer = regexp.MustCompile(`\t+`)
)
type hostFile struct {
profiles profileMap
}
type profileMap map[string]hostLines
type hostLines []string
// ReadHostFile open a file an read content into a hostFile struct
func ReadHostFile(file string) (*hostFile, error) {
fromFile, err := os.Open(file)
if err != nil {
return nil, err
}
return Read(fromFile, false)
}
// ReadHostFileStrict open a file an read content into a hostFile struct. removes all comments.
func ReadHostFileStrict(file string) (*hostFile, error) {
_, err := os.Stat(file)
if err != nil {
return nil, err
}
fromFile, err := os.Open(file)
if err != nil {
return nil, err
}
return Read(fromFile, true)
}
// Read returns hosts file content grouped by profiles.
// If you pass strict=true it would remove all comments.
func Read(r io.Reader, strict bool) (*hostFile, error) {
h := &hostFile{
profiles: profileMap{},
}
ln := 0
s := bufio.NewScanner(r)
open := ""
for s.Scan() {
ln++
b := s.Bytes()
switch {
case profileStart.Match(b):
open = strings.TrimSpace(strings.Split(string(b), "# profile")[1])
h.profiles[open] = []string{}
case profileEnd.Match(b):
open = ""
case open != "":
if strict && !IsHostLine(string(b)) {
// skip
} else {
h.profiles[open] = append(h.profiles[open], string(b))
}
default:
h.profiles["default"] = append(h.profiles["default"], string(b))
}
if err := s.Err(); err != nil {
return nil, err
}
}
return h, nil
}
// ReadFromArgs read arguments into a hostFile struct
func ReadFromArgs(domains []string, ip string) *hostFile {
dom := make([]string, len(domains))
for k, d := range domains {
dom[k] = fmt.Sprintf("%s %s", ip, d)
}
newData := &hostFile{
profiles: profileMap{
"default": dom,
},
}
return newData
}
// IsHostLine checks if a line is a host line or a comment line.
func IsHostLine(line string) bool {
p := strings.Split(cleanLine(line), " ")
i := 0
if p[0] == "#" && len(p) > 1 {
i = 1
}
ip := net.ParseIP(p[i])
return ip != nil
}
func cleanLine(line string) string {
return tabReplacer.ReplaceAllString(spaceRemover.ReplaceAllString(line, " "), " ")
}
// IsDisabled check if a line starts with a # comment marker.
func IsDisabled(line string) bool {
return disableRe.MatchString(line)
}
// EnableLine removes the # comment marker of the line.
func EnableLine(line string) string {
return disableRe.ReplaceAllString(line, "")
}
func banner() string {
return `
##################################################################
# Content under this line is handled by hostctl. DO NOT EDIT.
##################################################################`
}
func containsBanner(str string) bool {
b := banner()
m, _ := regexp.MatchString(b, str)
return m
}
// WriteToFile write hosts content to file
func WriteToFile(f *os.File, h *hostFile) error {
err := f.Truncate(0)
if err != nil {
return err
}
for _, l := range h.profiles["default"] {
_, err = f.WriteString(l + "\n")
if err != nil {
return err
}
}
content := strings.Join(h.profiles["default"], "\n")
if !containsBanner(content) {
f.WriteString("\n" + banner() + "\n")
}
for n, p := range h.profiles {
if n == "default" {
continue
}
err = addProfile(f, n, p)
if err != nil {
return err
}
}
return nil
}
func addProfile(f *os.File, profile string, hl hostLines) error {
_, err := f.WriteString(fmt.Sprintf("# profile %s\n", profile))
if err != nil {
return err
}
for _, l := range hl {
_, err = f.WriteString(l + "\n")
if err != nil {
return err
}
}
_, err = f.WriteString("# end\n")
return err
}
func getHostData(dst, profile string) (*hostFile, error) {
h := &hostFile{
profiles: profileMap{},
}
if dst == "" {
return h, MissingDestError
}
h, err := ReadHostFile(dst)
if err != nil {
return h, err
}
_, ok := h.profiles[profile]
if profile != "" && !ok {
return h, fmt.Errorf("unknown profile '%s'", profile)
}
return h, nil
}
func writeHostData(dst string, h *hostFile) error {
dstFile, err := os.OpenFile(dst, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0644)
if err != nil {
return err
}
return WriteToFile(dstFile, h)
}