CyberStrikeAI/internal/handler/fofa_test.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"cyberstrike-ai/internal/config"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

func TestFofaSearchUsesAPIKeyWithoutEmail(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("FOFA_API_KEY", "")
	t.Setenv("FOFA_EMAIL", "legacy@example.com")

	var receivedEmail string
	var receivedKey string
	fofaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		receivedEmail = r.URL.Query().Get("email")
		receivedKey = r.URL.Query().Get("key")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"error":false,"size":1,"page":1,"results":[["https://example.com"]]}`))
	}))
	defer fofaServer.Close()

	h := NewFofaHandler(&config.Config{
		FOFA: config.FofaConfig{
			BaseURL: fofaServer.URL,
			APIKey:  "test-api-key",
		},
	}, zap.NewNop())

	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)
	body := `{"query":"domain=\"example.com\"","fields":"host"}`
	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
	ctx.Request.Header.Set("Content-Type", "application/json")

	h.Search(ctx)

	if recorder.Code != http.StatusOK {
		t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
	if receivedEmail != "" {
		t.Fatalf("FOFA request unexpectedly included email = %q", receivedEmail)
	}
	if receivedKey != "test-api-key" {
		t.Fatalf("FOFA request key = %q, want %q", receivedKey, "test-api-key")
	}

	var response fofaSearchResponse
	if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.ResultsCount != 1 {
		t.Fatalf("results_count = %d, want 1", response.ResultsCount)
	}
}

func TestSafeFofaRequestErrorDoesNotExposeURLOrAPIKey(t *testing.T) {
	const secretURL = "https://fofa.info/api/v1/search/all?key=secret-api-key"
	err := &url.Error{
		Op:  http.MethodGet,
		URL: secretURL,
		Err: context.DeadlineExceeded,
	}

	status, message, timeout := safeFofaRequestError(err)

	if status != http.StatusGatewayTimeout {
		t.Fatalf("status = %d, want %d", status, http.StatusGatewayTimeout)
	}
	if !timeout {
		t.Fatal("timeout = false, want true")
	}
	if strings.Contains(message, "secret-api-key") || strings.Contains(message, secretURL) {
		t.Fatalf("safe error exposed request URL or API key: %q", message)
	}
}

func TestShodanSearchReportsShortfallWhenTotalExceedsMatches(t *testing.T) {
	gin.SetMode(gin.TestMode)
	t.Setenv("SHODAN_API_KEY", "")

	shodanServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/shodan/host/search" {
			t.Fatalf("unexpected path: %s", r.URL.Path)
		}
		if got := r.URL.Query().Get("key"); got != "test-shodan-key" {
			t.Fatalf("Shodan key = %q, want test-shodan-key", got)
		}
		page := r.URL.Query().Get("page")
		count := 0
		switch page {
		case "1":
			count = 100
		case "2":
			count = 3
		default:
			count = 0
		}
		matches := make([]map[string]interface{}, 0, count)
		for i := 0; i < count; i++ {
			matches = append(matches, map[string]interface{}{
				"ip_str": fmt.Sprintf("192.0.2.%d", i+1),
				"port":   80,
			})
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"total":   104,
			"matches": matches,
		})
	}))
	defer shodanServer.Close()

	h := NewFofaHandler(&config.Config{
		Shodan: config.SpaceSearchConfig{
			BaseURL: shodanServer.URL,
			APIKey:  "test-shodan-key",
		},
	}, zap.NewNop())

	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)
	body := `{"provider":"shodan","query":"product:nginx","fields":"ip_str,port","size":1000,"page":1}`
	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/fofa/search", strings.NewReader(body))
	ctx.Request.Header.Set("Content-Type", "application/json")

	h.Search(ctx)

	if recorder.Code != http.StatusOK {
		t.Fatalf("Search() status = %d, body = %s", recorder.Code, recorder.Body.String())
	}
	var response fofaSearchResponse
	if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.Total != 104 || response.ResultsCount != 103 {
		t.Fatalf("counts: total=%d results_count=%d, want 104/103", response.Total, response.ResultsCount)
	}
	if response.ExpectedCount != 104 || response.Shortfall != 1 {
		t.Fatalf("shortfall: expected=%d shortfall=%d, want 104/1", response.ExpectedCount, response.Shortfall)
	}
	if response.Warning == "" {
		t.Fatal("warning should explain shortfall")
	}
}

func TestExtractInfoCollectJSONObject(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name string
		in   string
		want string
	}{
		{
			name: "plain json",
			in:   `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`,
			want: `{"query":"title:\"CyberStrikeAI\"","warnings":[]}`,
		},
		{
			name: "fenced json",
			in:   "```json\n{\"query\":\"product:nginx\"}\n```",
			want: `{"query":"product:nginx"}`,
		},
		{
			name: "prefixed explanation",
			in:   "解析结果如下：\n{\"query\":\"ssl.cert.subject.cn:example.com\",\"explanation\":\"ok\"}\n请确认。",
			want: `{"query":"ssl.cert.subject.cn:example.com","explanation":"ok"}`,
		},
		{
			name: "braces inside string",
			in:   "结果：{\"query\":\"title:\\\"{admin}\\\"\",\"warnings\":[\"check\"]}",
			want: `{"query":"title:\"{admin}\"","warnings":["check"]}`,
		},
	}
	for _, tc := range cases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got, err := extractInfoCollectJSONObject(tc.in)
			if err != nil {
				t.Fatalf("extractInfoCollectJSONObject() error = %v", err)
			}
			if got != tc.want {
				t.Fatalf("extractInfoCollectJSONObject() = %q, want %q", got, tc.want)
			}
		})
	}
}