package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"math/big"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/ribbybibby/ssl_exporter/config"
	"github.com/ribbybibby/ssl_exporter/test"
	"golang.org/x/crypto/ocsp"
)

// TestProbeHandlerHTTPS tests a typical HTTPS probe
func TestProbeHandlerHTTPS(t *testing.T) {
	server, certPEM, _, caFile, teardown, err := test.SetupHTTPSServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober: "https",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.URL, "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"https\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"https\"} 1`")
	}

	// Check notAfter and notBefore metrics for the peer certificate
	if err := checkDates(certPEM, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}

	// Check TLS version metric
	ok := strings.Contains(rr.Body.String(), "ssl_tls_version_info{version=\"TLS 1.3\"} 1")
	if !ok {
		t.Errorf("expected `ssl_tls_version_info{version=\"TLS 1.3\"} 1`")
	}

	// Check that empty OCSP response is reported
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_stapled 0"); !ok {
		t.Errorf("expected `ssl_ocsp_response_stapled 0`")
	}
}

// TestProbeHandlerHTTPSTimeout tests that the probe respects the timeout set in
// the module configuration
func TestProbeHandlerHTTPSTimeout(t *testing.T) {
	server, _, _, caFile, teardown, err := test.SetupHTTPSServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(3 * time.Second)
		fmt.Fprintln(w, "Hello world")
	})

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober:  "https",
				Timeout: 1 * time.Second,
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.URL, "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}

	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"https\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"https\"} 1`")
	}
}

// TestProbeHandlerHTTPSVerifiedChains checks that metrics are generated
// correctly for the verified chains
func TestProbeHandlerHTTPSVerifiedChains(t *testing.T) {
	rootPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf(err.Error())
	}

	rootCertExpiry := time.Now().AddDate(0, 0, 5)
	rootCertTmpl := test.GenerateCertificateTemplate(rootCertExpiry)
	rootCertTmpl.IsCA = true
	rootCertTmpl.SerialNumber = big.NewInt(1)
	rootCert, rootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(rootCertTmpl, rootPrivateKey)

	olderRootCertExpiry := time.Now().AddDate(0, 0, 3)
	olderRootCertTmpl := test.GenerateCertificateTemplate(olderRootCertExpiry)
	olderRootCertTmpl.IsCA = true
	olderRootCertTmpl.SerialNumber = big.NewInt(2)
	olderRootCert, olderRootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(olderRootCertTmpl, rootPrivateKey)

	oldestRootCertExpiry := time.Now().AddDate(0, 0, 1)
	oldestRootCertTmpl := test.GenerateCertificateTemplate(oldestRootCertExpiry)
	oldestRootCertTmpl.IsCA = true
	oldestRootCertTmpl.SerialNumber = big.NewInt(3)
	oldestRootCert, oldestRootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(oldestRootCertTmpl, rootPrivateKey)

	serverCertExpiry := time.Now().AddDate(0, 0, 4)
	serverCertTmpl := test.GenerateCertificateTemplate(serverCertExpiry)
	serverCertTmpl.SerialNumber = big.NewInt(4)
	serverCert, serverCertPem, serverKey := test.GenerateSignedCertificate(serverCertTmpl, olderRootCert, rootPrivateKey)

	verifiedChains := [][]*x509.Certificate{
		[]*x509.Certificate{
			serverCert,
			rootCert,
		},
		[]*x509.Certificate{
			serverCert,
			olderRootCert,
		},
		[]*x509.Certificate{
			serverCert,
			oldestRootCert,
		},
	}

	caCertPem := bytes.Join([][]byte{oldestRootCertPem, olderRootCertPem, rootCertPem}, []byte(""))

	server, caFile, teardown, err := test.SetupHTTPSServerWithCertAndKey(
		caCertPem,
		serverCertPem,
		serverKey,
	)
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober: "https",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.URL, "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check verifiedNotAfter and verifiedNotBefore metrics
	if err := checkVerifiedChainDates(verifiedChains, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}
}

