Merge pull request #47 from cerberauth/refactor-and-tests

refactor: make the structure easier to understand

cmd/scan/root_test.go

@@ -0,0 +1,18 @@
+package scan_test
+
+import (
+	"testing"
+
+	"github.com/cerberauth/vulnapi/cmd/scan"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewScanCmd(t *testing.T) {
+	scanCmd := scan.NewScanCmd()
+
+	assert.NotNil(t, scanCmd)
+
+	// Assert that NewCURLScanCmd and NewOpenAPIScanCmd commands are added
+	assert.NotNil(t, scanCmd.Commands())
+	assert.Len(t, scanCmd.Commands(), 2)
+}

internal/auth/bearer_test.go

@@ -0,0 +1,84 @@
+package auth_test
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/cerberauth/vulnapi/internal/auth"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewAuthorizationBearerSecurityScheme(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	assert.Equal(t, auth.HttpType, ss.Type)
+	assert.Equal(t, auth.BearerScheme, ss.Scheme)
+	assert.Equal(t, auth.InHeader, ss.In)
+	assert.Equal(t, name, ss.Name)
+	assert.Equal(t, &value, ss.ValidValue)
+	assert.Equal(t, "", ss.AttackValue)
+}
+
+func TestBearerSecurityScheme_GetHeaders(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	headers := ss.GetHeaders()
+
+	assert.Equal(t, http.Header{
+		"Authorization": []string{"Bearer abc123"},
+	}, headers)
+}
+
+func TestBearerSecurityScheme_GetCookies(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	cookies := ss.GetCookies()
+
+	assert.Empty(t, cookies)
+}
+
+func TestBearerSecurityScheme_GetValidValue(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	validValue := ss.GetValidValue()
+
+	assert.Equal(t, value, validValue)
+}
+
+func TestBearerSecurityScheme_SetAttackValue(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	attackValue := "xyz789"
+	ss.SetAttackValue(attackValue)
+
+	assert.Equal(t, attackValue, ss.AttackValue)
+}
+
+func TestBearerSecurityScheme_GetAttackValue(t *testing.T) {
+	name := "token"
+	value := "abc123"
+
+	ss := auth.NewAuthorizationBearerSecurityScheme(name, &value)
+
+	attackValue := "xyz789"
+	ss.SetAttackValue(attackValue)
+
+	result := ss.GetAttackValue()
+
+	assert.Equal(t, attackValue, result)
+}

internal/auth/scheme_test.go

