func GetProcessNameMap() map[uint32]string {
defer metrics("GetProcessNameMap")(time.Now())
snapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPPROCESS, 0)
if err != nil {
common.Error("Fail to syscall CreateToolhelp32Snapshot: %v", err)
return nil
}
defer syscall.CloseHandle(snapshot)
var procEntry syscall.ProcessEntry32
procEntry.Size = uint32(unsafe.Sizeof(procEntry))
if err = syscall.Process32First(snapshot, &procEntry); err != nil {
common.Error("Fail to syscall Process32First: %v", err)
return nil
}
processNameMap := make(map[uint32]string)
for {
processNameMap[procEntry.ProcessID] = parseProcessName(procEntry.ExeFile)
if err = syscall.Process32Next(snapshot, &procEntry); err != nil {
if err == syscall.ERROR_NO_MORE_FILES {
return processNameMap
}
common.Error("Fail to syscall Process32Next: %v", err)
return nil
}
}
}
// netstat -ano | findstr 202.89.233.104
func getTCPTable() *MIB_TCPTABLE2 {
getTCPTable2 := syscall.NewLazyDLL("Iphlpapi.dll").NewProc("GetTcpTable2")
var n uint32
if err, _, _ := getTCPTable2.Call(uintptr(unsafe.Pointer(&MIB_TCPTABLE2{})), uintptr(unsafe.Pointer(&n)), 1); syscall.Errno(err) != syscall.ERROR_INSUFFICIENT_BUFFER {
common.Error("Error calling GetTcpTable2: %v", syscall.Errno(err))
}
b := make([]byte, n)
if err, _, _ := getTCPTable2.Call(uintptr(unsafe.Pointer(&b[0])), uintptr(unsafe.Pointer(&n)), 1); err != 0 {
common.Error("Error calling GetTcpTable2: %v", syscall.Errno(err))
}
table := newTCPTable(NewClassReader(b))
return table
}
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
}
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: ", err)
}
}
func watchPidFile() {
batcher, err := New(time.Millisecond * 50)
if err != nil {
common.Error("Fail to initialize batcher")
return
}
if err := batcher.Add(pidFilePath); err != nil {
common.Error("Fail to add pid file: %s", pidFilePath)
return
}
for events := range batcher.Events {
for _, event := range events {
if event.Op&fsnotify.Write == fsnotify.Write {
common.Info("modified file: %v", event)
trigger <- true
break
}
}
}
}
func newTextEdit() TextEdit {
return TextEdit{
AssignTo: &(context.hostConfigText),
StretchFactor: 3,
OnKeyUp: func(key walk.Key) {
current := context.treeView.CurrentItem()
if current != nil {
file := "conf/hosts/" + current.Text() + ".hosts"
if err := ioutil.WriteFile(file, []byte(context.hostConfigText.Text()), os.ModeExclusive); err != nil {
common.Error("Error writing to system hosts file: ", err)
}
}
},
}
}
func newNotify() {
var err error
context.notifyIcon, err = walk.NewNotifyIcon()
if err != nil {
common.Error("Error invoking NewNotifyIcon: %v", err)
}
icon, _ := walk.NewIconFromFile("res/lily.ico")
if err := context.notifyIcon.SetIcon(icon); err != nil {
common.Error("Error setting notify icon: %v", err)
}
if err := context.notifyIcon.SetToolTip("Click for info or use the context menu to exit."); err != nil {
common.Error("Fail to set tooltip: %v", err)
}
f := func() {
if !context.mw.Visible() {
context.mw.Show()
} else {
context.mw.SwitchToThisWindow()
}
}
go core.Triggered(f)
context.notifyIcon.MouseUp().Attach(func(x, y int, button walk.MouseButton) {
if button == walk.LeftButton {
f()
}
// if err := context.notifyIcon.ShowCustom(
// "Walk NotifyIcon Example",
// "There are multiple ShowX methods sporting different icons."); err != nil {
// common.Error("Fail to show custom notify: %v", err)
// }
})
exitAction := walk.NewAction()
if err := exitAction.SetText("退出"); err != nil {
common.Error("Error setting exitAction text: %v", err)
}
exitAction.Triggered().Attach(func() {
context.notifyIcon.Dispose()
// os.Exit(-1)
walk.App().Exit(0)
})
if err := context.notifyIcon.ContextMenu().Actions().Add(exitAction); err != nil {
common.Error("Error Adding exitAction: %v", err)
}
if err := context.notifyIcon.SetVisible(true); err != nil {
common.Error("Error setting notify visible: %v", err)
}
// if err := context.notifyIcon.ShowInfo("Walk NotifyIcon Example", "Click the icon to show again."); err != nil {
// common.Error("Error showing info: %v", err)
// }
}
func readHostConfigMap(path string) map[string]string {
hostConfigMap := make(map[string]string)
file, err := os.Open(path)
if err != nil {
common.Info("Fail to open system_hosts: %s", err)
}
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 writePidFile(pid int) {
if err := ioutil.WriteFile(pidFilePath, []byte(fmt.Sprint(pid)), os.ModeExclusive); err != nil {
common.Error("Error writing to system hosts file: ", err)
}
}
func newToolBar() ToolBar {
tb := ToolBar{
ButtonStyle: ToolBarButtonImageBeforeText,
Items: []MenuItem{
Action{
AssignTo: &(context.addButton),
Text: "新增",
Image: "res/add.png",
// Enabled: Bind("isSpecialMode && enabledCB.Checked"),
OnTriggered: func() {
var dlg *walk.Dialog
var hostsNameEdit *walk.LineEdit
Dialog{
AssignTo: &dlg,
Title: "新增",
MinSize: Size{300, 150},
Layout: VBox{},
Children: []Widget{
Composite{
Layout: Grid{Columns: 2},
Children: []Widget{
Label{
ColumnSpan: 2,
Text: "Hosts名字:",
},
LineEdit{
AssignTo: &hostsNameEdit,
ColumnSpan: 2,
// Text: Bind("PatienceField"),
},
},
},
HSpacer{},
Composite{
Layout: HBox{},
Children: []Widget{
PushButton{
// AssignTo: &acceptPB,
Text: "确定",
OnClicked: func() {
// if err := db.Submit(); err != nil {
// log.Print(err)
// return
// }
item := &model.HostConfigItem{Name: hostsNameEdit.Text(), Icon: common.IconMap[common.ICON_NEW]}
conf.Config.HostConfigModel.Insert(item)
context.treeView.SetCurrentItem(item)
dlg.Accept()
},
},
PushButton{
// AssignTo: &cancelPB,
Text: "取消",
OnClicked: func() { dlg.Cancel() },
},
},
},
},
}.Run(context.mw)
},
},
//FIXME 去除刷新按钮是因为点击以后， 双击hosts不再生效
// Action{
// Text: "刷新",
// Image: "res/refresh.png",
// // Enabled: Bind("isSpecialMode && enabledCB.Checked"),
// OnTriggered: func() {
// conf.Load()
// },
// },
// Action{
// Text: "修改",
// Image: "res/pencil.png",
// // Enabled: Bind("isSpecialMode && enabledCB.Checked"),
// // OnTriggered: mw.specialAction_Triggered,
// },
Action{
AssignTo: &(context.deleteButton),
Text: "删除",
Image: "res/delete.png",
// Enabled: Bind("isSpecialMode && enabledCB.Checked"),
OnTriggered: func() {
if context.treeView.CurrentItem() == nil {
walk.MsgBox(context.mw, "删除hosts", "请选择左边列表后再删除", walk.MsgBoxIconInformation)
context.deleteButton.SetEnabled(false)
return
}
current := context.treeView.CurrentItem().(*model.HostConfigItem)
message := fmt.Sprintf("确定要删除hosts '%s'?", current.Text())
ret := walk.MsgBox(context.mw, "删除hosts", message, walk.MsgBoxYesNo)
if ret == win.IDYES {
if !conf.Config.HostConfigModel.Remove(current) {
common.Error("Fail to remove current item: %v", current.Text())
// TODO notify user
return
}
if context.treeView.Model().RootCount() > 0 {
context.treeView.SetCurrentItem(context.treeView.Model().RootAt(0))
}
file := "conf/hosts/" + current.Text() + ".hosts"
if err := os.Remove(file); err != nil {
common.Error("Fail to delete file: %s", file)
// TODO notify user
return
}
common.Info("Succeed to delete file: %s", file)
}
},
},
},
}
return tb
}