From f5ba0198495072bce27d605d4b024ffdd7f54e4f Mon Sep 17 00:00:00 2001
From: estse <estse@estse.estse.com>
Date: Sat, 10 Sep 2022 02:52:29 +0200
Subject: [PATCH] Adds response status code filter

---
 README.md |  8 ++++++++
 main.go   | 13 +++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 86405d6..d0df214 100644
--- a/README.md
+++ b/README.md
@@ -68,6 +68,14 @@ Sometimes you don't care about checking HTTP if HTTPS is working. You can do tha
 ▶ cat domains.txt | httprobe --prefer-https
 ```
 
+## Filter Status Code
+
+Show only results with given response status code. You can do that with the `--status-codes` flag:
+
+```
+▶ cat domains.txt | httprobe --status-codes 200,301
+```
+
 ## Docker
 
 Build the docker container:
diff --git a/main.go b/main.go
index 60e5004..19280b2 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"strconv"
 	"sync"
 	"time"
 )
@@ -52,6 +53,10 @@ func main() {
 	var method string
 	flag.StringVar(&method, "method", "GET", "HTTP method to use")
 
+	// Filter status code
+	var statusCodes string
+	flag.StringVar(&statusCodes, "status-codes", "", "Status codes to display.")
+
 	flag.Parse()
 
 	// make an actual time.Duration out of the timeout
@@ -96,7 +101,7 @@ func main() {
 
 				// always try HTTPS first
 				withProto := "https://" + url
-				if isListening(client, withProto, method) {
+				if isListening(client, withProto, method, statusCodes) {
 					output <- withProto
 
 					// skip trying HTTP if --prefer-https is set
@@ -120,7 +125,7 @@ func main() {
 		go func() {
 			for url := range httpURLs {
 				withProto := "http://" + url
-				if isListening(client, withProto, method) {
+				if isListening(client, withProto, method, statusCodes) {
 					output <- withProto
 					continue
 				}
@@ -210,7 +215,7 @@ func main() {
 	outputWG.Wait()
 }
 
-func isListening(client *http.Client, url, method string) bool {
+func isListening(client *http.Client, url, method string, statusCodes string) bool {
 
 	req, err := http.NewRequest(method, url, nil)
 	if err != nil {
@@ -226,7 +231,7 @@ func isListening(client *http.Client, url, method string) bool {
 		resp.Body.Close()
 	}
 
-	if err != nil {
+	if err != nil || !strings.Contains(statusCodes, strconv.Itoa(resp.StatusCode)) && statusCodes != "" {
 		return false
 	}