utils/auth.py

#!/usr/bin/env python3
"""
auth.py

Authentication utilities for MorphGuard using SQLite database.
"""
import os
import json
import time
from typing import Dict, Any, List, Optional, Tuple
from werkzeug.security import generate_password_hash, check_password_hash
from datetime import datetime

# Try to import psycopg2
try:
    import psycopg2
    from psycopg2.extras import RealDictCursor
    HAS_POSTGRES = True
except ImportError:
    HAS_POSTGRES = False
    RealDictCursor = None

import sqlite3

IS_HF_SPACE = os.environ.get('HF_SPACE', '0') == '1'
USE_SQLITE = IS_HF_SPACE or not HAS_POSTGRES

# Database configuration
DB_HOST = os.environ.get('MORPHGUARD_DB_HOST', 'localhost')
DB_PORT = int(os.environ.get('MORPHGUARD_DB_PORT', 5432))
DB_NAME = os.environ.get('MORPHGUARD_DB_NAME', 'morphguard')
DB_USER = os.environ.get('MORPHGUARD_DB_USER', 'morphguard')
DB_PASS = os.environ.get('MORPHGUARD_DB_PASS', 'morphguard')

def _is_sqlite(conn):
    return isinstance(conn, sqlite3.Connection)

def _get_conn():
    if USE_SQLITE:
        try:
            conn = sqlite3.connect('users.db')
            conn.row_factory = sqlite3.Row
            return conn
        except Exception as e:
            print(f"SQLite connection failed: {e}")
            return None
    else:
        try:
            conn = psycopg2.connect(
                host=DB_HOST,
                port=DB_PORT,
                dbname=DB_NAME,
                user=DB_USER,
                password=DB_PASS
            )
            return conn
        except Exception as e:
            print(f"Database connection failed: {e}")
            # Try falling back to SQLite if Postgres connection fails
            try:
                conn = sqlite3.connect('users.db')
                conn.row_factory = sqlite3.Row
                return conn
            except Exception:
                return None

def _execute(conn, query: str, params: tuple = (), fetch_all: bool = False, fetch_one: bool = False, commit: bool = False):
    is_sqlite = _is_sqlite(conn)
    if is_sqlite:
        query = query.replace('%s', '?')
        new_params = []
        for p in params:
            if isinstance(p, datetime):
                new_params.append(p.timestamp())
            elif isinstance(p, bool):
                new_params.append(1 if p else 0)
            else:
                new_params.append(p)
        params = tuple(new_params)
        cur = conn.cursor()
    else:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        
    cur.execute(query, params)
    
    result = None
    if fetch_all:
        rows = cur.fetchall()
        result = [dict(r) for r in rows]
    elif fetch_one:
        row = cur.fetchone()
        result = dict(row) if row else None
        
    if commit:
        conn.commit()
        
    return cur, result

