From c27bba786e0db2676327ce84b1cf9a314d9b8cec Mon Sep 17 00:00:00 2001
From: Ramana Reddy <ramanaredy.manda@gmail.com>
Date: Thu, 22 Feb 2024 20:11:12 +0530
Subject: [PATCH] Add health check option

---
 pkg/runner/options.go |  1 +
 pkg/runner/runner.go  |  7 +++++++
 pkg/runner/util.go    | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/pkg/runner/options.go b/pkg/runner/options.go
index 77f2068..93448b7 100644
--- a/pkg/runner/options.go
+++ b/pkg/runner/options.go
@@ -41,4 +41,5 @@ type Options struct {
 	Silent             bool
 	Verbose            bool
 	Debug              bool
+	HealthCheck        bool
 }
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
index d936d17..b3f1c3b 100644
--- a/pkg/runner/runner.go
+++ b/pkg/runner/runner.go
@@ -146,6 +146,7 @@ func ParseOptions() *Options {
 		flagset.BoolVar(&options.Silent, "silent", false, "Silent"),
 		flagset.BoolVar(&options.Verbose, "verbose", false, "Verbose"),
 		flagset.BoolVar(&options.Debug, "debug", false, "Debug"),
+		flagset.BoolVarP(&options.HealthCheck, "health-check", "hc", false, "run diagnostic check up"),
 	)
 
 	if err := flagset.Parse(); err != nil {
@@ -160,6 +161,12 @@ func ParseOptions() *Options {
 		os.Exit(0)
 	}
 
+	if options.HealthCheck {
+		showBanner()
+		gologger.Print().Msgf("%s\n", DoHealthCheck(options))
+		os.Exit(0)
+	}
+
 	// api key hierarchy: cli flag > env var > .pdcp/credential file
 	if options.PdcpAuth == "true" {
 		AuthWithPDCP()
diff --git a/pkg/runner/util.go b/pkg/runner/util.go
index 522af36..0a79af9 100644
--- a/pkg/runner/util.go
+++ b/pkg/runner/util.go
@@ -1,14 +1,48 @@
 package runner
 
 import (
+	"fmt"
+	"net"
 	"os"
 	"os/exec"
 	"runtime"
+	"strings"
 
 	"github.com/projectdiscovery/cvemap/pkg/types"
 	fileutil "github.com/projectdiscovery/utils/file"
 )
 
+func DoHealthCheck(opts Options) string {
+	var test strings.Builder
+	test.WriteString(fmt.Sprintf("Version: %s\n", Version))
+	test.WriteString(fmt.Sprintf("Operating System: %s\n", runtime.GOOS))
+	test.WriteString(fmt.Sprintf("Architecture: %s\n", runtime.GOARCH))
+	test.WriteString(fmt.Sprintf("Go Version: %s\n", runtime.Version()))
+	test.WriteString(fmt.Sprintf("Compiler: %s\n", runtime.Compiler))
+
+	var testResult string
+	c4, err := net.Dial("tcp4", "cve.projectdiscovery.io:443")
+	if err == nil && c4 != nil {
+		c4.Close()
+	}
+	testResult = "Ok"
+	if err != nil {
+		testResult = fmt.Sprintf("Ko (%s)", err)
+	}
+	test.WriteString(fmt.Sprintf("IPv4 connectivity to cve.projectdiscovery.io:443 => %s\n", testResult))
+
+	u4, err := net.Dial("udp4", "cve.projectdiscovery.io:53")
+	if err == nil && u4 != nil {
+		u4.Close()
+	}
+	testResult = "Ok"
+	if err != nil {
+		testResult = fmt.Sprintf("Ko (%s)", err)
+	}
+	test.WriteString(fmt.Sprintf("IPv4 UDP connectivity to cve.projectdiscovery.io:53 => %s\n", testResult))
+	return test.String()
+}
+
 func getLatestVersionCVSSScore(cvss types.CvssMetrics) float64 {
 	var highestScore float64
 	if cvss.Cvss2 != nil {