func TestProbeHandlerHTTPSNoServer(t *testing.T) {
	rr, err := probe("localhost:6666", "https", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

// TestProbeHandlerHTTPSEmptyTarget tests a https probe with an empty target
func TestProbeHandlerHTTPSEmptyTarget(t *testing.T) {
	rr, err := probe("", "https", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	if rr.Code != 400 {
		t.Fatalf("expected 400 status code, got %v", rr.Code)
	}
}

// TestProbeHandlerHTTPSSpaces tests an invalid address with spaces in it
func TestProbeHandlerHTTPSSpaces(t *testing.T) {
	rr, err := probe("with spaces", "https", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

// TestProbeHandlerHTTPSHTTP tests a https probe against a http server
func TestProbeHandlerHTTPSHTTP(t *testing.T) {
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "Hello world")
	}))

	server.Start()
	defer server.Close()

	u, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf(err.Error())
	}

	rr, err := probe(u.Host, "https", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

func TestProbeHandlerHTTPSClientAuthWrongClientCert(t *testing.T) {
	server, serverCertPEM, _, caFile, teardown, err := test.SetupHTTPSServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	// Configure client auth on the server
	certPool := x509.NewCertPool()
	certPool.AppendCertsFromPEM(serverCertPEM)

	server.TLS.ClientAuth = tls.RequireAndVerifyClientCert
	server.TLS.RootCAs = certPool
	server.TLS.ClientCAs = certPool

	server.StartTLS()
	defer server.Close()

	// Create a different cert/key pair that won't be accepted by the server
	certPEM, keyPEM := test.GenerateTestCertificate(time.Now().AddDate(0, 0, 1))

	// Create cert file
	certFile, err := test.WriteFile("cert.pem", certPEM)
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer os.Remove(certFile)

	// Create key file
	keyFile, err := test.WriteFile("key.pem", keyPEM)
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer os.Remove(keyFile)

	conf := &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober: "https",
				TLSConfig: config.TLSConfig{
					CAFile:   caFile,
					CertFile: certFile,
					KeyFile:  keyFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

// TestProbeHandlerTCP tests a typical TCP probe
func TestProbeHandlerTCP(t *testing.T) {
	server, certPEM, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}

	// Check notAfter and notBefore metrics
	if err := checkDates(certPEM, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}

	// Check that empty OCSP response is reported
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_stapled 0"); !ok {
		t.Errorf("expected `ssl_ocsp_response_stapled 0`")
	}
}

// TestProbeHandlerTCPTimeout tests that the probe respects the timeout set in
// the module configuration
func TestProbeHandlerTCPTimeout(t *testing.T) {
	server, _, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLSWait(3 * time.Second)
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober:  "tcp",
				Timeout: 1 * time.Second,
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}

	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}

}

// TestProbeHandlerTCPOCSP tests a TCP probe with OCSP stapling
func TestProbeHandlerTCPOCSP(t *testing.T) {
	server, certPEM, keyPEM, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	block, _ := pem.Decode([]byte(certPEM))
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		t.Fatalf(err.Error())
	}

	block, _ = pem.Decode([]byte(keyPEM))
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		t.Fatalf(err.Error())
	}

	resp, err := ocsp.CreateResponse(cert, cert, ocsp.Response{SerialNumber: big.NewInt(64), Status: 1}, key)
	if err != nil {
		t.Fatalf(err.Error())
	}
	server.TLS.Certificates[0].OCSPStaple = resp

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	parsedResponse, err := ocsp.ParseResponse(resp, nil)
	if err != nil {
		t.Fatalf(err.Error())
	}

	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_stapled 1"); !ok {
		t.Errorf("expected `ssl_ocsp_response_stapled 1`")
	}

	status := strconv.Itoa(parsedResponse.Status)
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_status "+status); !ok {
		t.Errorf("expected `ssl_ocsp_response_status " + status + "`")
	}

	nextUpdate := strconv.FormatFloat(float64(parsedResponse.NextUpdate.Unix()), 'g', -1, 64)
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_next_update "+nextUpdate); !ok {
		t.Errorf("expected `ssl_ocsp_response_next_update " + nextUpdate + "`")
	}

	thisUpdate := strconv.FormatFloat(float64(parsedResponse.ThisUpdate.Unix()), 'g', -1, 64)
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_this_update "+thisUpdate); !ok {
		t.Errorf("expected `ssl_ocsp_response_this_update " + thisUpdate + "`")
	}

	revokedAt := strconv.FormatFloat(float64(parsedResponse.RevokedAt.Unix()), 'g', -1, 64)
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_revoked_at "+revokedAt); !ok {
		t.Errorf("expected `ssl_ocsp_response_revoked_at " + revokedAt + "`")
	}

	producedAt := strconv.FormatFloat(float64(parsedResponse.ProducedAt.Unix()), 'g', -1, 64)
	if ok := strings.Contains(rr.Body.String(), "ssl_ocsp_response_produced_at "+producedAt); !ok {
		t.Errorf("expected `ssl_ocsp_response_produced_at " + producedAt + "`")
	}
}

