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) } }) } }