525 lines
18 KiB
Python
525 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
TOTP Manager for Media Downloader
|
|
Handles Time-based One-Time Password (TOTP) operations for 2FA
|
|
|
|
Based on backup-central's implementation
|
|
"""
|
|
|
|
import sys
|
|
import pyotp
|
|
import qrcode
|
|
import io
|
|
import base64
|
|
import sqlite3
|
|
import hashlib
|
|
import bcrypt
|
|
import secrets
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional, Dict, List, Tuple
|
|
from cryptography.fernet import Fernet
|
|
from pathlib import Path
|
|
|
|
# Add parent path to allow imports from modules
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
from modules.universal_logger import get_logger
|
|
|
|
logger = get_logger('TOTPManager')
|
|
|
|
|
|
class TOTPManager:
|
|
def __init__(self, db_path: str = None):
|
|
if db_path is None:
|
|
db_path = str(Path(__file__).parent.parent.parent / 'database' / 'auth.db')
|
|
|
|
self.db_path = db_path
|
|
self.issuer = 'Media Downloader'
|
|
self.window = 1 # Allow 1 step (30s) tolerance for time skew
|
|
|
|
# Encryption key for TOTP secrets
|
|
self.encryption_key = self._get_encryption_key()
|
|
self.cipher_suite = Fernet(self.encryption_key)
|
|
|
|
# Rate limiting configuration
|
|
self.max_attempts = 5
|
|
self.rate_limit_window = timedelta(minutes=15)
|
|
self.lockout_duration = timedelta(minutes=30)
|
|
|
|
# Initialize database tables
|
|
self._init_database()
|
|
|
|
def _get_encryption_key(self) -> bytes:
|
|
"""Get or generate encryption key for TOTP secrets"""
|
|
key_file = Path(__file__).parent.parent.parent / '.totp_encryption_key'
|
|
|
|
if key_file.exists():
|
|
with open(key_file, 'rb') as f:
|
|
return f.read()
|
|
|
|
# Generate new key
|
|
key = Fernet.generate_key()
|
|
try:
|
|
with open(key_file, 'wb') as f:
|
|
f.write(key)
|
|
import os
|
|
os.chmod(key_file, 0o600)
|
|
except Exception as e:
|
|
logger.warning(f"Could not save TOTP encryption key: {e}", module="TOTP")
|
|
|
|
return key
|
|
|
|
def _init_database(self):
|
|
"""Initialize TOTP-related database tables"""
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Ensure TOTP columns exist in users table
|
|
try:
|
|
cursor.execute("ALTER TABLE users ADD COLUMN totp_secret TEXT")
|
|
except sqlite3.OperationalError:
|
|
pass # Column already exists
|
|
|
|
try:
|
|
cursor.execute("ALTER TABLE users ADD COLUMN totp_enabled INTEGER NOT NULL DEFAULT 0")
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
|
|
try:
|
|
cursor.execute("ALTER TABLE users ADD COLUMN totp_enrolled_at TEXT")
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
|
|
# TOTP rate limiting table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS totp_rate_limit (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT NOT NULL,
|
|
ip_address TEXT,
|
|
attempts INTEGER NOT NULL DEFAULT 1,
|
|
window_start TEXT NOT NULL,
|
|
locked_until TEXT,
|
|
UNIQUE(username, ip_address)
|
|
)
|
|
""")
|
|
|
|
# TOTP audit log
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS totp_audit_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT NOT NULL,
|
|
action TEXT NOT NULL,
|
|
success INTEGER NOT NULL,
|
|
ip_address TEXT,
|
|
user_agent TEXT,
|
|
details TEXT,
|
|
timestamp TEXT NOT NULL
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
|
|
def generate_secret(self, username: str) -> Dict:
|
|
"""
|
|
Generate a new TOTP secret and QR code for a user
|
|
|
|
Returns:
|
|
dict with secret, qrCodeDataURL, manualEntryKey
|
|
"""
|
|
try:
|
|
# Generate secret
|
|
secret = pyotp.random_base32()
|
|
|
|
# Create provisioning URI
|
|
totp = pyotp.TOTP(secret)
|
|
provisioning_uri = totp.provisioning_uri(
|
|
name=username,
|
|
issuer_name=self.issuer
|
|
)
|
|
|
|
# Generate QR code
|
|
qr = qrcode.QRCode(version=1, box_size=10, border=4)
|
|
qr.add_data(provisioning_uri)
|
|
qr.make(fit=True)
|
|
|
|
img = qr.make_image(fill_color="black", back_color="white")
|
|
|
|
# Convert to base64 data URL
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
img_str = base64.b64encode(buffer.getvalue()).decode()
|
|
qr_code_data_url = f"data:image/png;base64,{img_str}"
|
|
|
|
self._log_audit(username, 'totp_secret_generated', True, None, None,
|
|
'TOTP secret generated for user')
|
|
|
|
return {
|
|
'secret': secret,
|
|
'qrCodeDataURL': qr_code_data_url,
|
|
'manualEntryKey': secret,
|
|
'otpauthURL': provisioning_uri
|
|
}
|
|
except Exception as error:
|
|
self._log_audit(username, 'totp_secret_generation_failed', False, None, None,
|
|
f'Error: {str(error)}')
|
|
raise
|
|
|
|
def verify_token(self, secret: str, token: str) -> bool:
|
|
"""
|
|
Verify a TOTP token
|
|
|
|
Args:
|
|
secret: Base32 encoded secret
|
|
token: 6-digit TOTP code
|
|
|
|
Returns:
|
|
True if valid, False otherwise
|
|
"""
|
|
try:
|
|
totp = pyotp.TOTP(secret)
|
|
return totp.verify(token, valid_window=self.window)
|
|
except Exception as e:
|
|
logger.error(f"TOTP verification error: {e}", module="TOTP")
|
|
return False
|
|
|
|
def encrypt_secret(self, secret: str) -> str:
|
|
"""Encrypt TOTP secret for storage"""
|
|
encrypted = self.cipher_suite.encrypt(secret.encode())
|
|
return base64.b64encode(encrypted).decode()
|
|
|
|
def decrypt_secret(self, encrypted_secret: str) -> str:
|
|
"""Decrypt TOTP secret from storage"""
|
|
encrypted = base64.b64decode(encrypted_secret.encode())
|
|
decrypted = self.cipher_suite.decrypt(encrypted)
|
|
return decrypted.decode()
|
|
|
|
def generate_backup_codes(self, count: int = 10) -> List[str]:
|
|
"""
|
|
Generate backup codes for recovery
|
|
|
|
Args:
|
|
count: Number of codes to generate
|
|
|
|
Returns:
|
|
List of formatted backup codes
|
|
"""
|
|
codes = []
|
|
for _ in range(count):
|
|
# Generate 8-character hex code
|
|
code = secrets.token_hex(4).upper()
|
|
# Format as XXXX-XXXX
|
|
formatted = f"{code[:4]}-{code[4:8]}"
|
|
codes.append(formatted)
|
|
return codes
|
|
|
|
def hash_backup_code(self, code: str) -> str:
|
|
"""
|
|
Hash a backup code for storage using bcrypt
|
|
|
|
Args:
|
|
code: Backup code to hash
|
|
|
|
Returns:
|
|
Bcrypt hash of the code
|
|
"""
|
|
# Use bcrypt with default work factor (12 rounds)
|
|
return bcrypt.hashpw(code.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
|
|
|
|
def verify_backup_code(self, code: str, username: str) -> Tuple[bool, int]:
|
|
"""
|
|
Verify a backup code and mark it as used
|
|
|
|
Returns:
|
|
Tuple of (valid, remaining_codes)
|
|
"""
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Get unused backup codes for user
|
|
cursor.execute("""
|
|
SELECT id, code_hash FROM backup_codes
|
|
WHERE username = ? AND used = 0
|
|
""", (username,))
|
|
|
|
rows = cursor.fetchall()
|
|
|
|
# Check if code matches any unused code
|
|
for row_id, stored_hash in rows:
|
|
# Support both bcrypt (new) and SHA256 (legacy) hashes
|
|
is_match = False
|
|
|
|
if stored_hash.startswith('$2b$'):
|
|
# Bcrypt hash
|
|
try:
|
|
is_match = bcrypt.checkpw(code.encode('utf-8'), stored_hash.encode('utf-8'))
|
|
except Exception as e:
|
|
logger.error(f"Error verifying bcrypt backup code: {e}", module="TOTP")
|
|
continue
|
|
else:
|
|
# Legacy SHA256 hash
|
|
legacy_hash = hashlib.sha256(code.encode()).hexdigest()
|
|
is_match = (stored_hash == legacy_hash)
|
|
|
|
if is_match:
|
|
# Mark as used
|
|
cursor.execute("""
|
|
UPDATE backup_codes
|
|
SET used = 1, used_at = ?
|
|
WHERE id = ?
|
|
""", (datetime.now().isoformat(), row_id))
|
|
|
|
conn.commit()
|
|
|
|
# Count remaining codes
|
|
cursor.execute("""
|
|
SELECT COUNT(*) FROM backup_codes
|
|
WHERE username = ? AND used = 0
|
|
""", (username,))
|
|
|
|
remaining = cursor.fetchone()[0]
|
|
return True, remaining
|
|
|
|
return False, len(rows)
|
|
|
|
def enable_totp(self, username: str, secret: str, backup_codes: List[str]) -> bool:
|
|
"""
|
|
Enable TOTP for a user
|
|
|
|
Args:
|
|
username: Username
|
|
secret: TOTP secret
|
|
backup_codes: List of backup codes
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
try:
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Encrypt secret
|
|
encrypted_secret = self.encrypt_secret(secret)
|
|
|
|
# Update user
|
|
cursor.execute("""
|
|
UPDATE users
|
|
SET totp_secret = ?, totp_enabled = 1, totp_enrolled_at = ?
|
|
WHERE username = ?
|
|
""", (encrypted_secret, datetime.now().isoformat(), username))
|
|
|
|
# Delete old backup codes
|
|
cursor.execute("DELETE FROM backup_codes WHERE username = ?", (username,))
|
|
|
|
# Insert new backup codes
|
|
for code in backup_codes:
|
|
code_hash = self.hash_backup_code(code)
|
|
cursor.execute("""
|
|
INSERT INTO backup_codes (username, code_hash, created_at)
|
|
VALUES (?, ?, ?)
|
|
""", (username, code_hash, datetime.now().isoformat()))
|
|
|
|
conn.commit()
|
|
|
|
self._log_audit(username, 'totp_enabled', True, None, None,
|
|
f'TOTP enabled with {len(backup_codes)} backup codes')
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error enabling TOTP: {e}", module="TOTP")
|
|
return False
|
|
|
|
def disable_totp(self, username: str) -> bool:
|
|
"""
|
|
Disable TOTP for a user
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
try:
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
UPDATE users
|
|
SET totp_secret = NULL, totp_enabled = 0, totp_enrolled_at = NULL
|
|
WHERE username = ?
|
|
""", (username,))
|
|
|
|
# Delete backup codes
|
|
cursor.execute("DELETE FROM backup_codes WHERE username = ?", (username,))
|
|
|
|
conn.commit()
|
|
|
|
self._log_audit(username, 'totp_disabled', True, None, None, 'TOTP disabled')
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error disabling TOTP: {e}", module="TOTP")
|
|
return False
|
|
|
|
def get_totp_status(self, username: str) -> Dict:
|
|
"""
|
|
Get TOTP status for a user
|
|
|
|
Returns:
|
|
dict with enabled, enrolledAt
|
|
"""
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT totp_enabled, totp_enrolled_at FROM users
|
|
WHERE username = ?
|
|
""", (username,))
|
|
|
|
row = cursor.fetchone()
|
|
|
|
if not row:
|
|
return {'enabled': False, 'enrolledAt': None}
|
|
|
|
return {
|
|
'enabled': bool(row[0]),
|
|
'enrolledAt': row[1]
|
|
}
|
|
|
|
def check_rate_limit(self, username: str, ip_address: str) -> Tuple[bool, int, Optional[str]]:
|
|
"""
|
|
Check and enforce rate limiting for TOTP verification
|
|
|
|
Returns:
|
|
Tuple of (allowed, attempts_remaining, locked_until)
|
|
"""
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
now = datetime.now()
|
|
window_start = now - self.rate_limit_window
|
|
|
|
cursor.execute("""
|
|
SELECT attempts, window_start, locked_until
|
|
FROM totp_rate_limit
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (username, ip_address))
|
|
|
|
row = cursor.fetchone()
|
|
|
|
if not row:
|
|
# First attempt - create record
|
|
cursor.execute("""
|
|
INSERT INTO totp_rate_limit (username, ip_address, attempts, window_start)
|
|
VALUES (?, ?, 1, ?)
|
|
""", (username, ip_address, now.isoformat()))
|
|
conn.commit()
|
|
return True, self.max_attempts - 1, None
|
|
|
|
attempts, stored_window_start, locked_until = row
|
|
stored_window_start = datetime.fromisoformat(stored_window_start)
|
|
|
|
# Check if locked out
|
|
if locked_until:
|
|
locked_until_dt = datetime.fromisoformat(locked_until)
|
|
if locked_until_dt > now:
|
|
return False, 0, locked_until
|
|
else:
|
|
# Lockout expired - reset
|
|
cursor.execute("""
|
|
DELETE FROM totp_rate_limit
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (username, ip_address))
|
|
conn.commit()
|
|
return True, self.max_attempts, None
|
|
|
|
# Check if window expired
|
|
if stored_window_start < window_start:
|
|
# Reset window
|
|
cursor.execute("""
|
|
UPDATE totp_rate_limit
|
|
SET attempts = 1, window_start = ?, locked_until = NULL
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (now.isoformat(), username, ip_address))
|
|
conn.commit()
|
|
return True, self.max_attempts - 1, None
|
|
|
|
# Increment attempts
|
|
new_attempts = attempts + 1
|
|
|
|
if new_attempts >= self.max_attempts:
|
|
# Lock out
|
|
locked_until = (now + self.lockout_duration).isoformat()
|
|
cursor.execute("""
|
|
UPDATE totp_rate_limit
|
|
SET attempts = ?, locked_until = ?
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (new_attempts, locked_until, username, ip_address))
|
|
conn.commit()
|
|
return False, 0, locked_until
|
|
else:
|
|
cursor.execute("""
|
|
UPDATE totp_rate_limit
|
|
SET attempts = ?
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (new_attempts, username, ip_address))
|
|
conn.commit()
|
|
return True, self.max_attempts - new_attempts, None
|
|
|
|
def reset_rate_limit(self, username: str, ip_address: str):
|
|
"""Reset rate limit after successful authentication"""
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
DELETE FROM totp_rate_limit
|
|
WHERE username = ? AND ip_address = ?
|
|
""", (username, ip_address))
|
|
|
|
conn.commit()
|
|
|
|
def regenerate_backup_codes(self, username: str) -> List[str]:
|
|
"""
|
|
Regenerate backup codes for a user
|
|
|
|
Returns:
|
|
List of new backup codes
|
|
"""
|
|
# Generate new codes
|
|
new_codes = self.generate_backup_codes(10)
|
|
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Delete old codes
|
|
cursor.execute("DELETE FROM backup_codes WHERE username = ?", (username,))
|
|
|
|
# Insert new codes
|
|
for code in new_codes:
|
|
code_hash = self.hash_backup_code(code)
|
|
cursor.execute("""
|
|
INSERT INTO backup_codes (username, code_hash, created_at)
|
|
VALUES (?, ?, ?)
|
|
""", (username, code_hash, datetime.now().isoformat()))
|
|
|
|
conn.commit()
|
|
|
|
self._log_audit(username, 'backup_codes_regenerated', True, None, None,
|
|
f'Generated {len(new_codes)} new backup codes')
|
|
|
|
return new_codes
|
|
|
|
def _log_audit(self, username: str, action: str, success: bool,
|
|
ip_address: Optional[str], user_agent: Optional[str],
|
|
details: Optional[str] = None):
|
|
"""Log TOTP audit event"""
|
|
try:
|
|
with sqlite3.connect(self.db_path) as conn:
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
INSERT INTO totp_audit_log
|
|
(username, action, success, ip_address, user_agent, details, timestamp)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""", (username, action, int(success), ip_address, user_agent, details,
|
|
datetime.now().isoformat()))
|
|
|
|
conn.commit()
|
|
except Exception as e:
|
|
logger.error(f"Error logging TOTP audit: {e}", module="TOTP")
|