@@ -0,0 +1,43 @@
+package auth_test
+
+import (
+	"testing"
+
+	"github.com/cerberauth/vulnapi/internal/auth"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSchemeName_String(t *testing.T) {
+	scheme := auth.BasicScheme
+	assert.Equal(t, "basic", scheme.String())
+}
+
+func TestSchemeName_Set_Valid(t *testing.T) {
+	scheme := auth.SchemeName("")
+	err := scheme.Set("bearer")
+	assert.NoError(t, err)
+	assert.Equal(t, auth.BearerScheme, scheme)
+}
+
+func TestSchemeName_Set_Invalid(t *testing.T) {
+	scheme := auth.SchemeName("")
+	err := scheme.Set("invalid")
+	assert.Error(t, err)
+	assert.EqualError(t, err, `must be one of "basic", "bearer", "digest", "oauth", "privateToken"`)
+	assert.Equal(t, auth.SchemeName(""), scheme)
+}
+
+func TestSchemeName_Type(t *testing.T) {
+	scheme := auth.BasicScheme
+	assert.Equal(t, "scheme", scheme.Type())
+}
+
+func TestSchemeIn(t *testing.T) {
+	schemeIn := auth.InHeader
+	assert.Equal(t, "header", string(schemeIn))
+}
+
+func TestSchemeIn_String(t *testing.T) {
+	schemeIn := auth.InHeader
+	assert.Equal(t, "header", string(schemeIn))
+}

internal/auth/security_scheme.go

@@ -0,0 +1,48 @@
+package auth
+
+import (
+	"net/http"
+)
+
+type Type string
+
+const (
+	HttpType      Type = "http"
+	OAuth2        Type = "oauth2"
+	OpenIdConnect Type = "openIdConnect"
+	ApiKey        Type = "apiKey"
+)
+
+type SecurityScheme interface {
+	GetHeaders() http.Header
+	GetCookies() []*http.Cookie
+	GetValidValue() interface{}
+	SetAttackValue(v interface{})
+	GetAttackValue() interface{}
+}
+
+type NoAuthSecurityScheme struct{}
+
+var _ SecurityScheme = (*NoAuthSecurityScheme)(nil)
+
+func NewNoAuthSecurityScheme() *NoAuthSecurityScheme {
+	return &NoAuthSecurityScheme{}
+}
+
+func (ss *NoAuthSecurityScheme) GetHeaders() http.Header {
+	return http.Header{}
+}
+
+func (ss *NoAuthSecurityScheme) GetCookies() []*http.Cookie {
+	return []*http.Cookie{}
+}
+
+func (ss *NoAuthSecurityScheme) GetValidValue() interface{} {
+	return nil
+}
+
+func (ss *NoAuthSecurityScheme) SetAttackValue(v interface{}) {}
+
+func (ss *NoAuthSecurityScheme) GetAttackValue() interface{} {
+	return nil
+}

internal/auth/security_scheme_test.go

@@ -0,0 +1,45 @@
+package auth_test
+
+import (
+	"testing"
+
+	"github.com/cerberauth/vulnapi/internal/auth"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewNoAuthSecurityScheme(t *testing.T) {
+	ss := auth.NewNoAuthSecurityScheme()
+	assert.NotNil(t, ss)
+}
+
+func TestNoAuthSecurityScheme_GetHeaders(t *testing.T) {
+	ss := &auth.NoAuthSecurityScheme{}
+	headers := ss.GetHeaders()
+	assert.NotNil(t, headers)
+	assert.Empty(t, headers)
+}
+
+func TestNoAuthSecurityScheme_GetCookies(t *testing.T) {
+	ss := &auth.NoAuthSecurityScheme{}
+	cookies := ss.GetCookies()
+	assert.NotNil(t, cookies)
+	assert.Empty(t, cookies)
+}
+
+func TestNoAuthSecurityScheme_GetValidValue(t *testing.T) {
+	ss := &auth.NoAuthSecurityScheme{}
+	validValue := ss.GetValidValue()
+	assert.Nil(t, validValue)
+}
+
+func TestNoAuthSecurityScheme_SetAttackValue(t *testing.T) {
+	ss := &auth.NoAuthSecurityScheme{}
+	ss.SetAttackValue("attack value")
+	// No assertion as this method does not return anything
+}
+
+func TestNoAuthSecurityScheme_GetAttackValue(t *testing.T) {
+	ss := &auth.NoAuthSecurityScheme{}
+	attackValue := ss.GetAttackValue()
+	assert.Nil(t, attackValue)
+}

internal/rest_api/loader.go → internal/openapi/loader.go

@@ -1,4 +1,4 @@
-package restapi
+package openapi
 
 import (
 	"errors"

internal/auth/auth.go → internal/request/operation.go

@@ -1,24 +1,11 @@
-package auth
+package request
 
-import "net/http"
+import (
+	"net/http"
 
-type Type string
-
-const (
-	HttpType      Type = "http"
-	OAuth2        Type = "oauth2"
-	OpenIdConnect Type = "openIdConnect"
-	ApiKey        Type = "apiKey"
+	"github.com/cerberauth/vulnapi/internal/auth"
 )
 
-type SecurityScheme interface {
-	GetHeaders() http.Header
-	GetCookies() []*http.Cookie
-	GetValidValue() interface{}
-	SetAttackValue(v interface{})
-	GetAttackValue() interface{}
-}
-
 type Operations []Operation
 
 func (o Operations) Len() int      { return len(o) }
@@ -37,7 +24,7 @@ type Operation struct {
 	Headers *http.Header
 	Cookies []http.Cookie
 
-	SecuritySchemes []SecurityScheme
+	SecuritySchemes []auth.SecurityScheme
 }
 
 func (o Operation) Clone() Operation {

internal/request/operation_test.go

@@ -0,0 +1,39 @@
+package request_test
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/cerberauth/vulnapi/internal/request"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestOperation_Clone(t *testing.T) {
+	headers := http.Header{}
+	headers.Add("Content-Type", "application/json")
+
+	cookies := []http.Cookie{
+		{
+			Name:  "cookie1",
+			Value: "value1",
+		},
+		{
+			Name:  "cookie2",
+			Value: "value2",
+		},
+	}
+
+	operation := request.Operation{
+		Url:     "http://example.com",
+		Method:  "GET",
+		Headers: &headers,
+		Cookies: cookies,
+	}
+
+	clonedOperation := operation.Clone()
+
+	assert.Equal(t, operation.Url, clonedOperation.Url)
+	assert.Equal(t, operation.Method, clonedOperation.Method)
+	assert.Equal(t, operation.Headers, clonedOperation.Headers)
+	assert.Equal(t, operation.Cookies, clonedOperation.Cookies)
+}

internal/request/request.go

@@ -1,38 +1,47 @@
 package request
 
 import (
+	"io"
 	"net/http"
+	"time"
 
 	"github.com/cerberauth/vulnapi/internal/auth"
 )
 
-func NewRequest(method string, url string) (*http.Request, error) {
-	req, err := http.NewRequest(method, url, nil)
+var SharedClient = &http.Client{
+	Timeout: time.Second * 10,
+}
+
+type Request struct {
+	*http.Request
+	SecurityScheme *auth.SecurityScheme
+}
+
+func NewRequest(method string, url string, body io.Reader) (*Request, error) {
+	req, err := http.NewRequest(method, url, body)
 	if err != nil {
 		return nil, err
 	}
+	return &Request{req, nil}, nil
+}
 
-	req.Header.Set("User-Agent", "vulnapi/0.1")
-
-	return req, nil
+func (r *Request) WithSecurityScheme(ss *auth.SecurityScheme) *Request {
+	r.SecurityScheme = ss
+	return r
 }
 
-func DoRequest(client *http.Client, req *http.Request, ss auth.SecurityScheme) (*http.Request, *http.Response, error) {
-	if ss != nil {
-		for _, c := range ss.GetCookies() {
-			req.AddCookie(c)
-		}
+func (r *Request) Do() (*http.Response, error) {
+	r.Header.Set("User-Agent", "MyProject")
 
-		for n, h := range ss.GetHeaders() {
-			req.Header.Add(n, h[0])
+	if securityScheme := *r.SecurityScheme; securityScheme != nil {
+		for _, c := range securityScheme.GetCookies() {
+			r.AddCookie(c)
 		}
-	}
 
-	res, err := client.Do(req)
-	if err != nil {
-		return req, res, err
+		for n, h := range securityScheme.GetHeaders() {
+			r.Header.Add(n, h[0])
+		}
 	}
-	defer res.Body.Close()
 
-	return req, res, nil
+	return SharedClient.Do(r.Request)
 }