// TestProbeHandlerTCPVerifiedChains checks that metrics are generated
// correctly for the verified chains
func TestProbeHandlerTCPVerifiedChains(t *testing.T) {
	rootPrivateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf(err.Error())
	}

	rootCertExpiry := time.Now().AddDate(0, 0, 5)
	rootCertTmpl := test.GenerateCertificateTemplate(rootCertExpiry)
	rootCertTmpl.IsCA = true
	rootCertTmpl.SerialNumber = big.NewInt(1)
	rootCert, rootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(rootCertTmpl, rootPrivateKey)

	olderRootCertExpiry := time.Now().AddDate(0, 0, 3)
	olderRootCertTmpl := test.GenerateCertificateTemplate(olderRootCertExpiry)
	olderRootCertTmpl.IsCA = true
	olderRootCertTmpl.SerialNumber = big.NewInt(2)
	olderRootCert, olderRootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(olderRootCertTmpl, rootPrivateKey)

	oldestRootCertExpiry := time.Now().AddDate(0, 0, 1)
	oldestRootCertTmpl := test.GenerateCertificateTemplate(oldestRootCertExpiry)
	oldestRootCertTmpl.IsCA = true
	oldestRootCertTmpl.SerialNumber = big.NewInt(3)
	oldestRootCert, oldestRootCertPem := test.GenerateSelfSignedCertificateWithPrivateKey(oldestRootCertTmpl, rootPrivateKey)

	serverCertExpiry := time.Now().AddDate(0, 0, 4)
	serverCertTmpl := test.GenerateCertificateTemplate(serverCertExpiry)
	serverCertTmpl.SerialNumber = big.NewInt(4)
	serverCert, serverCertPem, serverKey := test.GenerateSignedCertificate(serverCertTmpl, olderRootCert, rootPrivateKey)

	verifiedChains := [][]*x509.Certificate{
		[]*x509.Certificate{
			serverCert,
			rootCert,
		},
		[]*x509.Certificate{
			serverCert,
			olderRootCert,
		},
		[]*x509.Certificate{
			serverCert,
			oldestRootCert,
		},
	}

	caCertPem := bytes.Join([][]byte{oldestRootCertPem, olderRootCertPem, rootCertPem}, []byte(""))

	server, caFile, teardown, err := test.SetupTCPServerWithCertAndKey(
		caCertPem,
		serverCertPem,
		serverKey,
	)
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check verifiedNotAfter and verifiedNotBefore metrics
	if err := checkVerifiedChainDates(verifiedChains, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}
}

//  TestProbeHandlerTCPNoServer tests against a tcp server that doesn't exist
func TestProbeHandlerTCPNoServer(t *testing.T) {
	rr, err := probe("localhost:6666", "tcp", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

// TestProbeHandlerTCPEmptyTarget tests a TCP probe with an empty target
func TestProbeHandlerTCPEmptyTarget(t *testing.T) {
	rr, err := probe("", "tcp", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	if rr.Code != 400 {
		t.Fatalf("expected 400 status code, got %v", rr.Code)
	}
}

// TestProbeHandlerTCPSpaces tests an invalid address with spaces in it
func TestProbeHandlerTCPSpaces(t *testing.T) {
	rr, err := probe("with spaces", "tcp", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

// TestProbeHandlerTCPHTTP tests a tcp probe against a HTTP server
func TestProbeHandlerTCPHTTP(t *testing.T) {
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "Hello world")
	}))

	server.Start()
	defer server.Close()

	u, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf(err.Error())
	}

	rr, err := probe(u.Host, "tcp", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

func TestProbeHandlerTCPExpired(t *testing.T) {
	server, _, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	// Create a certificate with a notAfter date in the past
	certPEM, keyPEM := test.GenerateTestCertificate(time.Now().AddDate(0, 0, -1))
	testcert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		t.Fatalf(err.Error())
	}
	server.TLS.Certificates = []tls.Certificate{testcert}

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}
}

