dashboard/backend/tests/integration/test_totp.py

@@ -0,0 +1,120 @@
+"""
+Integration tests for TOTP 2FA operations.
+"""
+import pytest
+import pyotp
+
+
+class TestTOTPSetup:
+    """Test TOTP 2FA setup endpoints."""
+
+    def test_generate_2fa_secret(self, test_app, admin_headers):
+        """Test generating a new 2FA secret."""
+        response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "secret" in data
+        assert "uri" in data
+        assert len(data["secret"]) == 32  # Base32 secret length
+        assert "otpauth://totp/" in data["uri"]
+        assert "NAS" in data["uri"] and "Dashboard" in data["uri"]
+
+    def test_generate_2fa_without_auth(self, test_app):
+        """Test generating 2FA secret without authentication."""
+        response = test_app.post("/api/auth/setup-2fa/generate")
+
+        assert response.status_code == 401
+
+    def test_verify_2fa_setup_success(self, test_app, admin_headers):
+        """Test verifying and enabling 2FA with valid code."""
+        # First generate a secret
+        gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
+        secret = gen_response.json()["secret"]
+
+        # Generate a valid TOTP code
+        totp = pyotp.TOTP(secret)
+        valid_code = totp.now()
+
+        # Verify the code
+        response = test_app.post(
+            "/api/auth/setup-2fa/verify",
+            headers=admin_headers,
+            json={"secret": secret, "code": valid_code}
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "message" in data
+        assert "enabled" in data["message"].lower()
+
+    def test_verify_2fa_setup_invalid_code(self, test_app, admin_headers):
+        """Test verifying 2FA with invalid code."""
+        response = test_app.post(
+            "/api/auth/setup-2fa/verify",
+            headers=admin_headers,
+            json={"secret": "JBSWY3DPEHPK3PXP", "code": "000000"}
+        )
+
+        assert response.status_code == 400
+        data = response.json()
+        assert "detail" in data
+        assert "invalid" in data["detail"].lower()
+
+    def test_verify_2fa_without_auth(self, test_app):
+        """Test verifying 2FA without authentication."""
+        response = test_app.post(
+            "/api/auth/setup-2fa/verify",
+            json={"secret": "JBSWY3DPEHPK3PXP", "code": "123456"}
+        )
+
+        assert response.status_code == 401
+
+    def test_disable_2fa(self, test_app, admin_headers):
+        """Test disabling 2FA."""
+        response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "message" in data
+        assert "disabled" in data["message"].lower()
+
+    def test_disable_2fa_without_auth(self, test_app):
+        """Test disabling 2FA without authentication."""
+        response = test_app.post("/api/auth/setup-2fa/disable")
+
+        assert response.status_code == 401
+
+
+class TestTOTPWorkflow:
+    """Test complete 2FA setup workflow."""
+
+    def test_complete_2fa_workflow(self, test_app, admin_headers, mock_config, temp_auth_file):
+        """Test the complete 2FA setup and disable workflow."""
+        # Step 1: Generate secret
+        gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
+        assert gen_response.status_code == 200
+        secret = gen_response.json()["secret"]
+
+        # Step 2: Verify and enable with valid code
+        totp = pyotp.TOTP(secret)
+        valid_code = totp.now()
+        verify_response = test_app.post(
+            "/api/auth/setup-2fa/verify",
+            headers=admin_headers,
+            json={"secret": secret, "code": valid_code}
+        )
+        assert verify_response.status_code == 200
+
+        # Step 3: Verify 2FA is enabled by checking auth_service
+        from auth_service import load_totp_secret
+        saved_secret = load_totp_secret()
+        assert saved_secret == secret
+
+        # Step 4: Disable 2FA
+        disable_response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
+        assert disable_response.status_code == 200
+
+        # Step 5: Verify 2FA is disabled
+        saved_secret = load_totp_secret()
+        assert saved_secret == ""