def init_db():
    conn = _get_conn()
    if conn:
        db_type = "SQLite" if _is_sqlite(conn) else "PostgreSQL"
        if _is_sqlite(conn):
            try:
                cur = conn.cursor()
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS users (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        username TEXT UNIQUE NOT NULL,
                        password_hash TEXT NOT NULL,
                        role TEXT NOT NULL,
                        email TEXT,
                        created_at REAL,
                        last_login REAL,
                        active INTEGER NOT NULL DEFAULT 1
                    );
                """)
                conn.commit()
                
                cur.execute("SELECT id FROM users WHERE username = ?;", ("admin",))
                if not cur.fetchone():
                    pwd_hash = generate_password_hash("morphguard_admin", method='scrypt')
                    cur.execute(
                        "INSERT INTO users (username, password_hash, role, created_at) VALUES (?, ?, ?, ?);",
                        ("admin", pwd_hash, "admin", datetime.now().timestamp())
                    )
                    conn.commit()
                    print("Default admin user created in SQLite database.")
            except Exception as e:
                print(f"Error initializing SQLite database: {e}")
        conn.close()
        print(f"Auth module connected to {db_type} successfully")
    else:
        print("WARNING: Auth module failed to connect to database")

# Initialize check
init_db()

def get_users() -> List[Dict[str, Any]]:
    """Fetch all users"""
    conn = _get_conn()
    if not conn: return []
    try:
        cur, rows = _execute(conn, 'SELECT id, username, role, email, last_login, active FROM users;', fetch_all=True)
        if _is_sqlite(conn) and rows:
            for row in rows:
                if row.get('last_login'):
                    try:
                        row['last_login'] = datetime.fromtimestamp(row['last_login'])
                    except Exception:
                        pass
                if 'active' in row:
                    row['active'] = bool(row['active'])
        return rows
    finally:
        conn.close()

def get_user_by_username(username: str) -> Optional[Dict[str, Any]]:
    """Get a user by username"""
    conn = _get_conn()
    if not conn: return None
    try:
        cur, row = _execute(conn, 'SELECT * FROM users WHERE username = %s;', (username,), fetch_one=True)
        if row and _is_sqlite(conn):
            row = dict(row)
            if row.get('created_at'):
                try:
                    row['created_at'] = datetime.fromtimestamp(row['created_at'])
                except Exception:
                    pass
            if row.get('last_login'):
                try:
                    row['last_login'] = datetime.fromtimestamp(row['last_login'])
                except Exception:
                    pass
            if 'active' in row:
                row['active'] = bool(row['active'])
        return row
    finally:
        conn.close()

def authenticate_user(username: str, password: str) -> Tuple[bool, Any]:
    """Authenticate user credentials"""
    user = get_user_by_username(username)
    if not user or not check_password_hash(user['password_hash'], password):
        return False, {'error': 'Invalid credentials.'}
    # Update last login
    conn = _get_conn()
    if conn:
        try:
            now = datetime.now()
            _execute(conn, 'UPDATE users SET last_login = %s WHERE id = %s;', (now, user['id']), commit=True)
        except Exception as e:
            print(f"Warning: Failed to update last_login: {e}")
        finally:
            conn.close()
    return True, user

def create_user(username: str, password: str, email: Optional[str] = None, role: str = 'user') -> Tuple[bool, Any]:
    """Create a new user"""
    if get_user_by_username(username):
        return False, {'error': 'Username already exists.'}
    pwd_hash = generate_password_hash(password, method='scrypt')
    conn = _get_conn()
    if not conn: return False, {'error': 'Database connection failed'}
    
    now = datetime.now()
    try:
        if _is_sqlite(conn):
            cur, _ = _execute(
                conn,
                'INSERT INTO users (username, password_hash, role, email, created_at) VALUES (%s, %s, %s, %s, %s);',
                (username, pwd_hash, role, email, now),
                commit=True
            )
            user_id = cur.lastrowid
        else:
            cur, row = _execute(
                conn,
                'INSERT INTO users (username, password_hash, role, email, created_at) VALUES (%s, %s, %s, %s, %s) RETURNING id;',
                (username, pwd_hash, role, email, now),
                fetch_one=True,
                commit=True
            )
            user_id = row['id'] if isinstance(row, dict) else row[0]
            
        return True, {'id': user_id, 'username': username, 'role': role}
    except Exception as e:
        conn.rollback()
        print(f"Error creating user: {e}")
        return False, {'error': str(e)}
    finally:
        conn.close()

def update_user(user_id: int, data: Dict[str, Any]) -> Tuple[bool, Any]:
    """Update user's password and/or role"""
    fields = []
    params = []
    if data.get('password'):
        pwd_hash = generate_password_hash(data['password'], method='scrypt')
        fields.append('password_hash = %s')
        params.append(pwd_hash)
    if data.get('role'):
        fields.append('role = %s')
        params.append(data['role'])
    if not fields:
        return False, {'error': 'No fields to update.'}
    params.append(user_id)
    
    conn = _get_conn()
    if not conn: return False, {'error': 'Database connection failed'}
    
    try:
        if _is_sqlite(conn):
            cur, row = _execute(conn, 'SELECT username FROM users WHERE id = %s;', (user_id,), fetch_one=True)
            if not row:
                return False, {'error': 'User not found.'}
            username = row['username']
            sql_query = f"UPDATE users SET {', '.join(fields)} WHERE id = %s;"
            _execute(conn, sql_query, params, commit=True)
            return True, {'username': username}
        else:
            sql_query = f"UPDATE users SET {', '.join(fields)} WHERE id = %s RETURNING username;"
            cur, row = _execute(conn, sql_query, params, fetch_one=True, commit=True)
            if row:
                return True, {'username': row['username']}
            return False, {'error': 'User not found.'}
    except Exception as e:
        conn.rollback()
        return False, {'error': str(e)}
    finally:
        conn.close()

def delete_user(user_id: int) -> bool:
    """Delete a user, preventing deletion of last admin"""
    conn = _get_conn()
    if not conn: return False
    try:
        cur, row = _execute(conn, 'SELECT role FROM users WHERE id = %s;', (user_id,), fetch_one=True)
        if not row:
            return False
        if row['role'] == 'admin':
            _, count_row = _execute(conn, "SELECT COUNT(*) AS cnt FROM users WHERE role = 'admin';", fetch_one=True)
            if count_row and count_row['cnt'] <= 1:
                return False
        _execute(conn, 'DELETE FROM users WHERE id = %s;', (user_id,), commit=True)
        return True
    finally:
        conn.close()

def get_all_users() -> List[Dict[str, Any]]:
    """Alias for get_users to match expected API usage"""
    return get_users()

def update_user_details(user_id: int, role: str = None) -> bool:
    """Helper for simple user updates"""
    data = {}
    if role:
        data['role'] = role
    success, _ = update_user(user_id, data)
    return success

def get_user_by_id(user_id: int) -> Optional[Dict[str, Any]]:
    """Fetch user by ID"""
    conn = _get_conn()
    if not conn: return None
    try:
        cur, row = _execute(conn, 'SELECT * FROM users WHERE id = %s;', (user_id,), fetch_one=True)
        if row and _is_sqlite(conn):
            row = dict(row)
            if row.get('created_at'):
                try:
                    row['created_at'] = datetime.fromtimestamp(row['created_at'])
                except Exception:
                    pass
            if row.get('last_login'):
                try:
                    row['last_login'] = datetime.fromtimestamp(row['last_login'])
                except Exception:
                    pass
            if 'active' in row:
                row['active'] = bool(row['active'])
        return row
    finally:
        conn.close()

def update_user_status(user_id: int, active: bool) -> bool:
    """Toggle user active status"""
    conn = _get_conn()
    if not conn: return False
    try:
        _execute(conn, 'UPDATE users SET active = %s WHERE id = %s;', (active, user_id), commit=True)
        return True
    except Exception as e:
        conn.rollback()
        return False
    finally:
        conn.close()