From 7454c885064a1c2b8f7fe2ec27dff850fdc0aad3 Mon Sep 17 00:00:00 2001
From: guessi <guessi@gmail.com>
Date: Sun, 10 May 2020 17:03:11 +0800
Subject: [PATCH] Make it possible to change timeout value

---
 config.go |  8 +++-----
 main.go   |  9 ++++++++-
 utils.go  | 13 +++++++------
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/config.go b/config.go
index b686c6f..58ab62d 100644
--- a/config.go
+++ b/config.go
@@ -1,9 +1,7 @@
 package main
 
-import "time"
-
 const (
-	dialerTimeout = 5 * time.Second
-	defaultPort   = 443
-	protocol      = "tcp"
+	defaultDialerTimeout = 5
+	defaultPort          = 443
+	protocol             = "tcp"
 )
diff --git a/main.go b/main.go
index ad07d2c..16621ab 100644
--- a/main.go
+++ b/main.go
@@ -18,9 +18,16 @@ func main() {
 				Usage:    "config file",
 				Required: true,
 			},
+			&cli.IntFlag{
+				Name:     "timeout",
+				Aliases:  []string{"t"},
+				Value:    defaultDialerTimeout,
+				Usage:    "dialer timeout in second(s)",
+				Required: false,
+			},
 		},
 		Action: func(c *cli.Context) error {
-			prettyPrintCertsInfo(c.String("config"))
+			prettyPrintCertsInfo(c.String("config"), c.Int("timeout"))
 			return nil
 		},
 	}
diff --git a/utils.go b/utils.go
index f007bfe..6c86920 100644
--- a/utils.go
+++ b/utils.go
@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/jedib0t/go-pretty/table"
 	"github.com/jedib0t/go-pretty/text"
@@ -37,10 +38,10 @@ func readConfig(config string) Config {
 	return c
 }
 
-func getPeerCertificates(h string, port int) ([]*x509.Certificate, error) {
+func getPeerCertificates(h string, port int, timeout int) ([]*x509.Certificate, error) {
 	conn, err := tls.DialWithDialer(
 		&net.Dialer{
-			Timeout: dialerTimeout,
+			Timeout: time.Duration(timeout) * time.Second,
 		},
 		protocol,
 		h+":"+strconv.Itoa(port),
@@ -58,9 +59,9 @@ func getPeerCertificates(h string, port int) ([]*x509.Certificate, error) {
 	return conn.ConnectionState().PeerCertificates, nil
 }
 
-func getCells(t table.Writer, host string, port int, wg *sync.WaitGroup) {
+func getCells(t table.Writer, host string, port, timeout int, wg *sync.WaitGroup) {
 	defer wg.Done()
-	certs, err := getPeerCertificates(host, port)
+	certs, err := getPeerCertificates(host, port, timeout)
 	if err != nil {
 		fmt.Printf("err: %s\n", err)
 		return // skip if target host invalid
@@ -82,7 +83,7 @@ func getCells(t table.Writer, host string, port int, wg *sync.WaitGroup) {
 	}
 }
 
-func prettyPrintCertsInfo(config string) {
+func prettyPrintCertsInfo(config string, timeout int) {
 	rc := readConfig(config)
 	if len(rc.Hosts) <= 0 {
 		fmt.Printf("key not found, or empty input\n")
@@ -116,7 +117,7 @@ func prettyPrintCertsInfo(config string) {
 		}
 		wg.Add(1)
 
-		go getCells(t, ts[0], p, &wg)
+		go getCells(t, ts[0], p, timeout, &wg)
 	}
 	wg.Wait()