Add internal domains flag (#68)

* feat: add internal domains flag to scan options * fix: fix docs
Bearer · Oct 21, 2022 · 19a2ac1 · 19a2ac1
1 parent 8b76166
commit 19a2ac1
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 16 deletions.
diff --git a/docs/reference/commands.md b/docs/reference/commands.md
@@ -38,6 +38,7 @@ curio scan [FLAGS] [PATH]
 - `--debug` enable debug logs
 - `--disable-domain-resolution` skip attempt to resolve detected domains during classification (default false)
 - `--domain-resolution-timeout` set timeout when attempting to resolve detected domains during classification
+- `--internal-domains` define regular expressions for better classification of private or unreachable domains eg. --internal-domains="*.my-company.com,private.sh"
 
 #### Worker Flags
 

diff --git a/pkg/classification/classification.go b/pkg/classification/classification.go
@@ -1,8 +1,6 @@
 package classification
 
 import (
-	"regexp"
-
 	"github.com/bearer/curio/pkg/classification/db"
 	"github.com/bearer/curio/pkg/classification/dependencies"
 	"github.com/bearer/curio/pkg/classification/interfaces"
@@ -26,8 +24,8 @@ type Config struct {
 func NewClassifier(config *Config) (*Classifier, error) {
 	interfacesClassifier, err := interfaces.New(
 		interfaces.Config{
-			Recipes:                db.Default(),
-			InternalDomainMatchers: []*regexp.Regexp{},
+			Recipes:         db.Default(),
+			InternalDomains: config.Config.Scan.InternalDomains,
 			DomainResolver: url.NewDomainResolver(
 				!config.Config.Scan.DisableDomainResolution,
 				config.Config.Scan.DomainResolutionTimeout,

diff --git a/pkg/classification/interfaces/interfaces.go b/pkg/classification/interfaces/interfaces.go
@@ -35,9 +35,9 @@ type Classifier struct {
 }
 
 type Config struct {
-	Recipes                []db.Recipe
-	InternalDomainMatchers []*regexp.Regexp
-	DomainResolver         *url.DomainResolver
+	Recipes         []db.Recipe
+	InternalDomains []string
+	DomainResolver  *url.DomainResolver
 }
 
 type Recipe struct {
@@ -58,8 +58,10 @@ type RecipeURLMatch struct {
 }
 
 var ErrInvalidRecipes = errors.New("invalid interface recipe")
+var ErrInvalidInternalDomainRegexp = errors.New("could not parse internal domains as regexp")
 
 func New(config Config) (*Classifier, error) {
+	// prepare regular expressions for recipes
 	var preparedRecipes []Recipe
 	for _, recipe := range config.Recipes {
 		preparedRecipe := Recipe{
@@ -81,19 +83,30 @@ func New(config Config) (*Classifier, error) {
 		preparedRecipes = append(preparedRecipes, preparedRecipe)
 	}
 
+	// parse internal domains as regular expressions
+	var internalDomainMatchers []*regexp.Regexp
+	for _, internalDomain := range config.InternalDomains {
+		internalDomainMatcher, err := regexp.Compile(internalDomain)
+		if err != nil {
+			return nil, ErrInvalidInternalDomainRegexp
+		}
+
+		internalDomainMatchers = append(internalDomainMatchers, internalDomainMatcher)
+	}
+
 	return &Classifier{
 		Recipes:                preparedRecipes,
-		InternalDomainMatchers: config.InternalDomainMatchers,
+		InternalDomainMatchers: internalDomainMatchers,
 		DomainResolver:         config.DomainResolver,
 	}, nil
 }
 
 func NewDefault() (*Classifier, error) {
 	return New(
 		Config{
-			Recipes:                db.Default(),
-			InternalDomainMatchers: []*regexp.Regexp{},
-			DomainResolver:         url.NewDomainResolverDefault(),
+			Recipes:         db.Default(),
+			InternalDomains: []string{},
+			DomainResolver:  url.NewDomainResolverDefault(),
 		},
 	)
 }

diff --git a/pkg/classification/interfaces/interfaces_test.go b/pkg/classification/interfaces/interfaces_test.go
@@ -1,7 +1,6 @@
 package interfaces_test
 
 import (
-	"regexp"
 	"testing"
 
 	"github.com/bearer/curio/pkg/classification/db"
@@ -195,10 +194,8 @@ func TestInterface(t *testing.T) {
 
 	classifier, err := interfaces.New(
 		interfaces.Config{
-			Recipes: db.Default(),
-			InternalDomainMatchers: []*regexp.Regexp{
-				regexp.MustCompile(`https://my.internal.domain.com`),
-			},
+			Recipes:         db.Default(),
+			InternalDomains: []string{"https://my.internal.domain.com"},
 		},
 	)
 	if err != nil {

diff --git a/pkg/commands/app.go b/pkg/commands/app.go
@@ -184,6 +184,7 @@ func NewConfigCommand() *cobra.Command {
 		SkipPathFlag:                &flag.SkipPathFlag,
 		DisableDomainResolutionFlag: &flag.DisableDomainResolutionFlag,
 		DomainResolutionTimeoutFlag: &flag.DomainResolutionTimeoutFlag,
+		InternalDomainsFlag:         &flag.InternalDomainsFlag,
 	}
 
 	configFlags := &flag.Flags{

diff --git a/pkg/flag/scan_flags.go b/pkg/flag/scan_flags.go
@@ -27,13 +27,20 @@ var (
 		Value:      3 * time.Second,
 		Usage:      "set timeout when attempting to resolve detected domains during classification (default 3 seconds), eg. --domain-resolution-timeout=TODO",
 	}
+	InternalDomainsFlag = Flag{
+		Name:       "internal-domains",
+		ConfigName: "scan.internal-domains",
+		Value:      []string{},
+		Usage:      "define regular expressions for better classification of private or unreachable domains eg. --internal-domains=\"/*.my-company.com/,/private.sh/\"",
+	}
 )
 
 type ScanFlagGroup struct {
 	SkipPathFlag                *Flag
 	DebugFlag                   *Flag
 	DisableDomainResolutionFlag *Flag
 	DomainResolutionTimeoutFlag *Flag
+	InternalDomainsFlag         *Flag
 }
 
 type ScanOptions struct {
@@ -42,6 +49,7 @@ type ScanOptions struct {
 	Debug                   bool          `json:"debug"`
 	DisableDomainResolution bool          `json:"disable_domain_resolution"`
 	DomainResolutionTimeout time.Duration `json:"domain_resolution_timeout"`
+	InternalDomains         []string      `json:"internal_domains"`
 }
 
 func NewScanFlagGroup() *ScanFlagGroup {
@@ -50,6 +58,7 @@ func NewScanFlagGroup() *ScanFlagGroup {
 		DebugFlag:                   &DebugFlag,
 		DisableDomainResolutionFlag: &DisableDomainResolutionFlag,
 		DomainResolutionTimeoutFlag: &DomainResolutionTimeoutFlag,
+		InternalDomainsFlag:         &InternalDomainsFlag,
 	}
 }
 
@@ -63,6 +72,7 @@ func (f *ScanFlagGroup) Flags() []*Flag {
 		f.DebugFlag,
 		f.DisableDomainResolutionFlag,
 		f.DomainResolutionTimeoutFlag,
+		f.InternalDomainsFlag,
 	}
 }
 
@@ -77,6 +87,7 @@ func (f *ScanFlagGroup) ToOptions(args []string) (ScanOptions, error) {
 		Debug:                   getBool(f.DebugFlag),
 		DisableDomainResolution: getBool(f.DisableDomainResolutionFlag),
 		DomainResolutionTimeout: getDuration(f.DomainResolutionTimeoutFlag),
+		InternalDomains:         getStringSlice(f.InternalDomainsFlag),
 		Target:                  target,
 	}, nil
 }