/
core.go
147 lines (135 loc) · 3.8 KB
/
core.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
package core
import (
"bufio"
"io/ioutil"
"os"
"strings"
"github.com/tinycedar/lily/common"
"github.com/tinycedar/lily/conf"
)
var systemHosts = os.Getenv("SystemRoot") + "/System32/drivers/etc/hosts"
// var batcher *Batcher
func FireHostsSwitch() bool {
common.Info("============================== Fire hosts switch ==============================")
// if batcher != nil {
// batcher.Close()
// }
return doProcess()
// batcher = initSystemHostsWatcher()
// go startSystemHostsWatcher()
}
func doProcess() bool {
success := true
hostsIPMap := getHostsIpMap()
overwriteSystemHosts()
processNameMap := GetProcessNameMap()
table := getTCPTable()
for i := uint32(0); i < uint32(table.dwNumEntries); i++ {
row := table.table[i]
ip := row.displayIP(row.dwRemoteAddr)
port := row.displayPort(row.dwRemotePort)
if row.dwOwningPid <= 0 {
continue
}
if port != 80 && port != 443 {
continue
}
process := strings.ToLower(processNameMap[uint32(row.dwOwningPid)])
if hostsIPMap[ip] || common.BrowserMap[process] {
if err := CloseTCPEntry(row); err != nil {
common.Error("Fail to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
success = false
} else {
common.Info("Succeed to close TCP connections: Process = %v, Pid = %v, Addr = %v:%v", process, row.dwOwningPid, ip, port)
}
}
}
return success
}
// find all the ip of system and current hosts
func getHostsIpMap() map[string]bool {
ipMap := make(map[string]bool)
for _, v := range readHostConfigMap(systemHosts) {
ipMap[v] = true
}
model := conf.Config.HostConfigModel
index := conf.Config.CurrentHostIndex
if length := len(model.Roots); index >= 0 && length > 0 && index < length {
path := "conf/hosts/" + model.RootAt(index).Text() + ".hosts"
for _, v := range readHostConfigMap(path) {
ipMap[v] = true
}
}
return ipMap
}
func overwriteSystemHosts() {
bytes := ReadCurrentHostConfig()
if bytes == nil {
return
}
if err := ioutil.WriteFile(systemHosts, bytes, os.ModeExclusive); err != nil {
common.Error("Error writing to system hosts file: %v", err)
}
}
func ReadCurrentHostConfig() []byte {
model := conf.Config.HostConfigModel
index := conf.Config.CurrentHostIndex
if length := len(model.Roots); index < 0 || length <= 0 || index >= length {
return nil
}
bytes, err := ioutil.ReadFile("conf/hosts/" + model.RootAt(index).Text() + ".hosts")
if err != nil {
common.Info("Error reading host config: %v", err)
return nil
}
return bytes
}
func readHostConfigMap(path string) map[string]string {
hostConfigMap := make(map[string]string)
file, err := os.Open(path)
if err != nil {
common.Error("Fail to open system_hosts: %s", err)
return hostConfigMap
}
defer file.Close()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line != "" && !strings.HasPrefix(line, "#") {
config := strings.Fields(line)
if len(config) == 2 {
hostConfigMap[config[1]] = config[0]
}
}
}
if err := scanner.Err(); err != nil {
common.Error("Fail to read system_hosts: %s", err)
}
return hostConfigMap
}
// func initSystemHostsWatcher() *Batcher {
// batcher, err := New(time.Millisecond * 300)
// if err != nil {
// common.Error("Fail to initialize batcher")
// }
// if err = batcher.Add(systemHosts); err != nil {
// common.Error("Fail to add system hosts: %s", systemHosts)
// }
// return batcher
// }
// func startSystemHostsWatcher() {
// if batcher == nil {
// common.Error("Fail to start system hosts watcher, watcher is nil")
// return
// }
// for events := range batcher.Events {
// for _, event := range events {
// if event.Op&fsnotify.Write == fsnotify.Write {
// common.Info("modified file: %v", event)
// doProcess()
// break
// }
// }
// }
// // never return
// }