From 5f8e5d0dc5df536fec7cc6189a3e54a6d3a9e60e Mon Sep 17 00:00:00 2001
From: Ramana Reddy <90540245+RamanaReddy0M@users.noreply.github.com>
Date: Tue, 27 Jun 2023 20:37:47 +0530
Subject: [PATCH] fix tls-grab option with ztls mode (#1259)

* fix tls-grab option with ztls mode

* using canonical concatenation

* mod tidy

---------

Co-authored-by: mzack <marco.rivoli.nvh@gmail.com>
---
 common/httpx/httpx.go |  8 ++++++--
 common/httpx/tls.go   | 39 +++++++++++++++++++++++++++++++++++++++
 go.mod                |  2 +-
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/common/httpx/httpx.go b/common/httpx/httpx.go
index e9627623c..6d3ee165f 100644
--- a/common/httpx/httpx.go
+++ b/common/httpx/httpx.go
@@ -264,8 +264,12 @@ get_response:
 	resp.Lines = len(strings.Split(respbodystr, "\n"))
 
 	if !h.Options.Unsafe && h.Options.TLSGrab {
-		// extracts TLS data if any
-		resp.TLSData = h.TLSGrab(httpresp)
+		if h.Options.ZTLS {
+			resp.TLSData = h.ZTLSGrab(httpresp)
+		} else {
+			// extracts TLS data if any
+			resp.TLSData = h.TLSGrab(httpresp)
+		}
 	}
 
 	resp.CSPData = h.CSPGrab(&resp)
diff --git a/common/httpx/tls.go b/common/httpx/tls.go
index cd7fd1547..14077b8bd 100644
--- a/common/httpx/tls.go
+++ b/common/httpx/tls.go
@@ -3,10 +3,13 @@ package httpx
 import (
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
 	"net"
 	"net/http"
 
 	"github.com/projectdiscovery/tlsx/pkg/tlsx/clients"
+	"github.com/projectdiscovery/tlsx/pkg/tlsx/ztls"
+	zmaptls "github.com/zmap/zcrypto/tls"
 )
 
 // versionToTLSVersionString converts tls version to version string
@@ -48,6 +51,42 @@ func (h *HTTPX) TLSGrab(r *http.Response) *clients.Response {
 	return response
 }
 
+func (h *HTTPX) ZTLSGrab(r *http.Response) *clients.Response {
+	host := r.Request.URL.Host
+	hostname, port, _ := net.SplitHostPort(host)
+	if hostname == "" {
+		hostname = host
+	}
+	if port == "" {
+		port = "443"
+	}
+	// canonical net concatenation
+	host = net.JoinHostPort(hostname, fmt.Sprint(port))
+	tlsConn, err := h.Dialer.DialTLS(r.Request.Context(), "tcp", host)
+	if err != nil {
+		return nil
+	}
+	ztlsConn, ok := (tlsConn).(*zmaptls.Conn)
+	if !ok {
+		return nil
+	}
+	ztlsState := ztlsConn.ConnectionState()
+	if len(ztlsState.PeerCertificates) == 0 {
+		return nil
+	}
+	response := &clients.Response{
+		Host:                hostname,
+		ProbeStatus:         true,
+		Port:                port,
+		Version:             versionToTLSVersionString[ztlsState.Version],
+		Cipher:              tls.CipherSuiteName(ztlsState.CipherSuite),
+		TLSConnection:       "ztls",
+		CertificateResponse: ztls.ConvertCertificateToResponse(&clients.Options{}, hostname, ztlsState.PeerCertificates[0]),
+		ServerName:          ztlsState.ServerName,
+	}
+	return response
+}
+
 func convertCertificateToResponse(hostname string, cert *x509.Certificate) *clients.CertificateResponse {
 	response := &clients.CertificateResponse{
 		SubjectAN:    cert.DNSNames,
diff --git a/go.mod b/go.mod
index ed40202e6..af2e285b7 100644
--- a/go.mod
+++ b/go.mod
@@ -48,6 +48,7 @@ require (
 	github.com/projectdiscovery/tlsx v1.1.0
 	github.com/projectdiscovery/utils v0.0.38
 	github.com/stretchr/testify v1.8.4
+	github.com/zmap/zcrypto v0.0.0-20230205235340-d51ce4775101
 	go.uber.org/multierr v1.11.0
 	golang.org/x/exp v0.0.0-20230420155640-133eef4313cb
 )
@@ -136,7 +137,6 @@ require (
 	github.com/yuin/goldmark-emoji v1.0.1 // indirect
 	github.com/yusufpapurcu/wmi v1.2.3 // indirect
 	github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248 // indirect
-	github.com/zmap/zcrypto v0.0.0-20230205235340-d51ce4775101 // indirect
 	golang.org/x/crypto v0.10.0 // indirect
 	golang.org/x/mod v0.10.0 // indirect
 	golang.org/x/oauth2 v0.9.0 // indirect