From baac7a691456fc14ca2bb34cd6f6ea238f0eadba Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sat, 22 Apr 2023 22:49:41 -0400
Subject: [PATCH] Refactor print

---
 README.md |  6 +++---
 main.go   | 32 ++++++++++++++++++--------------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 1d45d17..ad31844 100644
--- a/README.md
+++ b/README.md
@@ -9,9 +9,9 @@ Usage
 
 Recommend to run this tool locally. It may cause VPS to be flagged if you run scanner in the cloud.
 ```
-./RealiTLScanner -addr www.microsoft.com
-./RealiTLScanner -addr 20.53.203.50
-./RealiTLScanner -addr 2607:f8b0:4004:c1b::65 -thread 10
+./RealiTLScanner -addr www.microsoft.com -showFail -o
+./RealiTLScanner -addr 20.53.203.50 -showFail -o
+./RealiTLScanner -addr 2607:f8b0:4004:c1b::65 -thread 10 -showFail -o
 ```
 
 Example
diff --git a/main.go b/main.go
index 7546918..4bce16e 100644
--- a/main.go
+++ b/main.go
@@ -19,7 +19,6 @@ func main() {
 	timeOutPtr := flag.Int("timeOut", 10, "Time out of a scan")
 	showFailPtr := flag.Bool("showFail", false, "Is Show fail logs")
 	flag.Parse()
-	fmt.Println("Reality TLS Scanner running: ", *addrPtr, ":", *portPtr)
 	s := Scanner{
 		addr:           *addrPtr,
 		port:           *portPtr,
@@ -33,7 +32,7 @@ func main() {
 		s.logFile, _ = os.OpenFile("results.txt", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
 	}
 	location, _ := time.LoadLocation("Asia/Shanghai")
-	s.logFile.WriteString("start scan at : " + time.Now().In(location).String() + "\n")
+	s.Print(fmt.Sprint("Reality TLS Scanner running: ", *addrPtr, ":", *portPtr, " start scan at : ", time.Now().In(location).String()))
 	defer s.logFile.Close()
 	s.Run()
 }
@@ -59,7 +58,7 @@ func (s *Scanner) Run() {
 	}
 	conn, err := net.DialTimeout("tcp", str+":"+s.port, s.timeout)
 	if err != nil {
-		fmt.Println("Dial failed: ", err)
+		s.Print(fmt.Sprint("Dial failed: ", err))
 	} else {
 		addr = conn.RemoteAddr().(*net.TCPAddr).IP
 		line := "" + conn.RemoteAddr().String() + " \t"
@@ -70,19 +69,19 @@ func (s *Scanner) Run() {
 		})
 		err = c.Handshake()
 		if err != nil {
-			fmt.Println("", line, "TLS handshake failed: ", err)
+			s.Print(fmt.Sprint("", line, "TLS handshake failed: ", err))
 		} else {
 			state := c.ConnectionState()
 			alpn := state.NegotiatedProtocol
 			if alpn == "" {
 				alpn = "  "
 			}
-			fmt.Println(fmt.Sprint("", line, "----- Found TLS v", TlsDic[state.Version], "\tALPN ", alpn, "\t", state.PeerCertificates[0].Subject))
+			s.Print(fmt.Sprint("", line, "----- Found TLS v", TlsDic[state.Version], "\tALPN ", alpn, "\t", state.PeerCertificates[0].Subject))
 			c.Close()
 		}
 	}
 	if addr == nil {
-		fmt.Println("Invalid address format")
+		s.Print("Invalid address format")
 		return
 	}
 
@@ -115,7 +114,9 @@ func (s *Scanner) Scan(increment bool) {
 	}
 	conn, err := net.DialTimeout("tcp", str+":"+s.port, s.timeout)
 	if err != nil {
-		fmt.Println("Dial failed: ", err)
+		if s.showFail {
+			s.Print(fmt.Sprint("Dial failed: ", err))
+		}
 	} else {
 		line := "" + conn.RemoteAddr().String() + " \t"
 		conn.SetDeadline(time.Now().Add(s.timeout))
@@ -126,7 +127,7 @@ func (s *Scanner) Scan(increment bool) {
 		err = c.Handshake()
 		if err != nil {
 			if s.showFail {
-				fmt.Println("", line, "TLS handshake failed: ", err)
+				s.Print(fmt.Sprint("", line, "TLS handshake failed: ", err))
 			}
 		} else {
 			defer c.Close()
@@ -135,18 +136,21 @@ func (s *Scanner) Scan(increment bool) {
 			if alpn == "" {
 				alpn = "  "
 			}
-			outStr := fmt.Sprint("", line, "----- Found TLS v", TlsDic[state.Version], "\tALPN ", alpn, "\t", state.PeerCertificates[0].Subject)
-			if state.Version == 0x0304 && alpn == "h2" {
-				if s.output {
-					s.logFile.WriteString(outStr + "\n")
-				}
+			if s.showFail || (state.Version == 0x0304 && alpn == "h2") {
+				s.Print(fmt.Sprint("", line, "----- Found TLS v", TlsDic[state.Version], "\tALPN ", alpn, "\t", state.PeerCertificates[0].Subject))
 			}
-			fmt.Println(outStr)
 		}
 	}
 	go s.Scan(increment)
 }
 
+func (s *Scanner) Print(outStr string) {
+	if s.output {
+		s.logFile.WriteString(outStr + "\n")
+	}
+	fmt.Println(outStr)
+}
+
 func nextIP(ip net.IP, increment bool) net.IP {
 	// Convert to big.Int and increment
 	ipb := big.NewInt(0).SetBytes([]byte(ip))