func TestProbeHandlerTCPExpiredInsecure(t *testing.T) {
	server, _, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	// Create a certificate with a notAfter date in the past
	certPEM, keyPEM := test.GenerateTestCertificate(time.Now().AddDate(0, 0, -1))
	testcert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		t.Fatalf(err.Error())
	}
	server.TLS.Certificates = []tls.Certificate{testcert}

	server.StartTLS()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"tcp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile:             caFile,
					InsecureSkipVerify: true,
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "tcp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1")
	if !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}
}

// TestProbeHandlerDefaultModule tests that the default module uses the tcp prober
func TestProbeHandlerDefaultModule(t *testing.T) {
	rr, err := probe("localhost:6666", "", config.DefaultConfig)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}
}

func TestProbeHandlerProxy(t *testing.T) {
	server, _, _, caFile, teardown, err := test.SetupHTTPSServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartTLS()
	defer server.Close()

	// Test with a proxy that doesn't exist first
	badProxyURL, err := url.Parse("http://localhost:6666")
	if err != nil {
		t.Fatalf(err.Error())
	}

	conf := &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober: "https",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
				HTTPS: config.HTTPSProbe{
					// Check with a bad proxy url initially
					ProxyURL: config.URL{URL: badProxyURL},
				},
			},
		},
	}

	rr, err := probe(server.URL, "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 0"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 0`")
	}

	// Test with an actual proxy server
	proxyServer, err := test.SetupHTTPProxyServer()
	if err != nil {
		t.Fatalf(err.Error())
	}

	proxyServer.Start()
	defer proxyServer.Close()

	proxyURL, err := url.Parse(proxyServer.URL)
	if err != nil {
		t.Fatalf(err.Error())
	}

	conf = &config.Config{
		Modules: map[string]config.Module{
			"https": config.Module{
				Prober: "https",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
				HTTPS: config.HTTPSProbe{
					// Check with a valid URL
					ProxyURL: config.URL{URL: proxyURL},
				},
			},
		},
	}

	rr, err = probe(server.URL, "https", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}
}

// TestProbeHandlerTCPStartTLSSMTP tests STARTTLS with a smtp server
func TestProbeHandlerTCPStartTLSSMTP(t *testing.T) {
	server, certPEM, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartSMTP()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"smtp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
				TCP: config.TCPProbe{
					StartTLS: "smtp",
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "smtp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}

	// Check notAfter and notBefore metrics
	if err := checkDates(certPEM, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}
}

// TestProbeHandlerTCPStartTLSFTP tests STARTTLS with a ftp server
func TestProbeHandlerTCPStartTLSFTP(t *testing.T) {
	server, certPEM, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartFTP()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"ftp": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
				TCP: config.TCPProbe{
					StartTLS: "ftp",
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "ftp", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}

	// Check notAfter and notBefore metrics
	if err := checkDates(certPEM, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}
}

// TestProbeHandlerTCPStartTLSIMAP tests STARTTLS with an imap server
func TestProbeHandlerTCPStartTLSIMAP(t *testing.T) {
	server, certPEM, _, caFile, teardown, err := test.SetupTCPServer()
	if err != nil {
		t.Fatalf(err.Error())
	}
	defer teardown()

	server.StartIMAP()
	defer server.Close()

	conf := &config.Config{
		Modules: map[string]config.Module{
			"imap": config.Module{
				Prober: "tcp",
				TLSConfig: config.TLSConfig{
					CAFile: caFile,
				},
				TCP: config.TCPProbe{
					StartTLS: "imap",
				},
			},
		},
	}

	rr, err := probe(server.Listener.Addr().String(), "imap", conf)
	if err != nil {
		t.Fatalf(err.Error())
	}

	// Check success metric
	if ok := strings.Contains(rr.Body.String(), "ssl_tls_connect_success 1"); !ok {
		t.Errorf("expected `ssl_tls_connect_success 1`")
	}

	// Check probe metric
	if ok := strings.Contains(rr.Body.String(), "ssl_prober{prober=\"tcp\"} 1"); !ok {
		t.Errorf("expected `ssl_prober{prober=\"tcp\"} 1`")
	}

	// Check notAfter and notBefore metrics
	if err := checkDates(certPEM, rr.Body.String()); err != nil {
		t.Errorf(err.Error())
	}
}

