From 4f209f227ff2e24d7ee7eb84f4a336bae6e0e7c8 Mon Sep 17 00:00:00 2001
From: Dhrumil Mistry <56185972+dmdhrumilmistry@users.noreply.github.com>
Date: Wed, 28 Aug 2024 01:46:56 +0530
Subject: [PATCH] add -pr flag path regex filtering

---
 src/cmd/offat/flag.go |  3 +++
 src/cmd/offat/main.go | 11 ++++++++++-
 src/pkg/tgen/tgen.go  | 19 +++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/cmd/offat/flag.go b/src/cmd/offat/flag.go
index 840db16..2ddd792 100644
--- a/src/cmd/offat/flag.go
+++ b/src/cmd/offat/flag.go
@@ -24,6 +24,9 @@ type FlagConfig struct {
 	QueryParams         KeyValueMap
 	Proxy               *string
 
+	// API test filter
+	PathRegex *string
+
 	// SSRF Test
 	SsrfUrl *string
 
diff --git a/src/cmd/offat/main.go b/src/cmd/offat/main.go
index 1ea33c7..e30d0ec 100644
--- a/src/cmd/offat/main.go
+++ b/src/cmd/offat/main.go
@@ -33,6 +33,8 @@ func main() {
 	config.DisableSchemaDefaultsValidation = flag.Bool("ds", false, "disable schema defaults validation for OAS files")
 	config.DisableSchemaPatternValidation = flag.Bool("dp", false, "disable schema patterns validation for OAS files")
 
+	config.PathRegex = flag.String("pr", "", "run tests for paths matching given regex pattern")
+
 	config.SsrfUrl = flag.String("ssrf", "", "injects user defined SSRF url payload in http request components")
 
 	config.RequestsPerSecond = flag.Int("r", 60, "number of requests per second")
@@ -124,8 +126,15 @@ func main() {
 		SsrfUrl: *config.SsrfUrl,
 	}
 
-	// generate and run api tests
+	// generate api tests
 	apiTests := apiTestHandler.GenerateTests()
+
+	// filter api tests
+	if *config.PathRegex != "" {
+		apiTests = apiTestHandler.FilterTests(apiTests, *config.PathRegex)
+	}
+
+	// run api tests
 	trunner.RunApiTests(&apiTestHandler, hc, hc.Client.FHClient, apiTests)
 	log.Info().Msgf("Total Requests: %d", len(apiTests))
 
diff --git a/src/pkg/tgen/tgen.go b/src/pkg/tgen/tgen.go
index f9d8715..7fe0b13 100644
--- a/src/pkg/tgen/tgen.go
+++ b/src/pkg/tgen/tgen.go
@@ -1,6 +1,8 @@
 package tgen
 
 import (
+	"regexp"
+
 	_ "github.com/OWASP/OFFAT/src/pkg/logging"
 	"github.com/OWASP/OFFAT/src/pkg/parser"
 	"github.com/OWASP/OFFAT/src/pkg/utils"
@@ -27,6 +29,23 @@ type TGenHandler struct {
 	SsrfUrl string
 }
 
+func (t *TGenHandler) FilterTests(apiTests []*ApiTest, pathRegex string) []*ApiTest {
+	var filteredTests []*ApiTest
+	for _, apiTest := range apiTests {
+		match, err := regexp.MatchString(pathRegex, apiTest.Path)
+		if err != nil {
+			log.Error().Err(err).Msgf("Failed to match %v regex with endpoint path %v", pathRegex, apiTest.Path)
+			continue
+		}
+
+		if match {
+			filteredTests = append(filteredTests, apiTest)
+		}
+	}
+
+	return filteredTests
+}
+
 func (t *TGenHandler) GenerateTests() []*ApiTest {
 	tests := []*ApiTest{}