misc: Add notification on switch

Author: tinycedar

common/constants.go

@@ -14,7 +14,21 @@ const (
 )
 
 var IconMap = make(map[int]*walk.Bitmap)
-var BrowserMap = make(map[string]bool)
+var BrowserMap = map[string]bool{
+	"chrome.exe":          true,
+	"firefox.exe":         true,
+	"opera.exe":           true,
+	"iexplore.exe":        true,
+	"microsoftedge.exe":   true,
+	"microsoftedgecp.exe": true,
+	"sogouexplorer.exe":   true,
+	"qqbrowser.exe":       true,
+	"360se.exe":           true,
+	"360chrome.exe":       true,
+	"liebao.exe":          true,
+	"maxthon.exe":         true,
+	"ucbrowser.exe":       true,
+}
 
 func Init() {
 	IconMap[ICON_ADD], _ = walk.NewBitmapFromFile("res/add.png")
@@ -25,18 +39,4 @@ func Init() {
 	IconMap[ICON_PLUS], _ = walk.NewBitmapFromFile("res/plus.png")
 	IconMap[ICON_REFRESH], _ = walk.NewBitmapFromFile("res/refresh.png")
 	IconMap[ICON_SETTING], _ = walk.NewBitmapFromFile("res/setting.png")
-
-	BrowserMap["chrome.exe"] = true
-	BrowserMap["firefox.exe"] = true
-	BrowserMap["opera.exe"] = true
-	BrowserMap["iexplore.exe"] = true
-	BrowserMap["microsoftedge.exe"] = true
-	BrowserMap["microsoftedgecp.exe"] = true
-	BrowserMap["sogouexplorer.exe"] = true
-	BrowserMap["qqbrowser.exe"] = true
-	BrowserMap["360se.exe"] = true
-	BrowserMap["360chrome.exe"] = true
-	BrowserMap["liebao.exe"] = true
-	BrowserMap["maxthon.exe"] = true
-	BrowserMap["ucbrowser.exe"] = true
 }

core/core.go

@@ -16,58 +16,58 @@ const (
 
 // var batcher *Batcher
 
-func FireHostsSwitch() {
+func FireHostsSwitch() bool {
 	common.Info("============================== Fire hosts switch ==============================")
 	// if batcher != nil {
 	// 	batcher.Close()
 	// }
-	doProcess()
+	return doProcess()
 	// batcher = initSystemHostsWatcher()
 	// go startSystemHostsWatcher()
 }
 
 // 1. Find collection of same domain names between system hosts and currentHostIndex
 // 2. Disconnect the TCP connections(http:80 & https:443) of collection found above
-func doProcess() {
-	overlapHostConfigMap := getHostsIpMap()
+func doProcess() bool {
+	success := true
+	hostsIpMap := getHostsIpMap()
 	overwriteSystemHosts()
-	common.Info("overlapHostConfigMap: %v", overlapHostConfigMap)
-	if len(overlapHostConfigMap) == 0 {
-		return
-	}
 	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
 		}
-		ip := row.displayIP(row.dwRemoteAddr)
-		port := row.displayPort(row.dwRemotePort)
-		if _, ok := overlapHostConfigMap[ip]; !ok {
+		if port != 80 && port != 443 {
 			continue
 		}
-		if common.BrowserMap[OpenProcess(uint32(row.dwOwningPid))] || port == 80 || port == 443 {
+		supported := common.BrowserMap[strings.ToLower(OpenProcess(uint32(row.dwOwningPid)))]
+		if hostsIpMap[ip] || supported {
 			if err := CloseTCPEntry(row); err != nil {
 				common.Error("Fail to close TCP connections: Pid = %v, Addr = %v:%v\n", row.dwOwningPid, ip, port)
+				success = false
 			} else {
 				common.Info("Succeed to close TCP connections: Pid = %v, Addr = %v:%v", 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 k := range readHostConfigMap(systemHosts) {
-		ipMap[k] = true
+	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 k := range readHostConfigMap(path) {
-			ipMap[k] = true
+		for _, v := range readHostConfigMap(path) {
+			ipMap[v] = true
 		}
 	}
 	return ipMap
@@ -152,29 +152,18 @@ func readHostConfigMap(path string) map[string]string {
 	for scanner.Scan() {
 		line := strings.TrimSpace(scanner.Text())
 		if line != "" && !strings.HasPrefix(line, "#") {
-			config := trimDuplicateSpaces(line)
+			config := strings.Fields(line)
 			if len(config) == 2 {
 				hostConfigMap[config[1]] = config[0]
 			}
 		}
 	}
 	if err := scanner.Err(); err != nil {
-		common.Info("Fail to read system_hosts: %s", err)
+		common.Error("Fail to read system_hosts: %s", err)
 	}
 	return hostConfigMap
 }
 
-func trimDuplicateSpaces(line string) []string {
-	temp := []string{}
-	line = strings.TrimSpace(line)
-	for _, v := range strings.SplitN(line, " ", 2) {
-		if trimed := strings.TrimSpace(v); trimed != "" {
-			temp = append(temp, trimed)
-		}
-	}
-	return temp
-}
-
 // func initSystemHostsWatcher() *Batcher {
 // 	batcher, err := New(time.Millisecond * 300)
 // 	if err != nil {

core/sys_process.go

@@ -50,7 +50,7 @@ func OpenProcess(pid uint32) string {
 	var dwDesiredAccess uint32 = 0x0400 | 0x0010
 	openPid, _, _ := procOpenProcess.Call(uintptr(unsafe.Pointer(&dwDesiredAccess)), 0, uintptr(pid))
 	if openPid <= 0 {
-		common.Info("Fail to open process: Pid = %v", pid)
+		common.Error("Fail to open process: Pid = %v", pid)
 		return ""
 	}
 	defer closeHandle(openPid)
@@ -71,7 +71,13 @@ func getProcessName(pid uintptr) string {
 	var cbNeeded uint32 = 1024
 	var processName = make([]byte, cbNeeded)
 	procQueryFullProcessImageName.Call(uintptr(pid), 0, uintptr(unsafe.Pointer(&processName[0])), uintptr(unsafe.Pointer(&cbNeeded)))
-	return string(processName[0:cbNeeded])
+	processPath := string(processName[0:cbNeeded])
+	for i := len(processPath) - 2; i >= 0; i-- {
+		if processPath[i] == '\\' {
+			return processPath[i+1:]
+		}
+	}
+	return processPath
 }
 
 func getProcessName2(pid uintptr) string {

gui/mainwindow.go

@@ -18,6 +18,7 @@ type widgetContext struct {
 	hostConfigText *walk.TextEdit
 	addButton      *walk.Action
 	deleteButton   *walk.Action
+	notifyIcon     *walk.NotifyIcon
 }
 
 // InitMainWindow initialize main window
@@ -43,12 +44,13 @@ func InitMainWindow() {
 	}
 	setWindowIcon(context.mw)
 	setXY(context.mw)
-	newNotify(context.mw)
+	newNotify()
 	setTreeViewBackground(context.treeView)
 	showCurrentItem(context.hostConfigText)
 	context.mw.Closing().Attach(func(canceled *bool, reason walk.CloseReason) {
 		*canceled = true
 		context.mw.Hide()
+		context.notifyIcon.ShowCustom("当前程序最小化到托盘", "Best hosts manager ever")
 	})
 	(*context.deleteButton).SetEnabled(false)
 	context.mw.Run()

gui/notify.go

@@ -5,27 +5,28 @@ import (
 	"github.com/tinycedar/lily/common"
 )
 
-func newNotify(mw *walk.MainWindow) {
-	ni, err := walk.NewNotifyIcon()
+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 := ni.SetIcon(icon); err != nil {
+	if err := context.notifyIcon.SetIcon(icon); err != nil {
 		common.Error("Error setting notify icon: %v", err)
 	}
-	if err := ni.SetToolTip("Click for info or use the context menu to exit."); err != nil {
+	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)
 	}
-	ni.MouseUp().Attach(func(x, y int, button walk.MouseButton) {
+	context.notifyIcon.MouseUp().Attach(func(x, y int, button walk.MouseButton) {
 		if button == walk.LeftButton {
-			if !mw.Visible() {
-				mw.Show()
+			if !context.mw.Visible() {
+				context.mw.Show()
 			} else {
 				//TODO
 			}
 		}
-		// if err := ni.ShowCustom(
+		// 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)
@@ -36,16 +37,16 @@ func newNotify(mw *walk.MainWindow) {
 		common.Error("Error setting exitAction text: %v", err)
 	}
 	exitAction.Triggered().Attach(func() {
-		defer ni.Dispose()
+		defer context.notifyIcon.Dispose()
 		walk.App().Exit(0)
 	})
-	if err := ni.ContextMenu().Actions().Add(exitAction); err != nil {
+	if err := context.notifyIcon.ContextMenu().Actions().Add(exitAction); err != nil {
 		common.Error("Error Adding exitAction: %v", err)
 	}
-	if err := ni.SetVisible(true); err != nil {
+	if err := context.notifyIcon.SetVisible(true); err != nil {
 		common.Error("Error setting notify visible: %v", err)
 	}
-	// if err := ni.ShowInfo("Walk NotifyIcon Example", "Click the icon to show again."); err != nil {
+	// if err := context.notifyIcon.ShowInfo("Walk NotifyIcon Example", "Click the icon to show again."); err != nil {
 	// 	common.Error("Error showing info: %v", err)
 	// }
 }

gui/treeview.go

@@ -48,8 +48,11 @@ func newTreeView() TreeView {
 			} else {
 				ioutil.WriteFile("conf/config.json", configJSON, os.ModeExclusive)
 			}
-			// fire hosts switch
-			core.FireHostsSwitch()
+			if core.FireHostsSwitch() {
+				context.notifyIcon.ShowCustom("切换"+current.Text()+"成功！", "Best hosts manager ever")
+			} else {
+				context.notifyIcon.ShowError("切换"+current.Text()+"失败！", "Best hosts manager ever")
+			}
 		},
 	}
 }