686 lines
22 KiB
Go
686 lines
22 KiB
Go
|
package handler_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"encoding/json"
|
||
|
|
"net/http"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/user-management-system/internal/auth"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestTOTPHandler_GetTOTPStatus(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpstatususer", "totpstatus@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpstatususer", "UserPass123!")
|
||
|
|
|
||
|
|
resp, body := doGet(server.URL+"/api/v1/auth/2fa/status", token)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("expected status %d, got %d, body: %s", http.StatusOK, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
|
||
|
|
var result map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(body), &result); err != nil {
|
||
|
|
t.Fatalf("failed to parse response: %v", err)
|
||
|
|
}
|
||
|
|
if result["code"] != float64(0) {
|
||
|
|
t.Errorf("expected code 0, got %v", result["code"])
|
||
|
|
}
|
||
|
|
|
||
|
|
data, ok := result["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected data in response, got %s", body)
|
||
|
|
}
|
||
|
|
if data["enabled"] != false {
|
||
|
|
t.Errorf("expected enabled=false for new user, got %v", data["enabled"])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_GetTOTPStatus_Unauthorized(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
resp, _ := doGet(server.URL+"/api/v1/auth/2fa/status", "")
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
||
|
|
t.Errorf("expected status %d, got %d", http.StatusUnauthorized, resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_SetupTOTP(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpsetupuser", "totpsetup@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpsetupuser", "UserPass123!")
|
||
|
|
|
||
|
|
resp, body := doGet(server.URL+"/api/v1/auth/2fa/setup", token)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("expected status %d, got %d, body: %s", http.StatusOK, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
|
||
|
|
var result map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(body), &result); err != nil {
|
||
|
|
t.Fatalf("failed to parse response: %v", err)
|
||
|
|
}
|
||
|
|
if result["code"] != float64(0) {
|
||
|
|
t.Errorf("expected code 0, got %v", result["code"])
|
||
|
|
}
|
||
|
|
|
||
|
|
data, ok := result["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected data in response, got %s", body)
|
||
|
|
}
|
||
|
|
if data["secret"] == nil || data["secret"] == "" {
|
||
|
|
t.Errorf("expected secret in setup response, got %+v", data)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_SetupTOTP_Unauthorized(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
resp, _ := doGet(server.URL+"/api/v1/auth/2fa/setup", "")
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
||
|
|
t.Errorf("expected status %d, got %d", http.StatusUnauthorized, resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_EnableTOTP(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpenableuser", "totpenable@test.com", "UserPass123!")
|
||
|
|
_ = userID
|
||
|
|
_ = secret
|
||
|
|
|
||
|
|
// setupEnabledTOTPUser already enables TOTP, so let's just verify the user can login with TOTP
|
||
|
|
// Actually, we need a fresh user to test enable
|
||
|
|
registerUser(server.URL, "totpenableuser2", "totpenable2@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpenableuser2", "UserPass123!")
|
||
|
|
|
||
|
|
// Setup TOTP
|
||
|
|
setupResp, setupBody := doGet(server.URL+"/api/v1/auth/2fa/setup", token)
|
||
|
|
defer setupResp.Body.Close()
|
||
|
|
if setupResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("setup failed: status=%d body=%s", setupResp.StatusCode, setupBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
var setupResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(setupBody), &setupResult); err != nil {
|
||
|
|
t.Fatalf("failed to parse setup response: %v", err)
|
||
|
|
}
|
||
|
|
setupData, ok := setupResult["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected setup data, got %s", setupBody)
|
||
|
|
}
|
||
|
|
newSecret, ok := setupData["secret"].(string)
|
||
|
|
if !ok || newSecret == "" {
|
||
|
|
t.Fatalf("expected secret in setup response, got %s", setupBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Generate valid code
|
||
|
|
code, err := auth.NewTOTPManager().GenerateCurrentCode(newSecret)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("failed to generate TOTP code: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Enable TOTP
|
||
|
|
enableResp, enableBody := doPost(server.URL+"/api/v1/auth/2fa/enable", token, map[string]interface{}{
|
||
|
|
"code": code,
|
||
|
|
})
|
||
|
|
defer enableResp.Body.Close()
|
||
|
|
|
||
|
|
if enableResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("expected status %d, got %d, body: %s", http.StatusOK, enableResp.StatusCode, enableBody)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_EnableTOTP_InvalidCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpenableinv", "totpenableinv@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpenableinv", "UserPass123!")
|
||
|
|
|
||
|
|
// Setup TOTP first
|
||
|
|
setupResp, setupBody := doGet(server.URL+"/api/v1/auth/2fa/setup", token)
|
||
|
|
defer setupResp.Body.Close()
|
||
|
|
if setupResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("setup failed: status=%d body=%s", setupResp.StatusCode, setupBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Try enable with invalid code
|
||
|
|
enableResp, enableBody := doPost(server.URL+"/api/v1/auth/2fa/enable", token, map[string]interface{}{
|
||
|
|
"code": "000000",
|
||
|
|
})
|
||
|
|
defer enableResp.Body.Close()
|
||
|
|
|
||
|
|
if enableResp.StatusCode != http.StatusUnauthorized && enableResp.StatusCode != http.StatusInternalServerError {
|
||
|
|
t.Errorf("expected status 401 or 500 for invalid code, got %d, body: %s", enableResp.StatusCode, enableBody)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_EnableTOTP_MissingCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpenablemiss", "totpenablemiss@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpenablemiss", "UserPass123!")
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/enable", token, map[string]interface{}{})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected status %d, got %d, body: %s", http.StatusBadRequest, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_DisableTOTP(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpdisableuser", "totpdisable@test.com", "UserPass123!")
|
||
|
|
|
||
|
|
// Login again to get a fresh token (since TOTP is enabled, login may require TOTP)
|
||
|
|
deviceID := "test-device"
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpdisableuser",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
if loginResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("login failed: status=%d body=%s", loginResp.StatusCode, loginBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err != nil {
|
||
|
|
t.Fatalf("failed to parse login response: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// If requires_totp, we need to verify TOTP first
|
||
|
|
loginData, ok := loginResult["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected login data, got %s", loginBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
var token string
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
code, err := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("failed to generate TOTP code: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("totp verify failed: status=%d body=%s", verifyResp.StatusCode, verifyBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err != nil {
|
||
|
|
t.Fatalf("failed to parse verify response: %v", err)
|
||
|
|
}
|
||
|
|
verifyData, ok := verifyResult["data"].(map[string]interface{})
|
||
|
|
if ok && verifyData["access_token"] != nil {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Generate valid code for disable
|
||
|
|
code, err := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("failed to generate TOTP code: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/disable", token, map[string]interface{}{
|
||
|
|
"code": code,
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("expected status %d, got %d, body: %s", http.StatusOK, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify TOTP is disabled
|
||
|
|
statusResp, statusBody := doGet(server.URL+"/api/v1/auth/2fa/status", token)
|
||
|
|
defer statusResp.Body.Close()
|
||
|
|
if statusResp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("status check failed: status=%d body=%s", statusResp.StatusCode, statusBody)
|
||
|
|
}
|
||
|
|
|
||
|
|
var statusResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(statusBody), &statusResult); err != nil {
|
||
|
|
t.Fatalf("failed to parse status response: %v", err)
|
||
|
|
}
|
||
|
|
statusData, ok := statusResult["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected status data, got %s", statusBody)
|
||
|
|
}
|
||
|
|
if statusData["enabled"] != false {
|
||
|
|
t.Errorf("expected enabled=false after disable, got %v", statusData["enabled"])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_DisableTOTP_InvalidCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpdisableinv", "totpdisableinv@test.com", "UserPass123!")
|
||
|
|
|
||
|
|
// Get token (might need TOTP verification)
|
||
|
|
deviceID := "test-device"
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpdisableinv",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
var token string
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err == nil {
|
||
|
|
if loginData, ok := loginResult["data"].(map[string]interface{}); ok {
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
code, _ := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode == http.StatusOK {
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err == nil {
|
||
|
|
if verifyData, ok := verifyResult["data"].(map[string]interface{}); ok {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/disable", token, map[string]interface{}{
|
||
|
|
"code": "000000",
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized && resp.StatusCode != http.StatusInternalServerError {
|
||
|
|
t.Errorf("expected status 401 or 500 for invalid code, got %d, body: %s", resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_VerifyTOTP(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpverifyuser", "totpverify@test.com", "UserPass123!")
|
||
|
|
|
||
|
|
// Get token (might need TOTP verification)
|
||
|
|
deviceID := "test-device"
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpverifyuser",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
var token string
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err == nil {
|
||
|
|
if loginData, ok := loginResult["data"].(map[string]interface{}); ok {
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
code, _ := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode == http.StatusOK {
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err == nil {
|
||
|
|
if verifyData, ok := verifyResult["data"].(map[string]interface{}); ok {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
code, err := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("failed to generate TOTP code: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/verify", token, map[string]interface{}{
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
t.Fatalf("expected status %d, got %d, body: %s", http.StatusOK, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
|
||
|
|
var result map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(body), &result); err != nil {
|
||
|
|
t.Fatalf("failed to parse response: %v", err)
|
||
|
|
}
|
||
|
|
if result["code"] != float64(0) {
|
||
|
|
t.Errorf("expected code 0, got %v", result["code"])
|
||
|
|
}
|
||
|
|
data, ok := result["data"].(map[string]interface{})
|
||
|
|
if !ok {
|
||
|
|
t.Fatalf("expected data in response, got %s", body)
|
||
|
|
}
|
||
|
|
if data["verified"] != true {
|
||
|
|
t.Errorf("expected verified=true, got %v", data["verified"])
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_VerifyTOTP_InvalidCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpverifyinv", "totpverifyinv@test.com", "UserPass123!")
|
||
|
|
|
||
|
|
// Get token
|
||
|
|
deviceID := "test-device"
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpverifyinv",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
var token string
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err == nil {
|
||
|
|
if loginData, ok := loginResult["data"].(map[string]interface{}); ok {
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
code, _ := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode == http.StatusOK {
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err == nil {
|
||
|
|
if verifyData, ok := verifyResult["data"].(map[string]interface{}); ok {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/verify", token, map[string]interface{}{
|
||
|
|
"code": "000000",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized && resp.StatusCode != http.StatusInternalServerError {
|
||
|
|
t.Errorf("expected status 401 or 500 for invalid code, got %d, body: %s", resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_VerifyTOTP_MissingCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpverifymiss", "totpverifymiss@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpverifymiss", "UserPass123!")
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/verify", token, map[string]interface{}{})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected status %d, got %d, body: %s", http.StatusBadRequest, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_VerifyTOTP_Unauthorized(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
resp, _ := doPost(server.URL+"/api/v1/auth/2fa/verify", "", map[string]interface{}{
|
||
|
|
"code": "123456",
|
||
|
|
"device_id": "test-device",
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
||
|
|
t.Errorf("expected status %d, got %d", http.StatusUnauthorized, resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_DisableTOTP_MissingCode(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpdisablemiss", "totpdisablemiss@test.com", "UserPass123!")
|
||
|
|
|
||
|
|
// Get token
|
||
|
|
deviceID := "test-device"
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpdisablemiss",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": deviceID,
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
var token string
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err == nil {
|
||
|
|
if loginData, ok := loginResult["data"].(map[string]interface{}); ok {
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
code, _ := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"code": code,
|
||
|
|
"device_id": deviceID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode == http.StatusOK {
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err == nil {
|
||
|
|
if verifyData, ok := verifyResult["data"].(map[string]interface{}); ok {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
resp, body := doPost(server.URL+"/api/v1/auth/2fa/disable", token, map[string]interface{}{})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected status %d, got %d, body: %s", http.StatusBadRequest, resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_DisableTOTP_Unauthorized(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
resp, _ := doPost(server.URL+"/api/v1/auth/2fa/disable", "", map[string]interface{}{
|
||
|
|
"code": "123456",
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
||
|
|
t.Errorf("expected status %d, got %d", http.StatusUnauthorized, resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_SetupTOTP_AlreadyEnabled(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
userID, secret := setupEnabledTOTPUser(t, server.URL, "totpsetupenabled", "totpsetupenabled@test.com", "UserPass123!")
|
||
|
|
_ = secret
|
||
|
|
|
||
|
|
// Get token after TOTP login
|
||
|
|
loginResp, loginBody := doPost(server.URL+"/api/v1/auth/login", "", map[string]interface{}{
|
||
|
|
"account": "totpsetupenabled",
|
||
|
|
"password": "UserPass123!",
|
||
|
|
"device_id": "test-device",
|
||
|
|
})
|
||
|
|
defer loginResp.Body.Close()
|
||
|
|
|
||
|
|
var token string
|
||
|
|
var loginResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(loginBody), &loginResult); err == nil {
|
||
|
|
if loginData, ok := loginResult["data"].(map[string]interface{}); ok {
|
||
|
|
if loginData["requires_totp"] == true {
|
||
|
|
tempToken, _ := loginData["temp_token"].(string)
|
||
|
|
code, _ := auth.NewTOTPManager().GenerateCurrentCode(secret)
|
||
|
|
verifyResp, verifyBody := doPost(server.URL+"/api/v1/auth/login/totp-verify", "", map[string]interface{}{
|
||
|
|
"user_id": userID,
|
||
|
|
"temp_token": tempToken,
|
||
|
|
"code": code,
|
||
|
|
"device_id": "test-device",
|
||
|
|
})
|
||
|
|
defer verifyResp.Body.Close()
|
||
|
|
if verifyResp.StatusCode == http.StatusOK {
|
||
|
|
var verifyResult map[string]interface{}
|
||
|
|
if err := json.Unmarshal([]byte(verifyBody), &verifyResult); err == nil {
|
||
|
|
if verifyData, ok := verifyResult["data"].(map[string]interface{}); ok {
|
||
|
|
token, _ = verifyData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
token, _ = loginData["access_token"].(string)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if token == "" {
|
||
|
|
t.Fatal("failed to get token after login")
|
||
|
|
}
|
||
|
|
|
||
|
|
// Try setup again - should still work or return appropriate response
|
||
|
|
resp, body := doGet(server.URL+"/api/v1/auth/2fa/setup", token)
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
// Setup may return 200 with new secret or error if already enabled
|
||
|
|
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("unexpected status %d, body: %s", resp.StatusCode, body)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_EnableTOTP_Unauthorized(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
resp, _ := doPost(server.URL+"/api/v1/auth/2fa/enable", "", map[string]interface{}{
|
||
|
|
"code": "123456",
|
||
|
|
})
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
||
|
|
t.Errorf("expected status %d, got %d", http.StatusUnauthorized, resp.StatusCode)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestTOTPHandler_InvalidJSON(t *testing.T) {
|
||
|
|
server, cleanup := setupHandlerTestServer(t)
|
||
|
|
defer cleanup()
|
||
|
|
|
||
|
|
registerUser(server.URL, "totpjsonuser", "totpjson@test.com", "UserPass123!")
|
||
|
|
token := getToken(server.URL, "totpjsonuser", "UserPass123!")
|
||
|
|
|
||
|
|
tests := []struct {
|
||
|
|
name string
|
||
|
|
path string
|
||
|
|
method string
|
||
|
|
}{
|
||
|
|
{"enable_invalid_json", "/api/v1/auth/2fa/enable", "POST"},
|
||
|
|
{"disable_invalid_json", "/api/v1/auth/2fa/disable", "POST"},
|
||
|
|
{"verify_invalid_json", "/api/v1/auth/2fa/verify", "POST"},
|
||
|
|
}
|
||
|
|
|
||
|
|
for _, tc := range tests {
|
||
|
|
t.Run(tc.name, func(t *testing.T) {
|
||
|
|
req, _ := http.NewRequest(tc.method, server.URL+tc.path, bytes.NewReader([]byte("not json")))
|
||
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
||
|
|
req.Header.Set("Content-Type", "application/json")
|
||
|
|
client := &http.Client{}
|
||
|
|
resp, err := client.Do(req)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("request failed: %v", err)
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusBadRequest {
|
||
|
|
t.Errorf("expected status %d for invalid JSON, got %d", http.StatusBadRequest, resp.StatusCode)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|