From 4912fec990c0f4e9a3ffe0a7d62678a98f0d6ffc Mon Sep 17 00:00:00 2001
From: mstxq17 <154380808@qq.com>
Date: Thu, 19 Oct 2023 12:11:31 +0800
Subject: [PATCH] fix: rule don't work if output to file

---
 cmd/root.go      | 9 +++++++--
 core/myBuffer.go | 9 ++++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index c1c872a..9a18da0 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -337,11 +337,13 @@ func runCommand(cmd *cobra.Command, args []string) {
 							if !filterExt(_url, myUrlFilter) {
 								outputBuffer.WriteString(_url, &customStringHandler, newLine)
 								//fmt.Println(_url)
+								_url = outputBuffer.TempString
 								found[_url] = true
 							}
 						} else {
 							outputBuffer.WriteString(_url, &customStringHandler, newLine)
 							//fmt.Println(_url)
+							_url = outputBuffer.TempString
 							found[_url] = true
 						}
 					}
@@ -363,6 +365,7 @@ func runCommand(cmd *cobra.Command, args []string) {
 					if _, ok := found[_domain]; !ok {
 						outputBuffer.WriteString(_domain, &customStringHandler, newLine)
 						//fmt.Println(_domain)
+						_domain = outputBuffer.TempString
 						found[_domain] = true
 					}
 				}
@@ -385,11 +388,13 @@ func runCommand(cmd *cobra.Command, args []string) {
 						if isPrivateIP(ipWithPort) == false {
 							//fmt.Println(ipWithPort)
 							outputBuffer.WriteString(ipWithPort, &customStringHandler, newLine)
+							ipWithPort = outputBuffer.TempString
 							found[ipWithPort] = true
 						}
 					} else {
 						outputBuffer.WriteString(ipWithPort, &customStringHandler, newLine)
 						//fmt.Println(ipWithPort)
+						ipWithPort = outputBuffer.TempString
 						found[ipWithPort] = true
 					}
 				}
@@ -411,7 +416,7 @@ func runCommand(cmd *cobra.Command, args []string) {
 		}(_output)
 		writer := bufio.NewWriter(_output)
 		for key := range found {
-			_, err := writer.WriteString(key + "\n")
+			_, err := writer.WriteString(key)
 			if err != nil {
 				return
 			}
@@ -468,7 +473,7 @@ func init() {
 	rootCmd.PersistentFlags().BoolVarP(&myDomain, "domain", "d", false, "search domain from stdin or file(搜索域名)")
 	rootCmd.PersistentFlags().BoolVarP(&myRootDomain, "root", "", false, "only output the rootDomain when searching domain(只显示主域名)")
 	rootCmd.PersistentFlags().BoolVarP(&myWithPort, "port", "p", false, "only filter out domain&ip:port (保留域名&ip和端口)")
-	rootCmd.PersistentFlags().StringVarP(&myRule, "rule", "r", "", "replacement rule (替换规则 https://{}/)")
+	rootCmd.PersistentFlags().StringVarP(&myRule, "rule", "r", "", "use custom replacement rule (自定义输出替换规则 https://{}/)")
 	rootCmd.PersistentFlags().StringVarP(&myFlag, "flag", "", "{}", "replacement identification (替换标志位)")
 	rootCmd.PersistentFlags().BoolVarP(&myUrl, "url", "u", false, "search url from stdin or file(搜索URL)")
 	rootCmd.PersistentFlags().StringVarP(&myUrlFilter, "filter", "", "", "filter url with some useless ext(排除指定后缀的URL)")
diff --git a/core/myBuffer.go b/core/myBuffer.go
index 19e1358..b67217a 100644
--- a/core/myBuffer.go
+++ b/core/myBuffer.go
@@ -9,8 +9,9 @@ type stringHandler interface {
 }
 
 type MyBuffer struct {
-	buffer   *bytes.Buffer
-	IsFilter bool
+	buffer     *bytes.Buffer
+	TempString string
+	IsFilter   bool
 }
 
 type CustomStringHandler struct {
@@ -39,8 +40,10 @@ func (_bytes *MyBuffer) WriteString(s string, handler stringHandler, newLine str
 	// change the action of WriteString method
 	// 修改 WriteString 方法的行为
 	if _bytes.IsFilter {
-		return _bytes.buffer.WriteString(handler.HandleString(s) + newLine)
+		_bytes.TempString = handler.HandleString(s) + newLine
+		return _bytes.buffer.WriteString(_bytes.TempString)
 	}
+	_bytes.TempString = s + newLine
 	return _bytes.buffer.WriteString(s + newLine)
 }