func checkDates(certPEM []byte, body string) error {
	// Check notAfter and notBefore metrics
	block, _ := pem.Decode(certPEM)
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return err
	}
	notAfter := strconv.FormatFloat(float64(cert.NotAfter.UnixNano()/1e9), 'g', -1, 64)
	if ok := strings.Contains(body, "ssl_cert_not_after{cn=\"example.ribbybibby.me\",dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\",emails=\",me@ribbybibby.me,example@ribbybibby.me,\",ips=\",127.0.0.1,::1,\",issuer_cn=\"example.ribbybibby.me\",ou=\",ribbybibbys org,\",serial_no=\"100\"} "+notAfter); !ok {
		return fmt.Errorf("expected `ssl_cert_not_after{cn=\"example.ribbybibby.me\",dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\",emails=\",me@ribbybibby.me,example@ribbybibby.me,\",ips=\",127.0.0.1,::1,\",issuer_cn=\"example.ribbybibby.me\",ou=\",ribbybibbys org,\",serial_no=\"100\"} " + notAfter + "`")
	}
	notBefore := strconv.FormatFloat(float64(cert.NotBefore.UnixNano()/1e9), 'g', -1, 64)
	if ok := strings.Contains(body, "ssl_cert_not_before{cn=\"example.ribbybibby.me\",dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\",emails=\",me@ribbybibby.me,example@ribbybibby.me,\",ips=\",127.0.0.1,::1,\",issuer_cn=\"example.ribbybibby.me\",ou=\",ribbybibbys org,\",serial_no=\"100\"} "+notBefore); !ok {
		return fmt.Errorf("expected `ssl_cert_not_before{cn=\"example.ribbybibby.me\",dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\",emails=\",me@ribbybibby.me,example@ribbybibby.me,\",ips=\",127.0.0.1,::1,\",issuer_cn=\"example.ribbybibby.me\",ou=\",ribbybibbys org,\",serial_no=\"100\"} " + notBefore + "`")
	}
	return nil
}

func checkVerifiedChainDates(verifiedChains [][]*x509.Certificate, body string) error {
	for i, chain := range verifiedChains {
		for _, cert := range chain {
			notAfter := strconv.FormatFloat(float64(cert.NotAfter.UnixNano()/1e9), 'g', -1, 64)
			notAfterMetric := "ssl_verified_cert_not_after{" + strings.Join([]string{
				"chain_no=\"" + strconv.Itoa(i) + "\"",
				"cn=\"example.ribbybibby.me\"",
				"dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\"",
				"emails=\",me@ribbybibby.me,example@ribbybibby.me,\"",
				"ips=\",127.0.0.1,::1,\"",
				"issuer_cn=\"example.ribbybibby.me\"",
				"ou=\",ribbybibbys org,\"",
				"serial_no=\"" + cert.SerialNumber.String() + "\"",
			}, ",") + "} " + notAfter
			if ok := strings.Contains(body, notAfterMetric); !ok {
				return fmt.Errorf("expected `%s` in: %s", notAfterMetric, body)
			}

			notBefore := strconv.FormatFloat(float64(cert.NotBefore.UnixNano()/1e9), 'g', -1, 64)
			notBeforeMetric := "ssl_verified_cert_not_before{" + strings.Join([]string{
				"chain_no=\"" + strconv.Itoa(i) + "\"",
				"cn=\"example.ribbybibby.me\"",
				"dnsnames=\",example.ribbybibby.me,example-2.ribbybibby.me,example-3.ribbybibby.me,\"",
				"emails=\",me@ribbybibby.me,example@ribbybibby.me,\"",
				"ips=\",127.0.0.1,::1,\"",
				"issuer_cn=\"example.ribbybibby.me\"",
				"ou=\",ribbybibbys org,\"",
				"serial_no=\"" + cert.SerialNumber.String() + "\"",
			}, ",") + "} " + notBefore
			if ok := strings.Contains(body, notBeforeMetric); !ok {
				return fmt.Errorf("expected `%s` in: %s", notBeforeMetric, body)
			}
		}
	}

	return nil
}

func probe(target, module string, conf *config.Config) (*httptest.ResponseRecorder, error) {
	uri := "/probe?target=" + target
	if module != "" {
		uri = uri + "&module=" + module
	}
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, err
	}

	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		probeHandler(w, r, conf)
	})

	handler.ServeHTTP(rr, req)

	return rr, nil
}
