Merge pull request #34 from StarFleetCPTN/development

feat: Implement Two-Factor Authentication (2FA) functionality and Bug Fixes

components/profile.templ

@@ -5,8 +5,112 @@ import (
 	"github.com/starfleetcptn/gomft/internal/db"
 )
 
+// Dialog component for 2FA disable confirmation
+templ TwoFactorDisableDialog() {
+	<div id="disable-2fa-dialog" class="hidden fixed inset-0 bg-secondary-900/50 dark:bg-secondary-900/80 backdrop-blur-sm z-50 flex items-center justify-center">
+		<div class="bg-white dark:bg-secondary-800 rounded-lg shadow-xl max-w-md w-full mx-4 overflow-hidden">
+			<div class="px-6 pt-5 pb-3 text-center">
+				<div class="flex justify-center mb-2">
+					<i class="fas fa-shield-alt text-yellow-400 text-3xl"></i>
+				</div>
+				<h3 class="text-xl font-medium text-secondary-900 dark:text-secondary-100">
+					Disable Two-Factor Authentication
+				</h3>
+			</div>
+			<div class="px-6 py-4">
+				<p class="text-secondary-700 dark:text-secondary-300 mb-4">
+					Are you sure you want to disable two-factor authentication? This will make your account less secure.
+				</p>
+				<div class="space-y-4">
+					<div>
+						<label for="current-password-2fa" class="block text-sm font-medium text-secondary-700 dark:text-secondary-300 mb-1">
+							<i class="fas fa-lock mr-1"></i> Current Password
+						</label>
+						<div class="relative">
+							<div class="absolute inset-y-0 left-0 pl-3 flex items-center pointer-events-none">
+								<i class="fas fa-key text-secondary-400 dark:text-secondary-600"></i>
+							</div>
+							<input 
+								type="password" 
+								id="current-password-2fa" 
+								name="current_password"
+								class="form-input pl-10 w-full" 
+								placeholder="Enter your current password"
+								required/>
+						</div>
+					</div>
+					<div id="disable-2fa-result"></div>
+				</div>
+			</div>
+			<div class="px-6 py-4 flex justify-end space-x-3">
+				<button type="button" class="btn-secondary" onclick="hideDisable2FADialog()">
+					Cancel
+				</button>
+				<button 
+					type="button" 
+					class="btn-danger"
+					onclick="submitDisable2FA()">
+					<i class="fas fa-times mr-1"></i>
+					Disable 2FA
+				</button>
+			</div>
+		</div>
+	</div>
+}
+
 templ Profile(ctx context.Context, user db.User) {
 	@LayoutWithContext("Profile", ctx) {
+		<script>
+			// Initialize dialog functionality
+			document.addEventListener('DOMContentLoaded', function() {
+				console.log('Initializing 2FA dialog functionality');
+				
+				// Global functions for dialog control
+				window.hideDisable2FADialog = function() {
+					document.getElementById('disable-2fa-dialog').classList.add('hidden');
+					document.getElementById('current-password-2fa').value = '';
+					document.getElementById('disable-2fa-result').innerHTML = '';
+				};
+
+				window.showDisable2FADialog = function() {
+					console.log('Showing 2FA disable dialog');
+					document.getElementById('disable-2fa-dialog').classList.remove('hidden');
+				};
+
+				window.submitDisable2FA = function() {
+					const password = document.getElementById('current-password-2fa').value;
+					if (!password) {
+						document.getElementById('disable-2fa-result').innerHTML = `
+							<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded" role="alert">
+								<span class="block sm:inline">Current password is required</span>
+							</div>`;
+						return;
+					}
+
+					htmx.ajax('POST', '/profile/2fa/disable', {
+						target: '#disable-2fa-result',
+						swap: 'innerHTML',
+						values: { current_password: password }
+					});
+				};
+
+				// Close dialog when clicking outside
+				document.getElementById('disable-2fa-dialog').addEventListener('click', function(e) {
+					if (e.target === this) {
+						hideDisable2FADialog();
+					}
+				});
+
+				// Close dialog on escape key
+				document.addEventListener('keydown', function(e) {
+					if (e.key === 'Escape' && !document.getElementById('disable-2fa-dialog').classList.contains('hidden')) {
+						hideDisable2FADialog();
+					}
+				});
+			});
+		</script>
+
+		@TwoFactorDisableDialog()
 		<div class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-8">
 			<div class="flex items-center justify-between mb-6">
 				<h1 class="text-3xl font-bold text-secondary-900 dark:text-secondary-100">
@@ -45,6 +149,35 @@ templ Profile(ctx context.Context, user db.User) {
 									}
 								</dd>
 							</div>
+							<div class="flex flex-col sm:flex-row">
+								<dt class="text-sm font-medium text-secondary-500 dark:text-secondary-400 sm:w-1/3 mb-1 sm:mb-0">Two-Factor Authentication</dt>
+								<dd class="text-sm text-secondary-900 dark:text-secondary-100 sm:w-2/3">
+									if user.TwoFactorEnabled {
+										<div class="flex items-center space-x-4">
+											<span class="badge badge-success">
+												<i class="fas fa-shield-alt mr-1"></i> Enabled
+											</span>
+											<button
+												type="button"
+												class="btn-danger btn-sm"
+												onclick="showDisable2FADialog()">
+												<i class="fas fa-times mr-1"></i>
+												Disable 2FA
+											</button>
+										</div>
+									} else {
+										<div class="flex items-center space-x-4">
+											<span class="badge badge-warning">
+												<i class="fas fa-shield-alt mr-1"></i> Disabled
+											</span>
+											<a href="/profile/2fa/setup" class="btn-primary btn-sm">
+												<i class="fas fa-lock mr-1"></i>
+												Enable 2FA
+											</a>
+										</div>
+									}
+								</dd>
+							</div>
 							<div class="flex flex-col sm:flex-row">
 								<dt class="text-sm font-medium text-secondary-500 dark:text-secondary-400 sm:w-1/3 mb-1 sm:mb-0">Theme</dt>
 								<dd class="text-sm text-secondary-900 dark:text-secondary-100 sm:w-2/3">

components/providers/destination/gphotos.templ

@@ -74,6 +74,7 @@ templ GooglePhotosDestinationForm() {
                 <input id="dest_read_only" name="dest_read_only" type="checkbox" 
                     class="sr-only"
                     x-model="destReadOnly"
+                    :value="destReadOnly ? 'true' : 'false'"
                 />
                 <div class="block bg-secondary-200 dark:bg-secondary-700 w-14 h-8 rounded-full"></div>
                 <div class="dot absolute left-1 top-1 bg-white dark:bg-secondary-100 w-6 h-6 rounded-full transition"
@@ -109,6 +110,7 @@ templ GooglePhotosDestinationForm() {
                 <input id="dest_include_archived" name="dest_include_archived" type="checkbox" 
                     class="sr-only"
                     x-model="destIncludeArchived"
+                    :value="destIncludeArchived ? 'true' : 'false'"
                 />
                 <div class="block bg-secondary-200 dark:bg-secondary-700 w-14 h-8 rounded-full"></div>
                 <div class="dot absolute left-1 top-1 bg-white dark:bg-secondary-100 w-6 h-6 rounded-full transition"

components/providers/source/gphotos.templ

@@ -74,6 +74,7 @@ templ GooglePhotosSourceForm() {
                 <input id="source_read_only" name="source_read_only" type="checkbox" 
                     class="sr-only"
                     x-model="sourceReadOnly"
+                    :value="sourceReadOnly ? 'true' : 'false'"
                 />
                 <div class="block bg-secondary-200 dark:bg-secondary-700 w-14 h-8 rounded-full"></div>
                 <div class="dot absolute left-1 top-1 bg-white dark:bg-secondary-100 w-6 h-6 rounded-full transition"
@@ -109,6 +110,7 @@ templ GooglePhotosSourceForm() {
                 <input id="source_include_archived" name="source_include_archived" type="checkbox" 
                     class="sr-only"
                     x-model="sourceIncludeArchived"
+                    :value="sourceIncludeArchived ? 'true' : 'false'"
                 />
                 <div class="block bg-secondary-200 dark:bg-secondary-700 w-14 h-8 rounded-full"></div>
                 <div class="dot absolute left-1 top-1 bg-white dark:bg-secondary-100 w-6 h-6 rounded-full transition"

components/two_factor_setup.templ

@@ -0,0 +1,150 @@
+package components
+
+import "context"
+
+type TwoFactorSetupData struct {
+	QRCodeURL    string
+	Secret       string
+	BackupCodes  []string
+	ErrorMessage string
+}
+
+templ TwoFactorSetup(ctx context.Context, data TwoFactorSetupData) {
+	@LayoutWithContext("Two-Factor Authentication Setup", ctx) {
+		<div class="min-h-screen bg-secondary-50 dark:bg-secondary-900 py-12">
+			<div class="max-w-3xl mx-auto px-4 sm:px-6 lg:px-8">
+				<div class="bg-white dark:bg-secondary-800 shadow rounded-lg p-6">
+					<div class="text-center mb-8">
+						<h2 class="text-3xl font-bold text-secondary-900 dark:text-secondary-100">Set Up Two-Factor Authentication</h2>
+						<p class="mt-2 text-secondary-600 dark:text-secondary-400">Enhance your account security with 2FA</p>
+					</div>
+
+					if data.ErrorMessage != "" {
+						<div class="bg-red-100 dark:bg-red-900 border border-red-400 dark:border-red-700 text-red-700 dark:text-red-300 px-4 py-3 rounded-lg mb-6" role="alert">
+							<div class="flex items-center">
+								<i class="fas fa-exclamation-circle mr-2"></i>
+								<span class="block sm:inline">{ data.ErrorMessage }</span>
+							</div>
+						</div>
+					}
+
+					<div class="space-y-8">
+						<div>
+							<h3 class="text-xl font-semibold text-secondary-900 dark:text-secondary-100 mb-4">1. Scan QR Code</h3>
+							<p class="text-secondary-600 dark:text-secondary-400 mb-4">
+								Scan this QR code with your authenticator app (Google Authenticator, Authy, etc.)
+							</p>
+							<div class="flex justify-center mb-4">
+								<img src={ data.QRCodeURL } alt="QR Code" class="border border-secondary-200 dark:border-secondary-700 rounded-lg p-2 bg-white"/>
+							</div>
+							<div class="text-center">
+								<p class="text-sm text-secondary-600 dark:text-secondary-400">
+									Can't scan the QR code? Use this code instead:
+								</p>
+								<code class="block mt-2 p-2 bg-secondary-100 dark:bg-secondary-700 rounded font-mono text-sm">
+									{ data.Secret }
+								</code>
+							</div>
+						</div>
+
+						<div>
+							<h3 class="text-xl font-semibold text-secondary-900 dark:text-secondary-100 mb-4">2. Verify Setup</h3>
+							<form 
+								method="POST"
+								action="/profile/2fa/verify"
+								class="space-y-4"
+								x-data="{ code: '', loading: false }"
+								@submit="loading = true">
+								<div>
+									<label for="code" class="block text-sm font-medium text-secondary-700 dark:text-secondary-300 mb-1">
+										Enter the 6-digit code from your authenticator app
+									</label>
+									<input 
+										type="text" 
+										id="code"
+										name="code"
+										x-model="code"
+										class="form-input block w-full"
+										pattern="[0-9]*"
+										inputmode="numeric"
+										maxlength="6"
+										required/>
+								</div>
+								<button 
+									type="submit"
+									class="btn-primary w-full"
+									x-bind:disabled="code.length !== 6 || loading">
+									<span x-show="!loading">Verify and Enable 2FA</span>
+									<span x-show="loading" class="flex items-center justify-center">
+										<svg class="animate-spin -ml-1 mr-3 h-5 w-5 text-white" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24">
+											<circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
+											<path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
+										</svg>
+										Verifying...
+									</span>
+								</button>
+							</form>
+						</div>
+
+						if len(data.BackupCodes) > 0 {
+							<div>
+								<h3 class="text-xl font-semibold text-secondary-900 dark:text-secondary-100 mb-4">3. Save Backup Codes</h3>
+								<p class="text-secondary-600 dark:text-secondary-400 mb-4">
+									Store these backup codes in a safe place. You can use them to access your account if you lose your authenticator device.
+								</p>
+								<div class="grid grid-cols-2 gap-4 mb-4">
+									for _, code := range data.BackupCodes {
+										<div class="p-2 bg-secondary-100 dark:bg-secondary-700 rounded font-mono text-sm text-center">
+											{ code }
+										</div>
+									}
+								</div>
+								<div class="text-center">
+									<button 
+										class="btn-secondary"
+										onclick="downloadBackupCodes(this)">
+										<i class="fas fa-download mr-2"></i>
+										Download Backup Codes
+									</button>
+									<script>
+										function downloadBackupCodes(button) {
+											// Get backup codes from the displayed elements
+											const codes = Array.from(
+												document.querySelectorAll('.bg-secondary-100.dark\\:bg-secondary-700')
+											).map(el => el.textContent.trim());
+											
+											// Create content for the file
+											const content = 
+												"2FA Backup Codes - Keep these safe!\n" +
+												"=====================================\n\n" +
+												codes.join("\n") +
+												"\n\n" +
+												"Generated: " + new Date().toISOString().split('T')[0] + "\n" +
+												"These codes can be used to access your account if you lose access to your authenticator app.\n" +
+												"Each code can only be used once. Keep these codes safe and secure.";
+											
+											// Create blob and download link
+											const blob = new Blob([content], { type: 'text/plain' });
+											const url = window.URL.createObjectURL(blob);
+											const a = document.createElement('a');
+											a.href = url;
+											a.download = '2fa-backup-codes.txt';
+											
+											// Trigger download
+											document.body.appendChild(a);
+											a.click();
+											
+											// Cleanup
+											window.URL.revokeObjectURL(url);
+											document.body.removeChild(a);
+										}
+									</script>
+								</div>
+							</div>
+						}
+					</div>
+				</div>
+			</div>
+		</div>
+	}
+} 

components/two_factor_verify.templ

@@ -0,0 +1,83 @@
+package components
+
+import "context"
+
+type TwoFactorVerifyData struct {
+	ErrorMessage string
+}
+
+templ TwoFactorVerify(ctx context.Context, data TwoFactorVerifyData) {
+	@LayoutWithContext("Two-Factor Authentication", ctx) {
+		<div class="min-h-screen bg-secondary-50 dark:bg-secondary-900 py-12">
+			<div class="max-w-md mx-auto px-4 sm:px-6 lg:px-8">
+				<div class="bg-white dark:bg-secondary-800 shadow rounded-lg p-6">
+					<div class="text-center mb-8">
+						<div class="inline-flex items-center justify-center w-20 h-20 rounded-full bg-primary-100 dark:bg-primary-900 mb-4">
+							<i class="fas fa-shield-alt text-primary-600 dark:text-primary-400 text-3xl"></i>
+						</div>
+						<h2 class="text-3xl font-bold text-secondary-900 dark:text-secondary-100">Two-Factor Authentication</h2>
+						<p class="mt-2 text-secondary-600 dark:text-secondary-400">Enter the code from your authenticator app</p>
+					</div>
+
+					if data.ErrorMessage != "" {
+						<div class="bg-red-100 dark:bg-red-900 border border-red-400 dark:border-red-700 text-red-700 dark:text-red-300 px-4 py-3 rounded-lg mb-6" role="alert">
+							<div class="flex items-center">
+								<i class="fas fa-exclamation-circle mr-2"></i>
+								<span class="block sm:inline">{ data.ErrorMessage }</span>
+							</div>
+						</div>
+					}
+
+					<form 
+						method="POST" 
+						action="/login/verify"
+						class="space-y-6"
+						x-data="{ code: '', loading: false }"
+						@submit="loading = true">
+						<div>
+							<label for="code" class="block text-sm font-medium text-secondary-700 dark:text-secondary-300 mb-1">
+								Authentication Code
+							</label>
+							<div class="relative">
+								<div class="absolute inset-y-0 left-0 pl-3 flex items-center pointer-events-none">
+									<i class="fas fa-key text-secondary-400 dark:text-secondary-600"></i>
+								</div>
+								<input 
+									type="text" 
+									id="code"
+									name="code"
+									x-model="code"
+									class="form-input pl-10 w-full"
+									pattern="[0-9]*"
+									inputmode="numeric"
+									maxlength="6"
+									placeholder="Enter 6-digit code"
+									required/>
+							</div>
+						</div>
+
+						<button 
+							type="submit"
+							class="btn-primary w-full"
+							x-bind:disabled="code.length !== 6 || loading">
+							<span x-show="!loading">Verify</span>
+							<span x-show="loading" class="flex items-center justify-center">
+								<svg class="animate-spin -ml-1 mr-3 h-5 w-5 text-white" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24">
+									<circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
+									<path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
+								</svg>
+								Verifying...
+							</span>
+						</button>
+
+						<div class="text-center">
+							<a href="/backup-code" class="text-sm text-primary-600 dark:text-primary-400 hover:text-primary-500 dark:hover:text-primary-300">
+								Use a backup code instead
+							</a>
+						</div>
+					</form>
+				</div>
+			</div>
+		</div>
+	}
+} 

go.mod

@@ -4,11 +4,13 @@ go 1.24.0
 
 require (
 	github.com/a-h/templ v0.3.833
+	github.com/gin-contrib/sessions v1.0.2
 	github.com/gin-gonic/gin v1.10.0
 	github.com/glebarez/sqlite v1.11.0
 	github.com/go-gormigrate/gormigrate/v2 v2.1.3
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/joho/godotenv v1.5.1
+	github.com/pquerna/otp v1.4.0
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/stretchr/testify v1.10.0
 	golang.org/x/crypto v0.35.0
@@ -17,13 +19,13 @@ require (
 )
 
 require (
+	github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc // indirect
 	github.com/bytedance/sonic v1.12.9 // indirect
 	github.com/bytedance/sonic/loader v0.2.3 // indirect
 	github.com/cloudwego/base64x v0.1.5 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.8 // indirect
-	github.com/gin-contrib/sessions v1.0.2 // indirect
 	github.com/gin-contrib/sse v1.0.0 // indirect
 	github.com/glebarez/go-sqlite v1.21.2 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect

go.sum

@@ -1,5 +1,7 @@
 github.com/a-h/templ v0.3.833 h1:L/KOk/0VvVTBegtE0fp2RJQiBm7/52Zxv5fqlEHiQUU=
 github.com/a-h/templ v0.3.833/go.mod h1:cAu4AiZhtJfBjMY0HASlyzvkrtjnHWPeEsyGK2YYmfk=
+github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI=
+github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
 github.com/bytedance/sonic v1.12.9 h1:Od1BvK55NnewtGaJsTDeAOSnLVO2BTSLOe0+ooKokmQ=
 github.com/bytedance/sonic v1.12.9/go.mod h1:uVvFidNmlt9+wa31S1urfwwthTWteBgG0hWuoKAXTx8=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
@@ -42,6 +44,8 @@ github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVI
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
+github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ=
 github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
@@ -77,6 +81,8 @@ github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNH
 github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/pquerna/otp v1.4.0 h1:wZvl1TIVxKRThZIBiwOOHOGP/1+nZyWBil9Y2XNEDzg=
+github.com/pquerna/otp v1.4.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg=
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
 github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=

internal/auth/totp.go

@@ -0,0 +1,158 @@
+package auth
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/base32"
+	"encoding/base64"
+	"fmt"
+	"image/png"
+	"strings"
+
+	// "github.com/pquerna/otp/base32"
+	"github.com/pquerna/otp/totp"
+)
+
+const (
+	// IssuerName is the name of the issuer that appears in authenticator apps
+	IssuerName = "GoMFT"
+	// SecretSize is the size of the TOTP secret in bytes
+	SecretSize = 20
+	// BackupCodeCount is the number of backup codes to generate
+	BackupCodeCount = 8
+	// BackupCodeLength is the length of each backup code
+	BackupCodeLength = 8
+)
+
+// GenerateTOTPSecret generates a new TOTP secret for a user
+func GenerateTOTPSecret(email string) (string, string, error) {
+	// Generate TOTP key using the library
+	key, err := totp.Generate(totp.GenerateOpts{
+		Issuer:      IssuerName,
+		AccountName: email,
+	})
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate TOTP key: %v", err)
+	}
+
+	// Generate QR code image
+	var buf bytes.Buffer
+	img, err := key.Image(256, 256)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate QR code image: %v", err)
+	}
+
+	// Encode image as PNG and convert to base64
+	err = png.Encode(&buf, img)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to encode QR code image: %v", err)
+	}
+
+	// Create data URL
+	dataURL := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes()))
+
+	return key.Secret(), dataURL, nil
+}
+
+// ValidateTOTPCode validates a TOTP code against a secret
+func ValidateTOTPCode(secret string, code string) bool {
+	// Remove any spaces from the code
+	code = strings.ReplaceAll(code, " ", "")
+
+	// Use the library's Validate function
+	return totp.Validate(code, secret)
+}
+
+// GenerateBackupCodes generates a set of backup codes
+func GenerateBackupCodes() ([]string, error) {
+	codes := make([]string, BackupCodeCount)
+	for i := 0; i < BackupCodeCount; i++ {
+		// Generate random bytes
+		bytes := make([]byte, BackupCodeLength/2)
+		_, err := rand.Read(bytes)
+		if err != nil {
+			return nil, fmt.Errorf("failed to generate backup code: %v", err)
+		}
+
+		// Convert to hex string
+		codes[i] = fmt.Sprintf("%x", bytes)
+	}
+	return codes, nil
+}
+
+// ValidateBackupCode validates a backup code against a list of codes
+func ValidateBackupCode(providedCode string, storedCodes string) bool {
+	if storedCodes == "" {
+		return false
+	}
+
+	// Remove any spaces and convert to lowercase
+	providedCode = strings.ToLower(strings.ReplaceAll(providedCode, " ", ""))
+
+	// Split stored codes
+	codes := strings.Split(storedCodes, ",")
+
+	// Check if the provided code matches any stored code
+	for _, code := range codes {
+		if code == providedCode {
+			return true
+		}
+	}
+
+	return false
+}
+
+// RemoveBackupCode removes a used backup code from the list
+func RemoveBackupCode(usedCode string, storedCodes string) string {
+	if storedCodes == "" {
+		return ""
+	}
+
+	usedCode = strings.ToLower(strings.ReplaceAll(usedCode, " ", ""))
+	codes := strings.Split(storedCodes, ",")
+
+	var newCodes []string
+	for _, code := range codes {
+		if code != usedCode {
+			newCodes = append(newCodes, code)
+		}
+	}
+
+	return strings.Join(newCodes, ",")
+}
+
+// GenerateQRCodeURL generates a QR code URL for an existing secret
+func GenerateQRCodeURL(secret string, email string) (string, error) {
+	// Decode the base32 secret
+	secretBytes, err := base32.StdEncoding.DecodeString(secret)
+	if err != nil {
+		return "", fmt.Errorf("failed to decode secret: %v", err)
+	}
+
+	key, err := totp.Generate(totp.GenerateOpts{
+		Issuer:      IssuerName,
+		AccountName: email,
+		Secret:      secretBytes,
+	})
+	if err != nil {
+		return "", fmt.Errorf("failed to generate TOTP key: %v", err)
+	}
+
+	// Generate QR code image
+	var buf bytes.Buffer
+	img, err := key.Image(256, 256)
+	if err != nil {
+		return "", fmt.Errorf("failed to generate QR code image: %v", err)
+	}
+
+	// Encode image as PNG and convert to base64
+	err = png.Encode(&buf, img)
+	if err != nil {
+		return "", fmt.Errorf("failed to encode QR code image: %v", err)
+	}
+
+	// Create data URL
+	dataURL := fmt.Sprintf("data:image/png;base64,%s", base64.StdEncoding.EncodeToString(buf.Bytes()))
+
+	return dataURL, nil
+}

internal/db/db.go

@@ -25,6 +25,9 @@ type User struct {
 	AccountLocked       *bool `gorm:"default:false"`
 	LockoutUntil        *time.Time
 	Theme               string `gorm:"default:'light'"`
+	TwoFactorSecret     string `gorm:"type:varchar(32)"`
+	TwoFactorEnabled    bool   `gorm:"default:false"`
+	BackupCodes         string `gorm:"type:text"` // Comma-separated backup codes
 	CreatedAt           time.Time
 	UpdatedAt           time.Time
 }

internal/db/migrations/003_add_2fa.go

@@ -0,0 +1,80 @@
+package migrations
+
+import (
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/go-gormigrate/gormigrate/v2"
+	"gorm.io/gorm"
+)
+
+// Add2FA creates a migration for adding Two-Factor Authentication fields
+func Add2FA() *gormigrate.Migration {
+	return &gormigrate.Migration{
+		ID: "003_add_2fa",
+		Migrate: func(tx *gorm.DB) error {
+			// Check if any tables exist (indicating an existing database)
+			var count int64
+			if err := tx.Raw("SELECT count(*) FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'").Scan(&count).Error; err != nil {
+				return fmt.Errorf("failed to check for existing tables: %v", err)
+			}
+
+			// If tables exist, create a backup
+			if count > 0 {
+				// Get the database path
+				sqlDB, err := tx.DB()
+				if err != nil {
+					return fmt.Errorf("failed to get underlying database: %v", err)
+				}
+
+				var seq int
+				var name, dbPath string
+				if err := sqlDB.QueryRow("PRAGMA database_list").Scan(&seq, &name, &dbPath); err != nil {
+					return fmt.Errorf("failed to get database path: %v", err)
+				}
+
+				// Create backup file with timestamp
+				backupFile := fmt.Sprintf("%s.backup.%s", dbPath, time.Now().Format("20060102_150405"))
+
+				// Read original database
+				data, err := os.ReadFile(dbPath)
+				if err != nil {
+					return fmt.Errorf("failed to read database for backup: %v", err)
+				}
+
+				// Write backup
+				if err := os.WriteFile(backupFile, data, 0600); err != nil {
+					return fmt.Errorf("failed to create database backup: %v", err)
+				}
+
+				fmt.Printf("Created database backup at: %s\n", backupFile)
+			}
+
+			// Add new columns for 2FA - one at a time for SQLite compatibility
+			if err := tx.Exec(`ALTER TABLE users ADD COLUMN two_factor_secret VARCHAR(32)`).Error; err != nil {
+				return err
+			}
+			if err := tx.Exec(`ALTER TABLE users ADD COLUMN two_factor_enabled BOOLEAN DEFAULT FALSE`).Error; err != nil {
+				return err
+			}
+			if err := tx.Exec(`ALTER TABLE users ADD COLUMN backup_codes TEXT`).Error; err != nil {
+				return err
+			}
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			// Remove 2FA columns - one at a time for SQLite compatibility
+			if err := tx.Exec(`ALTER TABLE users DROP COLUMN two_factor_secret`).Error; err != nil {
+				return err
+			}
+			if err := tx.Exec(`ALTER TABLE users DROP COLUMN two_factor_enabled`).Error; err != nil {
+				return err
+			}
+			if err := tx.Exec(`ALTER TABLE users DROP COLUMN backup_codes`).Error; err != nil {
+				return err
+			}
+			return nil
+		},
+	}
+}

internal/db/migrations/migrations.go

@@ -10,6 +10,7 @@ func InitMigrations(db *gorm.DB) *gormigrate.Gormigrate {
 	migrations := []*gormigrate.Migration{
 		InitialSchema(),
 		UpdateGDriveType(),
+		Add2FA(),
 	}
 
 	return gormigrate.New(db, gormigrate.DefaultOptions, migrations)

internal/web/handlers/auth_handlers.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/base64"
+	"fmt"
 	"log"
 	"net/http"
 	"strings"
@@ -56,11 +57,21 @@ func (h *Handlers) AuthMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// Safely extract claims with type assertions and defaults
+		userID, _ := claims["user_id"].(float64)
+		email, _ := claims["email"].(string)
+		username, _ := claims["username"].(string)
+		isAdmin, _ := claims["is_admin"].(bool)
+
 		// Set user information in the context
-		c.Set("userID", uint(claims["user_id"].(float64)))
-		c.Set("email", claims["email"].(string))
-		c.Set("username", claims["username"].(string))
-		c.Set("isAdmin", claims["is_admin"].(bool))
+		c.Set("userID", uint(userID))
+		if email != "" {
+			c.Set("email", email)
+		}
+		if username != "" {
+			c.Set("username", username)
+		}
+		c.Set("isAdmin", isAdmin)
 
 		c.Next()
 	}
@@ -142,10 +153,11 @@ func (h *Handlers) APIAdminMiddleware() gin.HandlerFunc {
 }
 
 // GenerateJWT generates a JWT token for the given user
-func (h *Handlers) GenerateJWT(userID uint, username string, isAdmin bool) (string, error) {
+func (h *Handlers) GenerateJWT(userID uint, email string, isAdmin bool) (string, error) {
 	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
 		"user_id":  userID,
-		"username": username,
+		"email":    email,
+		"username": strings.Split(email, "@")[0], // Use email prefix as username
 		"is_admin": isAdmin,
 		"exp":      time.Now().Add(time.Hour * 24).Unix(),
 	})
@@ -243,24 +255,30 @@ func (h *Handlers) HandleLogin(c *gin.Context) {
 		return
 	}
 
-	// Generate JWT token with all necessary user information
-	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
-		"user_id":  user.ID,
-		"email":    user.Email,
-		"username": strings.Split(user.Email, "@")[0], // Use email prefix as username
-		"is_admin": user.IsAdmin,
-		"exp":      time.Now().Add(time.Hour * 24).Unix(),
-	})
+	// Check if 2FA is enabled
+	if user.TwoFactorEnabled {
+		// Store user ID temporarily for 2FA verification
+		c.SetCookie("temp_user_id", fmt.Sprintf("%d", user.ID), 300, "/", "", false, true) // 5 minutes expiry
 
-	// Sign the token
-	tokenString, err := token.SignedString([]byte(h.JWTSecret))
+		// Redirect to 2FA verification page
+		c.Redirect(http.StatusFound, "/login/verify")
+		return
+	}
+
+	// If 2FA is not enabled, proceed with normal login
+	// Generate JWT token with all necessary user information
+	isAdmin := false
+	if user.IsAdmin != nil {
+		isAdmin = *user.IsAdmin
+	}
+	token, err := h.GenerateJWT(user.ID, user.Email, isAdmin)
 	if err != nil {
 		components.Login(components.CreateTemplateContext(c), "Authentication error").Render(c, c.Writer)
 		return
 	}
 
 	// Set token in cookie
