admin_auth.py

#!/usr/bin/env python3
"""
Admin Authentication Middleware for UnForkRAG

This module provides Flask middleware for securing admin endpoints with:
- Session-based authentication
- Login/logout functionality
- CSRF protection
- Rate limiting for login attempts
- Admin activity logging
"""

import os
import json
import time
import secrets
from functools import wraps
from datetime import datetime, timedelta
from typing import Dict, Optional, List, Tuple, Any
from flask import Flask, request, session, jsonify, redirect, url_for, flash

# Import the log_api decorator from the main server
try:
    from unfork_server import log_api
except ImportError:
    # Fallback if not available
    def log_api(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper

from admin_config import admin_config


class AdminAuth:
    """Admin authentication middleware"""
    
    def __init__(self, app: Flask):
        self.app = app
        self.sessions = {}  # session_id -> user_info
        self.failed_attempts = {}  # ip -> {count, last_attempt}
        
        # Register routes
        self._register_routes()
    
    def _register_routes(self):
        """Register authentication routes"""
        @self.app.route('/admin/login', methods=['GET', 'POST'])
        @log_api
        def admin_login():
            if request.method == 'POST':
                return self._handle_login()
            return self._render_login_page()
        
        @self.app.route('/admin/logout', methods=['GET'])
        @log_api
        def admin_logout():
            return self._handle_logout()
        
        @self.app.route('/admin/change-password', methods=['GET', 'POST'])
        @log_api
        def admin_change_password():
            if request.method == 'POST':
                return self._handle_change_password()
            return self._render_change_password_page()
    
    def _handle_login(self):
        """Handle login POST request"""
        data = request.get_json() or request.form
        username = data.get('username', '')
        password = data.get('password', '')
        remember = data.get('remember', False)
        
        # Check rate limiting for this IP
        if self._is_rate_limited():
            return jsonify({'error': 'Too many login attempts. Please try again later.'}), 429
        
        # Authenticate user
        if not username or not password:
            return jsonify({'error': 'Username and password required'}), 400
        
        user = admin_config.authenticate_user(username, password)
        if user:
            # Create session
            session_id = secrets.token_hex(32)
            expires_at = datetime.now() + timedelta(hours=24 if remember else 1)
            
            self.sessions[session_id] = {
                'username': username,
                'user_info': user,
                'created_at': datetime.now().isoformat(),
                'expires_at': expires_at.isoformat(),
                'ip': request.remote_addr
            }
            
            # Clear failed attempts for this IP
            self._clear_failed_attempts()
            
            # Log successful login
            self._log_admin_access('login', username, True)
            
            response = jsonify({'success': True, 'redirect': '/admin'})
            response.set_cookie('admin_session', session_id, 
                              max_age=86400 if remember else 3600,
                              httponly=True, secure=False)  # Set secure=True in production
            return response
        
        # Failed login
        self._record_failed_attempt()
        self._log_admin_access('login', username or '', False)
        
        return jsonify({'error': 'Invalid credentials'}), 401
    
    def _handle_logout(self):
        """Handle logout"""
        session_id = request.cookies.get('admin_session')
        if session_id and session_id in self.sessions:
            username = self.sessions[session_id]['username']
            del self.sessions[session_id]
            self._log_admin_access('logout', username, True)
        
        response = redirect('/admin/login')
        response.set_cookie('admin_session', '', expires=0)
        return response
    
    def _handle_change_password(self):
        """Handle password change"""
        session_id = request.cookies.get('admin_session')
        if not session_id or session_id not in self.sessions:
            return jsonify({'error': 'Not authenticated'}), 401
        
        data = request.get_json() or request.form
        current_password = data.get('current_password', '')
        new_password = data.get('new_password', '')
        username = self.sessions[session_id]['username']
        
        # Verify current password
        user_info = admin_config.get_user(username)
        if not user_info or not admin_config.verify_password(current_password, 
                                          user_info.get('password_hash', '')):
            return jsonify({'error': 'Current password is incorrect'}), 400
        
        # Change password
        if admin_config.change_password(username, current_password, new_password):
            self._log_admin_access('password_change', username, True)
            return jsonify({'success': True})
        
        return jsonify({'error': 'Failed to change password'}), 500
    
    def _is_rate_limited(self) -> bool:
        """Check if current IP is rate limited"""
        ip = request.remote_addr
        now = time.time()
        config = admin_config.get_security_config()
        max_attempts = config.get('max_login_attempts', 5)
        lockout_duration = config.get('lockout_duration', 900)
        
        if ip not in self.failed_attempts:
            return False
        
        attempts = self.failed_attempts[ip]
        
        # Check if lockout period has expired
        if now - attempts['last_attempt'] > lockout_duration:
            del self.failed_attempts[ip]
            return False
        
        return attempts['count'] >= max_attempts
    
    def _record_failed_attempt(self):
        """Record a failed login attempt"""
        ip = request.remote_addr
        now = time.time()
        
        if ip not in self.failed_attempts:
            self.failed_attempts[ip] = {'count': 0, 'last_attempt': 0}
        
        self.failed_attempts[ip]['count'] += 1
        self.failed_attempts[ip]['last_attempt'] = now
    
    def _clear_failed_attempts(self):
        """Clear failed attempts for current IP"""
        ip = request.remote_addr
        if ip in self.failed_attempts:
            del self.failed_attempts[ip]
    
    def _is_session_valid(self, session_id: str) -> bool:
        """Check if session is valid and not expired"""
        if session_id not in self.sessions:
            return False
        
        session_data = self.sessions[session_id]
        expires_at = datetime.fromisoformat(session_data['expires_at'])
        
        if datetime.now() > expires_at:
            del self.sessions[session_id]
            return False
        
        return True
    
    def _get_current_user(self) -> Optional[Dict]:
        """Get current authenticated user"""
        session_id = request.cookies.get('admin_session')
        if not session_id or not self._is_session_valid(session_id):
            return None
        
        return self.sessions[session_id]['user_info']
    
    def _log_admin_access(self, action: str, username: str, success: bool):
        """Log admin access attempts"""
        if not admin_config.get_monitoring_config().get('log_admin_access', True):
            return
        
        log_entry = {
            'timestamp': datetime.now().isoformat(),
            'action': action,
            'username': username,
            'success': success,
            'ip': request.remote_addr,
            'user_agent': request.headers.get('User-Agent', '')
        }
        
        # Log to file
        log_file = 'admin_access.log'
        try:
            with open(log_file, 'a') as f:
                f.write(json.dumps(log_entry) + '\n')
        except Exception as e:
            print(f"Error logging admin access: {e}")
    
    def _render_login_page(self) -> str:
        """Render login page"""
        return '''
<!doctype html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>UnForkRAG Admin Login</title>
    <style>
        body { font-family: system-ui, -apple-system, Segoe UI, Roboto, Arial; margin: 0; padding: 0; background: #f3f4f6; }
        .container { max-width: 400px; margin: 100px auto; padding: 20px; background: white; border-radius: 8px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
        h1 { font-size: 24px; margin-bottom: 20px; text-align: center; color: #1f2937; }
        .form-group { margin-bottom: 15px; }
        label { display: block; margin-bottom: 5px; font-weight: 500; color: #374151; }
        input[type="text"], input[type="password"] { width: 100%; padding: 10px; border: 1px solid #d1d5db; border-radius: 4px; font-size: 16px; }
        .checkbox { display: flex; align-items: center; gap: 8px; margin-bottom: 20px; }
        button { width: 100%; padding: 12px; background: #2563eb; color: white; border: none; border-radius: 4px; font-size: 16px; cursor: pointer; }
        button:hover { background: #1d4ed8; }
        .error { color: #dc2626; font-size: 14px; margin-top: 10px; }
        .warning { color: #d97706; font-size: 12px; margin-top: 10px; }
    </style>
</head>
<body>
    <div class="container">
        <h1>Admin Login</h1>
        <form id="loginForm">
            <div class="form-group">
                <label for="username">Username</label>
                <input type="text" id="username" name="username" required>
            </div>
            <div class="form-group">
                <label for="password">Password</label>
                <input type="password" id="password" name="password" required>
            </div>
            <div class="checkbox">
                <input type="checkbox" id="remember" name="remember">
                <label for="remember">Remember me (24 hours)</label>
            </div>
            <button type="submit">Login</button>
        </form>
        <div id="message" class="error" style="display: none;"></div>
        <div class="warning">
            Default credentials: admin / admin123<br>
            Change password after first login!
        </div>
    </div>

    <script>
        document.getElementById('loginForm').addEventListener('submit', async (e) => {
            e.preventDefault();
            const formData = new FormData(e.target);
            const data = Object.fromEntries(formData);
            
            try {
                const response = await fetch('/admin/login', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify(data)
                });
                
                const result = await response.json();
                
                if (result.success) {
                    window.location.href = result.redirect;
                } else {
                    document.getElementById('message').style.display = 'block';
                    document.getElementById('message').textContent = result.error;
                }
            } catch (error) {
                document.getElementById('message').style.display = 'block';
                document.getElementById('message').textContent = 'Login failed';
            }
        });
    </script>
</body>
</html>
        '''
    
    def _render_change_password_page(self) -> str:
        """Render change password page"""
        return '''
<!doctype html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Change Password - UnForkRAG Admin</title>
    <style>
        body { font-family: system-ui, -apple-system, Segoe UI, Roboto, Arial; margin: 0; padding: 0; background: #f3f4f6; }
        .container { max-width: 400px; margin: 100px auto; padding: 20px; background: white; border-radius: 8px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
        h1 { font-size: 24px; margin-bottom: 20px; text-align: center; color: #1f2937; }
        .form-group { margin-bottom: 15px; }
        label { display: block; margin-bottom: 5px; font-weight: 500; color: #374151; }
        input[type="password"] { width: 100%; padding: 10px; border: 1px solid #d1d5db; border-radius: 4px; font-size: 16px; }
        button { width: 100%; padding: 12px; background: #2563eb; color: white; border: none; border-radius: 4px; font-size: 16px; cursor: pointer; }
        button:hover { background: #1d4ed8; }
        .message { margin-top: 10px; padding: 10px; border-radius: 4px; }
        .success { background: #dcfce7; color: #166534; }
        .error { background: #fee2e2; color: #991b1b; }
    </style>
</head>
<body>
    <div class="container">
        <h1>Change Password</h1>
        <form id="passwordForm">
            <div class="form-group">
                <label for="current_password">Current Password</label>
                <input type="password" id="current_password" name="current_password" required>
            </div>
            <div class="form-group">
                <label for="new_password">New Password</label>
                <input type="password" id="new_password" name="new_password" required>
            </div>
            <div class="form-group">
                <label for="confirm_password">Confirm New Password</label>
                <input type="password" id="confirm_password" name="confirm_password" required>
            </div>
            <button type="submit">Change Password</button>
        </form>
        <div id="message" class="message" style="display: none;"></div>
    </div>

    <script>
        document.getElementById('passwordForm').addEventListener('submit', async (e) => {
            e.preventDefault();
            
            const currentPassword = document.getElementById('current_password').value;
            const newPassword = document.getElementById('new_password').value;
            const confirmPassword = document.getElementById('confirm_password').value;
            
            if (newPassword !== confirmPassword) {
                showMessage('Passwords do not match', 'error');
                return;
            }
            
            try {
                const response = await fetch('/admin/change-password', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({
                        current_password: currentPassword,
                        new_password: newPassword
                    })
                });
                
                const result = await response.json();
                
                if (result.success) {
                    showMessage('Password changed successfully', 'success');
                    setTimeout(() => window.location.href = '/admin', 2000);
                } else {
                    showMessage(result.error || 'Failed to change password', 'error');
                }
            } catch (error) {
                showMessage('An error occurred', 'error');
            }
        });
        
        function showMessage(text, type) {
            const messageDiv = document.getElementById('message');
            messageDiv.textContent = text;
            messageDiv.className = `message ${type}`;
            messageDiv.style.display = 'block';
        }
    </script>
</body>
</html>
        '''
    
    def require_admin(self, f):
        """Decorator to require admin authentication"""
        @wraps(f)
        def decorated_function(*args, **kwargs):
            session_id = request.cookies.get('admin_session')
            if not session_id or not self._is_session_valid(session_id):
                return redirect('/admin/login')
            return f(*args, **kwargs)
        return decorated_function


def init_admin_auth(app: Flask) -> AdminAuth:
    """Initialize admin authentication for the Flask app"""
    return AdminAuth(app)