diff --git a/pkg/output/output.go b/pkg/output/output.go index 147db7e..9c96788 100644 --- a/pkg/output/output.go +++ b/pkg/output/output.go @@ -25,18 +25,24 @@ func (tw *TableWriter) Write(results []types.Result) error { } table := tablewriter.NewWriter(tw.writer) - table.SetHeader([]string{"TARGET", "SERVICE", "SPECIFICITY", "CATEGORY", "MODELS", "ERROR"}) + table.SetHeader([]string{"TARGET", "SERVICE", "SPECIFICITY", "CATEGORY", "AUTH", "MODELS", "ERROR"}) table.SetHeaderAlignment(tablewriter.ALIGN_LEFT) table.SetAlignment(tablewriter.ALIGN_LEFT) for _, result := range results { models := strings.Join(result.Models, ", ") + auth := "open" + if result.AuthRequired { + auth = "required" + } + table.Append([]string{ result.Target, result.Service, fmt.Sprintf("%d", result.Specificity), result.Category, + auth, models, result.Error, }) diff --git a/pkg/output/output_test.go b/pkg/output/output_test.go index 684adf7..35b1c98 100644 --- a/pkg/output/output_test.go +++ b/pkg/output/output_test.go @@ -50,10 +50,12 @@ func TestTableWriter_WriteSingleResult(t *testing.T) { assert.Contains(t, output, "TARGET") assert.Contains(t, output, "SERVICE") assert.Contains(t, output, "SPECIFICITY") + assert.Contains(t, output, "AUTH") assert.Contains(t, output, "https://api.openai.com") assert.Contains(t, output, "OpenAI API") assert.Contains(t, output, "75") + assert.Contains(t, output, "open") } func TestTableWriter_WriteMultipleResults(t *testing.T) { @@ -315,6 +317,48 @@ func TestNewWriter_JSONL(t *testing.T) { require.Len(t, lines, 1, "Should have 1 line") } +func TestTableWriter_AuthRequired(t *testing.T) { + buf := &bytes.Buffer{} + writer := NewTableWriter(buf) + + results := []types.Result{ + { + Target: "https://api.example.com/v1/models", + Service: "openai-compatible", + Category: "generic", + Specificity: 1, + AuthRequired: true, + }, + } + + err := writer.Write(results) + require.NoError(t, err) + + output := buf.String() + assert.Contains(t, output, "required") +} + +func TestTableWriter_AuthOpen(t *testing.T) { + buf := &bytes.Buffer{} + writer := NewTableWriter(buf) + + results := []types.Result{ + { + Target: "https://api.example.com/v1/models", + Service: "ollama", + Category: "self-hosted", + Specificity: 100, + AuthRequired: false, + }, + } + + err := writer.Write(results) + require.NoError(t, err) + + output := buf.String() + assert.Contains(t, output, "open") +} + func TestTableWriterModelsAndError(t *testing.T) { tests := []struct { name string diff --git a/pkg/scanner/scanner.go b/pkg/scanner/scanner.go index c8754b9..da184b8 100644 --- a/pkg/scanner/scanner.go +++ b/pkg/scanner/scanner.go @@ -74,17 +74,18 @@ func (s *Scanner) Scan(target string, probes []*types.Probe, augustus bool) []ty default: } - matched, matchedReq := s.matchProbe(target, p) + matched, mr := s.matchProbe(target, p) if !matched { return nil } result := types.Result{ - Target: target + matchedReq.Path, + Target: target + mr.Request.Path, Service: p.Name, - MatchedRequest: matchedReq.Path, + MatchedRequest: mr.Request.Path, Category: p.Category, Specificity: p.GetSpecificity(), + AuthRequired: mr.StatusCode == http.StatusUnauthorized || mr.StatusCode == http.StatusForbidden, } if p.Models != nil { @@ -117,63 +118,87 @@ func (s *Scanner) Scan(target string, probes []*types.Probe, augustus bool) []ty return results } -func (s *Scanner) matchProbe(target string, p *types.Probe) (bool, types.Request) { +type matchResult struct { + Request types.Request + StatusCode int +} + +func (s *Scanner) matchProbe(target string, p *types.Probe) (bool, matchResult) { if p.RequiresAll() { return s.matchProbeAll(target, p) } return s.matchProbeAny(target, p) } -func (s *Scanner) matchProbeAny(target string, p *types.Probe) (bool, types.Request) { +func (s *Scanner) matchProbeAny(target string, p *types.Probe) (bool, matchResult) { for _, req := range p.Requests { req.ApplyDefaults() - matched, err := s.DoRequest(target, req) - if err != nil || !matched { + statusCode, err := s.doRequestWithStatus(target, req) + if err != nil || statusCode == -1 { continue } - return true, req + return true, matchResult{Request: req, StatusCode: statusCode} } - return false, types.Request{} + return false, matchResult{} } -func (s *Scanner) matchProbeAll(target string, p *types.Probe) (bool, types.Request) { +func (s *Scanner) matchProbeAll(target string, p *types.Probe) (bool, matchResult) { if len(p.Requests) == 0 { - return false, types.Request{} + return false, matchResult{} } - var firstReq types.Request + var result matchResult for i, req := range p.Requests { req.ApplyDefaults() - matched, err := s.DoRequest(target, req) - if err != nil || !matched { - return false, types.Request{} + statusCode, err := s.doRequestWithStatus(target, req) + if err != nil || statusCode == -1 { + return false, matchResult{} } if i == 0 { - firstReq = req + result = matchResult{Request: req, StatusCode: statusCode} + } + if statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden { + result.StatusCode = statusCode } } - return true, firstReq + return true, result } func (s *Scanner) DoRequest(target string, req types.Request) (bool, error) { + statusCode, err := s.doRequestWithStatus(target, req) + if err != nil { + return false, err + } + if statusCode == -1 { + return false, nil + } + return true, nil +} + +// doRequestWithStatus returns (statusCode, error) where statusCode == -1 means +// the request failed or rules didn't parse. Callers use the status code to +// derive auth signals (401/403 → auth required). +func (s *Scanner) doRequestWithStatus(target string, req types.Request) (int, error) { resp, body, err := s.doHTTPRequest(target, req.Method, req.Path, req.Body, req.Headers) if err != nil { - return false, fmt.Errorf("executing request: %w", err) + return -1, fmt.Errorf("executing request: %w", err) } rules, err := req.GetRules() if err != nil { - return false, fmt.Errorf("parsing rules: %w", err) + return -1, fmt.Errorf("parsing rules: %w", err) } - matched := probe.MatchRules(resp, body, rules) - return matched, nil + if probe.MatchRules(resp, body, rules) { + return resp.StatusCode, nil + } + return -1, nil } func (s *Scanner) fetchModels(target string, cfg *types.ModelsConfig) ([]string, error) { diff --git a/pkg/scanner/scanner_test.go b/pkg/scanner/scanner_test.go index 2f60c28..4a832c6 100644 --- a/pkg/scanner/scanner_test.go +++ b/pkg/scanner/scanner_test.go @@ -3,6 +3,7 @@ package scanner import ( "crypto/tls" + "encoding/json" "fmt" "io" "net/http" @@ -1401,3 +1402,320 @@ func TestResponseSizeTruncation(t *testing.T) { // Body should be truncated to maxResponseSize assert.Equal(t, 512, len(body), "response body should be truncated at size limit") } + +// ============================================================================ +// Auth Required Signal Tests +// ============================================================================ + +func TestScan_AuthRequired_401(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "auth-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + {Type: "body.contains", Value: `"error"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.True(t, results[0].AuthRequired, "401 match should set AuthRequired") +} + +func TestScan_AuthRequired_403(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error":"forbidden"}`)) + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "auth-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 403}, + {Type: "body.contains", Value: `"error"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.True(t, results[0].AuthRequired, "403 match should set AuthRequired") +} + +func TestScan_AuthNotRequired_200(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"object":"list","data":[]}`)) + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "open-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 200}, + {Type: "body.contains", Value: `"data"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.False(t, results[0].AuthRequired, "200 match should not set AuthRequired") +} + +func TestScan_AuthRequired_FallsThrough(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/v1/models": + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"object":"list","data":[{"id":"gpt-4"}]}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "fallthrough-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + }, + }, + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 200}, + {Type: "body.contains", Value: `"data"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.False(t, results[0].AuthRequired, "should not be auth-required when 200 request matched") +} + +func TestScan_AuthRequired_JSONL_Output(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "jsonl-auth-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + {Type: "body.contains", Value: `"error"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + require.Len(t, results, 1) + + encoded, err := json.Marshal(results[0]) + require.NoError(t, err) + assert.Contains(t, string(encoded), `"auth_required":true`) +} + +func TestScan_AuthNotRequired_JSONL_Output(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`ok`)) + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "jsonl-open-test", + Category: "llm", + Specificity: 50, + Requests: []types.Request{ + { + Path: "/health", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 200}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + require.Len(t, results, 1) + + encoded, err := json.Marshal(results[0]) + require.NoError(t, err) + assert.Contains(t, string(encoded), `"auth_required":false`) +} + +func TestScan_RequireAll_AuthRequired(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/v1/models": + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + case "/v1/chat/completions": + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "require-all-auth", + Category: "llm", + Specificity: 75, + Require: "all", + Requests: []types.Request{ + { + Path: "/v1/models", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + {Type: "body.contains", Value: `"error"`}, + }, + }, + { + Path: "/v1/chat/completions", + Method: "POST", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + {Type: "body.contains", Value: `"error"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.True(t, results[0].AuthRequired, "require-all with 401 first request should set AuthRequired") +} + +func TestScan_RequireAll_AuthRequired_LaterRequest(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/": + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`FlowiseAI`)) + case "/api/v1/chatflows": + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"Unauthorized Access"}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer server.Close() + + s := NewScanner(WithTimeout(5 * time.Second)) + probes := []*types.Probe{ + { + Name: "flowise-pattern", + Category: "rag-orchestration", + Specificity: 90, + Require: "all", + Requests: []types.Request{ + { + Path: "/", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 200}, + {Type: "body.contains", Value: "flowiseai.com"}, + }, + }, + { + Path: "/api/v1/chatflows", + Method: "GET", + RawMatch: []rules.RawRule{ + {Type: "status", Value: 401}, + {Type: "body.contains", Value: `"Unauthorized Access"`}, + }, + }, + }, + }, + } + + results := s.Scan(server.URL, probes, false) + + require.Len(t, results, 1) + assert.True(t, results[0].AuthRequired, "require-all should detect auth from later 401 request even when first request is 200") + assert.Equal(t, "/", results[0].MatchedRequest, "matched request should still be the first request") +} diff --git a/pkg/types/result.go b/pkg/types/result.go index b64cf6a..c052b05 100644 --- a/pkg/types/result.go +++ b/pkg/types/result.go @@ -6,6 +6,7 @@ type Result struct { MatchedRequest string `json:"matched_request"` Category string `json:"category"` Specificity int `json:"specificity"` + AuthRequired bool `json:"auth_required"` Models []string `json:"models,omitempty"` GeneratorConfigs []GeneratorConfig `json:"generator_configs,omitempty"` Error string `json:"error,omitempty"`