-	c.SetCookie("jwt_token", tokenString, 86400, "/", "", false, true)
+	c.SetCookie("jwt_token", token, 86400, "/", "", false, true)
 	c.Redirect(http.StatusFound, "/dashboard")
 }
 

internal/web/handlers/routes.go

@@ -10,6 +10,8 @@ func (h *Handlers) RegisterRoutes(router *gin.Engine) {
 	router.GET("/", h.HandleHome)
 	router.GET("/login", h.HandleLoginPage)
 	router.POST("/login", h.HandleLogin)
+	router.GET("/login/verify", h.Handle2FAVerifyPage)
+	router.POST("/login/verify", h.Handle2FAVerify)
 	router.GET("/forgot-password", h.HandleForgotPasswordPage)
 	router.POST("/forgot-password", h.HandleForgotPassword)
 	router.GET("/reset-password", h.HandleResetPasswordPage)
@@ -22,6 +24,11 @@ func (h *Handlers) RegisterRoutes(router *gin.Engine) {
 	// Password change route - only accessed from profile page
 	authorized.POST("/change-password", h.HandleChangePassword)
 
+	// 2FA routes - under profile
+	authorized.GET("/profile/2fa/setup", h.Handle2FASetup)
+	authorized.POST("/profile/2fa/verify", h.Handle2FAVerifySetup)
+	authorized.POST("/profile/2fa/disable", h.Handle2FADisable)
+
 	{
 		authorized.GET("/dashboard", h.HandleDashboard)
 		authorized.GET("/configs", h.HandleConfigs)

internal/web/handlers/two_factor_handlers.go

@@ -0,0 +1,292 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+	"github.com/starfleetcptn/gomft/components"
+	"github.com/starfleetcptn/gomft/internal/auth"
+	"golang.org/x/crypto/bcrypt"
+)
+
+// Handle2FASetup handles the GET /profile/2fa/setup route
+func (h *Handlers) Handle2FASetup(c *gin.Context) {
+	// Get user from context
+	userID := c.GetUint("userID")
+
+	var user struct {
+		Email            string
+		TwoFactorEnabled bool
+	}
+
+	if err := h.DB.Table("users").Select("email, two_factor_enabled").Where("id = ?", userID).First(&user).Error; err != nil {
+		c.String(http.StatusInternalServerError, "Failed to get user")
+		return
+	}
+
+	// Check if 2FA is already enabled
+	if user.TwoFactorEnabled {
+		c.Redirect(http.StatusFound, "/profile")
+		return
+	}
+
+	// Generate TOTP secret and QR code URL
+	secret, qrCodeURL, err := auth.GenerateTOTPSecret(user.Email)
+	if err != nil {
+		c.String(http.StatusInternalServerError, "Failed to generate 2FA secret")
+		return
+	}
+
+	// Generate backup codes
+	backupCodes, err := auth.GenerateBackupCodes()
+	if err != nil {
+		c.String(http.StatusInternalServerError, "Failed to generate backup codes")
+		return
+	}
+
+	// Store secret and backup codes in session temporarily
+	c.SetCookie("2fa_setup_secret", secret, 3600, "/", "", false, true)
+	c.SetCookie("2fa_setup_backup_codes", strings.Join(backupCodes, ","), 3600, "/", "", false, true)
+
+	// Render setup page
+	data := components.TwoFactorSetupData{
+		QRCodeURL:    qrCodeURL,
+		Secret:       secret,
+		BackupCodes:  backupCodes,
+		ErrorMessage: "",
+	}
+	components.TwoFactorSetup(c.Request.Context(), data).Render(c, c.Writer)
+}
+
+// Handle2FAVerifySetup handles the POST /profile/2fa/verify route
+func (h *Handlers) Handle2FAVerifySetup(c *gin.Context) {
+	// Get user from context
+	userID := c.GetUint("userID")
+
+	var user struct {
+		Email string
+	}
+	if err := h.DB.Table("users").Select("email").Where("id = ?", userID).First(&user).Error; err != nil {
+		c.String(http.StatusInternalServerError, "Failed to get user")
+		return
+	}
+
+	// Get secret from session
+	secret, err := c.Cookie("2fa_setup_secret")
+	if err != nil {
+		c.String(http.StatusBadRequest, "Setup session expired")
+		return
+	}
+
+	// Get backup codes from session
+	backupCodes, err := c.Cookie("2fa_setup_backup_codes")
+	if err != nil {
+		c.String(http.StatusBadRequest, "Setup session expired")
+		return
+	}
+
+	// Verify the code
+	code := c.PostForm("code")
+	if !auth.ValidateTOTPCode(secret, code) {
+		// Regenerate QR code URL using the existing secret
+		qrCodeURL, err := auth.GenerateQRCodeURL(secret, user.Email)
+		if err != nil {
+			c.String(http.StatusInternalServerError, "Failed to generate QR code")
+			return
+		}
+
+		data := components.TwoFactorSetupData{
+			QRCodeURL:    qrCodeURL,
+			Secret:       secret,
+			BackupCodes:  strings.Split(backupCodes, ","),
+			ErrorMessage: "Invalid verification code. Please try again.",
+		}
+		components.TwoFactorSetup(c.Request.Context(), data).Render(c, c.Writer)
+		return
+	}
+
+	// Update user with 2FA settings
+	if err := h.DB.Table("users").Where("id = ?", userID).Updates(map[string]interface{}{
+		"two_factor_secret":  secret,
+		"two_factor_enabled": true,
+		"backup_codes":       backupCodes,
+	}).Error; err != nil {
+		c.String(http.StatusInternalServerError, "Failed to enable 2FA")
+		return
+	}
+
+	// Clear setup cookies
+	c.SetCookie("2fa_setup_secret", "", -1, "/", "", false, true)
+	c.SetCookie("2fa_setup_backup_codes", "", -1, "/", "", false, true)
+
+	// Redirect to profile with success message
+	c.Redirect(http.StatusFound, "/profile?message=2FA+enabled+successfully")
+}
+
+// Handle2FAVerifyPage handles the GET /login/verify route
+func (h *Handlers) Handle2FAVerifyPage(c *gin.Context) {
+	// Check if we have a temporary user ID
+	_, err := c.Cookie("temp_user_id")
+	if err != nil {
+		c.Redirect(http.StatusFound, "/login")
+		return
+	}
+
+	// Render verification page
+	data := components.TwoFactorVerifyData{
+		ErrorMessage: "",
+	}
+	components.TwoFactorVerify(c.Request.Context(), data).Render(c, c.Writer)
+}
+
+// Handle2FAVerify handles the POST /login/verify route
+func (h *Handlers) Handle2FAVerify(c *gin.Context) {
+	// Get user ID from cookie
+	tempUserID, err := c.Cookie("temp_user_id")
+	if err != nil {
+		c.Redirect(http.StatusFound, "/login")
+		return
+	}
+
+	// Parse user ID
+	var userID uint
+	if _, err := fmt.Sscanf(tempUserID, "%d", &userID); err != nil {
+		c.Redirect(http.StatusFound, "/login")
+		return
+	}
+
+	var user struct {
+		TwoFactorSecret string
+		BackupCodes     string
+		Email           string
+		IsAdmin         *bool
+	}
+	if err := h.DB.Table("users").Select("two_factor_secret, backup_codes, email, is_admin").Where("id = ?", userID).First(&user).Error; err != nil {
+		c.Redirect(http.StatusFound, "/login")
+		return
+	}
+
+	code := c.PostForm("code")
+
+	// First try TOTP code
+	if auth.ValidateTOTPCode(user.TwoFactorSecret, code) {
+		// Generate new JWT token and set cookie
+		isAdmin := false
+		if user.IsAdmin != nil {
+			isAdmin = *user.IsAdmin
+		}
+		token, err := h.GenerateJWT(userID, user.Email, isAdmin)
+		if err != nil {
+			c.String(http.StatusInternalServerError, "Failed to generate token")
+			return
+		}
+		c.SetCookie("jwt_token", token, 86400, "/", "", false, true)
+
+		// Clear temporary user ID cookie
+		c.SetCookie("temp_user_id", "", -1, "/", "", false, true)
+
+		c.Redirect(http.StatusFound, "/dashboard")
+		return
+	}
+
+	// Then try backup code
+	if auth.ValidateBackupCode(code, user.BackupCodes) {
+		// Remove used backup code
+		newBackupCodes := auth.RemoveBackupCode(code, user.BackupCodes)
+		if err := h.DB.Model("users").Where("id = ?", userID).Update("backup_codes", newBackupCodes).Error; err != nil {
+			c.String(http.StatusInternalServerError, "Failed to update backup codes")
+			return
+		}
+
+		// Generate new JWT token and set cookie
+		isAdmin := false
+		if user.IsAdmin != nil {
+			isAdmin = *user.IsAdmin
+		}
+		token, err := h.GenerateJWT(userID, user.Email, isAdmin)
+		if err != nil {
+			c.String(http.StatusInternalServerError, "Failed to generate token")
+			return
+		}
+		c.SetCookie("jwt_token", token, 86400, "/", "", false, true)
+
+		// Clear temporary user ID cookie
+		c.SetCookie("temp_user_id", "", -1, "/", "", false, true)
+
+		c.Redirect(http.StatusFound, "/dashboard")
+		return
+	}
+
+	// If neither code is valid, show error
+	data := components.TwoFactorVerifyData{
+		ErrorMessage: "Invalid verification code. Please try again.",
+	}
+	components.TwoFactorVerify(c.Request.Context(), data).Render(c, c.Writer)
+}
+
+// Handle2FADisable handles the POST /profile/2fa/disable route
+func (h *Handlers) Handle2FADisable(c *gin.Context) {
+	// Get user ID from context
+	userID := c.GetUint("userID")
+
+	// Get current password from form
+	currentPassword := c.PostForm("current_password")
+	if currentPassword == "" {
+		c.Data(http.StatusBadRequest, "text/html", []byte(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Current password is required</span>
+		</div>`))
+		return
+	}
+
+	// Get user from database
+	var user struct {
+		PasswordHash     string
+		TwoFactorEnabled bool
+	}
+	if err := h.DB.Table("users").Select("password_hash, two_factor_enabled").Where("id = ?", userID).First(&user).Error; err != nil {
+		c.Data(http.StatusInternalServerError, "text/html", []byte(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Failed to get user information</span>
+		</div>`))
+		return
+	}
+
+	// Verify current password
+	if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(currentPassword)); err != nil {
+		c.Data(http.StatusBadRequest, "text/html", []byte(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Current password is incorrect</span>
+		</div>`))
+		return
+	}
+
+	// Check if 2FA is already disabled
+	if !user.TwoFactorEnabled {
+		c.Data(http.StatusBadRequest, "text/html", []byte(`<div class="bg-yellow-100 border border-yellow-400 text-yellow-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Two-factor authentication is already disabled</span>
+		</div>`))
+		return
+	}
+
+	// Disable 2FA
+	if err := h.DB.Table("users").Where("id = ?", userID).Updates(map[string]interface{}{
+		"two_factor_enabled": false,
+		"two_factor_secret":  nil,
+		"backup_codes":       nil,
+	}).Error; err != nil {
+		c.Data(http.StatusInternalServerError, "text/html", []byte(`<div class="bg-red-100 border border-red-400 text-red-700 px-4 py-3 rounded mb-4" role="alert">
+			<span class="block sm:inline">Failed to disable two-factor authentication</span>
+		</div>`))
+		return
+	}
+
+	// Return success message
+	c.Data(http.StatusOK, "text/html", []byte(`<div class="bg-green-100 border border-green-400 text-green-700 px-4 py-3 rounded mb-4" role="alert">
+		<span class="block sm:inline">Two-factor authentication has been disabled</span>
+		<script>
+			setTimeout(function() {
+				window.location.reload();
+			}, 1500);
+		</script>
+	</div>`))
+}