Add a feature of writing to the hosts file; Add a feature of limiting…

… the number of scans.

README.md

@@ -27,11 +27,16 @@ Options:
 ## Configuration
 ```toml
 [General]
-Debug = true
-# workers
-Workers = 300
 # GoogleTranslate or Cloudflare
 UsedFor = "GoogleTranslate"
+# A boolean that turns on/off debug mode. true or false
+Debug = false
+# workers
+Workers = 300
+# Limit the maximum number of IPs scanned. No limit if it is less than or equal to 0.
+ScannedLimit = 0
+# Limit the maximum number of IPs found. No limit if it is less than or equal to 0.
+FoundLimit = 10
 
 [Ping]
 # avaivable values: icmp, tcp, udp
@@ -58,17 +63,27 @@ all = false
 [UsedFor]
 
 [UsedFor.Cloudflare]
+# All IP ranges of cloudflare
 IPRangesFile = "./data/cloudflare.json"
+# Customized IP ranges. If the file does not exist, will use IPRangesFile
 CustomIPRangesFile = "./data/cloudflare_custom_ip_ranges.txt"
+# Output the available IPs found
 IPOutputFile = "./data/output_cloudflare.txt"
+# A boolean that turns on/off scanning for IPv6. true or false.
 WithIPv6 = false
+# URL for testing HTTPS connection
 HttpsURL = "https://yezheng.pages.dev"
 
 [UsedFor.GoogleTranslate]
+# All IP ranges of google
 IPRangesFile = "./data/goog.json"
+# Customized IP ranges. If the file does not exist, will use IPRangesFile
 CustomIPRangesFile = "./data/google_translate_custom_ip_ranges.txt"
+# Output the available IPs found
 IPOutputFile = "./data/output_google_translate.txt"
+# # boolean that turns on/off scanning for IPv6. true or false.
 WithIPv6 = false
+# URL for testing HTTPS connection
 HttpsURL = "https://translate.google.com"
 ```
 

common/config.go

