Merge pull request #7 from someengineering/lloesche/token

Implement support for API tokens

auth/auth.go

@@ -1,7 +1,10 @@
 package auth
 
 import (
+	"bytes"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -40,3 +43,38 @@ func LoginAndGetJWT(apiEndpoint, username, password string) (string, error) {
 
 	return "", fmt.Errorf("JWT not found in response cookies")
 }
+
+func GetJWTFromToken(apiEndpoint, fixToken string) (string, error) {
+	tokenURL := fmt.Sprintf("%s/api/token/access", apiEndpoint)
+
+	body := map[string]string{"token": fixToken}
+	jsonBody, err := json.Marshal(body)
+	if err != nil {
+		return "", err
+	}
+
+	req, err := http.NewRequest("POST", tokenURL, bytes.NewBuffer(jsonBody))
+	if err != nil {
+		return "", err
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("received error response: %s", resp.Status)
+	}
+
+	bodyBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	jwt := string(bodyBytes)
+	return jwt, nil
+}

auth/auth_test.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -27,3 +28,41 @@ func TestLoginAndGetJWT(t *testing.T) {
 		t.Errorf("Expected JWT %s, got %s", expectedJWT, jwt)
 	}
 }
+
+func TestGetJWTFromToken(t *testing.T) {
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			t.Errorf("Expected 'POST' request, got '%s'", r.Method)
+		}
+
+		if r.Header.Get("Content-Type") != "application/json" {
+			t.Errorf("Expected 'Content-Type: application/json', got '%s'", r.Header.Get("Content-Type"))
+		}
+
+		var body map[string]string
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Errorf("Error decoding request body: %v", err)
+		}
+		expectedToken := "test_token"
+		if body["token"] != expectedToken {
+			t.Errorf("Expected token '%s', got '%s'", expectedToken, body["token"])
+		}
+
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("mock_jwt_token"))
+	}))
+	defer mockServer.Close()
+
+	apiEndpoint := mockServer.URL
+	fixToken := "test_token"
+	expectedJWT := "mock_jwt_token"
+
+	jwt, err := GetJWTFromToken(apiEndpoint, fixToken)
+	if err != nil {
+		t.Fatalf("Expected no error, got '%v'", err)
+	}
+
+	if jwt != expectedJWT {
+		t.Errorf("Expected JWT '%s', got '%s'", expectedJWT, jwt)
+	}
+}

cmd/root.go

@@ -23,6 +23,7 @@ var (
 
 	apiEndpoint string
 	fixToken    string
+	fixJWT      string
 	workspace   string
 	username    string
 	password    string
@@ -103,18 +104,26 @@ func executeSearch(cmd *cobra.Command, args []string) {
 
 	if fixToken == "" && username != "" && password != "" {
 		var err error
-		fixToken, err = auth.LoginAndGetJWT(apiEndpoint, username, password)
+		fixJWT, err = auth.LoginAndGetJWT(apiEndpoint, username, password)
 		if err != nil {
 			logrus.Errorln("Login error:", err)
 			return
 		}
 	}
-	if fixToken == "" {
+	if fixToken != "" {
+		var err error
+		fixJWT, err = auth.GetJWTFromToken(apiEndpoint, fixToken)
+		if err != nil {
+			logrus.Errorln("Login error:", err)
+			return
+		}
+	}
+	if fixJWT == "" {
 		logrus.Errorln("Either token or username and password are required")
 		os.Exit(1)
 	}
 
-	results, errs := search.SearchGraph(apiEndpoint, fixToken, workspace, searchStr, withEdges)
+	results, errs := search.SearchGraph(apiEndpoint, fixJWT, workspace, searchStr, withEdges)
 	firstResult := true
 	for result := range results {
 		var output string

search/search.go

@@ -14,7 +14,7 @@ type SearchRequest struct {
 	WithEdges bool   `json:"with_edges"`
 }
 
-func SearchGraph(apiEndpoint, fixToken, workspaceID, searchStr string, withEdges bool) (<-chan interface{}, <-chan error) {
+func SearchGraph(apiEndpoint, fixJWT, workspaceID, searchStr string, withEdges bool) (<-chan interface{}, <-chan error) {
 	results := make(chan interface{})
 	errs := make(chan error, 1)
 
@@ -39,7 +39,12 @@ func SearchGraph(apiEndpoint, fixToken, workspaceID, searchStr string, withEdges
 
 		req.Header.Set("Content-Type", "application/json")
 		req.Header.Set("Accept", "application/ndjson")
-		req.AddCookie(&http.Cookie{Name: "session_token", Value: fixToken})
+		req.AddCookie(&http.Cookie{
+			Name:     "session_token",
+			Value:    fixJWT,
+			HttpOnly: true,
+			Secure:   true,
+		})
 
 		client := &http.Client{}
 		resp, err := client.Do(req)