@@ -4,9 +4,11 @@ import "time"
 
 type Config struct {
 	General struct {
-		Debug   bool
-		Workers int
-		UsedFor string
+		UsedFor      string
+		Debug        bool
+		Workers      int
+		ScannedLimit int
+		FoundLimit   int
 	}
 	Ping struct {
 		Protocol string

common/record.go

@@ -35,6 +35,18 @@ func (records *ScanRecordArray) Swap(i, j int) {
 	(*records)[j] = tmp
 }
 
+func (result *ScanResult) Scanned() int {
+	result.recordMutex.Lock()
+	defer result.recordMutex.Unlock()
+	return int(result.scanned)
+}
+
+func (result *ScanResult) Found() int {
+	result.recordMutex.Lock()
+	defer result.recordMutex.Unlock()
+	return len(result.scanRecords)
+}
+
 func (result *ScanResult) AddRecord(record *ScanRecord) {
 	result.recordMutex.Lock()
 	if result.scanRecords == nil {

common/scan.go

@@ -115,6 +115,11 @@ func reqOneIP(destination string, destinationPort uint16, config *Config, record
 func testOne(ch chan string, config *Config, scanResult *ScanResult, wg *sync.WaitGroup) {
 	defer wg.Done()
 	for destination := range ch {
+		if (config.General.ScannedLimit > 0 && config.General.ScannedLimit < scanResult.Scanned()) ||
+			(config.General.FoundLimit > 0 && config.General.FoundLimit < scanResult.Found()) {
+			slog.Debug("The limit number of scans from configuration file has been reached, stop scanning!")
+			return
+		}
 		record := new(ScanRecord)
 		destinationPort := config.Ping.Port
 		success := pingOneIP(destination, destinationPort, config, record)
@@ -153,5 +158,5 @@ func Start(config *Config) {
 		return scanRecords[i].HttpRTT < scanRecords[j].HttpRTT
 	})
 	writeToFile(scanRecords, config)
-	printResult(scanRecords)
+	printResult(scanRecords, config)
 }

common/utils.go

@@ -2,10 +2,14 @@ package common
 
 import (
 	"bufio"
+	"fmt"
 	"github.com/csyezheng/ip-scanner/usedfor"
+	"io"
 	"log/slog"
 	"net/netip"
 	"os"
+	"runtime"
+	"strings"
 	"sync"
 )
 
@@ -90,21 +94,139 @@ func writeToFile(scanRecords ScanRecordArray, config *Config) {
 	}
 	w := bufio.NewWriter(f)
 	for _, record := range scanRecords {
-		w.WriteString(record.IP + "\n")
+		_, err := w.WriteString(record.IP + "\n")
+		if err != nil {
+			slog.Error("write to output file failed", "error", err)
+		}
+	}
+	err = w.Flush()
+	if err != nil {
+		slog.Error("flush failed", "error", err)
 	}
-	w.Flush()
 }
 
-func printResult(scanRecords ScanRecordArray) {
+func printResult(scanRecords ScanRecordArray, config *Config) {
 	if len(scanRecords) == 0 {
 		slog.Info("No found available ip!")
 		return
 	}
-	for i, record := range scanRecords {
-		if i < 10 {
-			slog.Info("Scan Result:", slog.String("IP", record.IP),
-				slog.String("Protocal", record.Protocol), slog.Float64("PingRTT", record.PingRTT),
-				slog.Float64("HttpRTT", record.HttpRTT))
+	head := scanRecords
+	if len(head) > 10 {
+		head = head[:10]
+	}
+	for _, record := range head {
+		fmt.Printf("%s\t%s\t%f\t%f\n", record.IP, record.Protocol, record.PingRTT, record.HttpRTT)
+	}
+	if config.General.UsedFor == "GoogleTranslate" {
+		fastestRecord := *scanRecords[0]
+		slog.Info("The fastest IP has been found:")
+		fmt.Printf("%v\t%s\n", fastestRecord.IP, "translate.googleapis.com")
+		fmt.Printf("%v\t%s\n", fastestRecord.IP, "translate.google.com")
+		if askForConfirmation() {
+			writeToHosts(fastestRecord.IP)
 		}
 	}
 }
+
+func askForConfirmation() bool {
+	var confirm string
+	fmt.Println("Whether to write to the hosts file (yes/no):")
+	fmt.Scanln(&confirm)
+	switch strings.ToLower(confirm) {
+	case "y", "yes":
+		return true
+	case "n", "no":
+		return false
+	default:
+		slog.Info("Please type (y)es or (n)o and then press enter:")
+		return askForConfirmation()
+	}
+}
+
+// writeToHosts: only use for Google Translate
+func writeToHosts(ip string) {
+	var hostsFile string
+	switch runtime.GOOS {
+	case "windows":
+		hostsFile = "C:\\Windows\\System32\\drivers\\etc\\hosts"
+	case "darwin":
+		hostsFile = "/private/etc/hosts"
+	case "linux":
+		hostsFile = "/etc/hosts"
+	default:
+		slog.Info("Your operating system is unknown, please configure hosts yourself.")
+		return
+	}
+	backupPath := "hosts"
+	err := Copy(hostsFile, backupPath)
+	if err != nil {
+		slog.Error("Backup hosts failed, please modify the hosts file yourself.", err)
+		return
+	}
+	err = modifyHosts(hostsFile, ip)
+	if err != nil {
+		slog.Error("Modify hosts failed, please modify the hosts file yourself.", err)
+		return
+	}
+	slog.Info("Successfully written to hosts file")
+}
+
+func Copy(srcPath, dstPath string) (err error) {
+	r, err := os.Open(srcPath)
+	if err != nil {
+		return err
+	}
+	// ignore error: file was opened read-only.
+	defer r.Close()
+	w, err := os.Create(dstPath)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		err := w.Close()
+		if err != nil {
+		}
+	}()
+	_, err = io.Copy(w, r)
+	return err
+}
+
+// modifyHosts: only use for Google Translate
+func modifyHosts(hostsFile string, ip string) error {
+	f, err := os.OpenFile(hostsFile, os.O_RDWR, 0644)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	lineSeparator := "\n"
+	if runtime.GOOS == "windows" {
+		lineSeparator = "\r\n"
+	}
+	var builder strings.Builder
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.Contains(line, "translate.googleapis.com") ||
+			strings.Contains(line, "translate.google.com") {
+			continue
+		}
+		builder.WriteString(line + lineSeparator)
+	}
+	builder.WriteString("")
+	builder.WriteString(fmt.Sprintf("%s\t%s", ip, "translate.googleapis.com") + lineSeparator)
+	builder.WriteString(fmt.Sprintf("%s\t%s", ip, "translate.google.com") + lineSeparator)
+	err = f.Truncate(0)
+	if err != nil {
+		return err
+	}
+	_, err = f.Seek(0, 0)
+	if err != nil {
+		return err
+	}
+	_, err = f.WriteString(builder.String())
+	if err != nil {
+		return err
+	}
+	return nil
+}

configs/config.toml

@@ -1,9 +1,14 @@
 [General]
+# GoogleTranslate or Cloudflare
+UsedFor = "GoogleTranslate"
+# A boolean that turns on/off debug mode. true or false
 Debug = false
 # workers
 Workers = 300
-# GoogleTranslate or Cloudflare
-UsedFor = "GoogleTranslate"
+# Limit the maximum number of IPs scanned. No limit if it is less than or equal to 0.
+ScannedLimit = 0
+# Limit the maximum number of IPs found. No limit if it is less than or equal to 0.
+FoundLimit = 10
 
 [Ping]
 # avaivable values: icmp, tcp, udp
@@ -30,15 +35,25 @@ all = false
 [UsedFor]
 
 [UsedFor.Cloudflare]
+# All IP ranges of cloudflare
 IPRangesFile = "./data/cloudflare.json"
+# Customized IP ranges. If the file does not exist, will use IPRangesFile
 CustomIPRangesFile = "./data/cloudflare_custom_ip_ranges.txt"
+# Output the available IPs found
 IPOutputFile = "./data/output_cloudflare.txt"
+# A boolean that turns on/off scanning for IPv6. true or false.
 WithIPv6 = false
+# URL for testing HTTPS connection
 HttpsURL = "https://yezheng.pages.dev"
 
 [UsedFor.GoogleTranslate]
+# All IP ranges of google
 IPRangesFile = "./data/goog.json"
+# Customized IP ranges. If the file does not exist, will use IPRangesFile
 CustomIPRangesFile = "./data/google_translate_custom_ip_ranges.txt"
+# Output the available IPs found
 IPOutputFile = "./data/output_google_translate.txt"
+# # boolean that turns on/off scanning for IPv6. true or false.
 WithIPv6 = false
+# URL for testing HTTPS connection
 HttpsURL = "https://translate.google.com"