feat(users): ✨ add user auth and the option for no expiry of… (#37)

Author: thezak48

Dockerfile

@@ -44,6 +44,7 @@ ENV PGID=1000
 ENV DB_PATH=/config/comparisons.db
 ENV UPLOADS_PATH=/data/uploads
 ENV RETENTION_DAYS=7
+ENV ADMIN_INVITATION_CODE=change-me-in-production
 
 # Expose port
 EXPOSE 8000

auth.py

@@ -0,0 +1,242 @@
+"""
+Authentication module for Comps.
+Handles user authentication, invitation codes, and session management.
+"""
+import hashlib
+import secrets
+import sqlite3
+import time
+from typing import Optional, Dict, Any
+from fastapi import Request, HTTPException, Depends, Cookie
+from fastapi.security import APIKeyCookie
+from jose import JWTError, jwt
+from datetime import datetime, timedelta
+
+# Constants
+DB_PATH = 'comparisons.db'
+SECRET_KEY = secrets.token_hex(32)  # Generate a random secret key
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7  # 1 week
+
+# Cookie security
+cookie_sec = APIKeyCookie(name="session")
+
+def hash_invitation_code(code: str) -> str:
+    """Hash an invitation code for secure storage"""
+    return hashlib.sha256(code.encode()).hexdigest()
+
+def create_invitation_code(created_by_id: int) -> str:
+    """Create a new invitation code"""
+    # Generate a random code
+    code = secrets.token_urlsafe(16)
+    
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    
+    # Store the code
+    c.execute(
+        'INSERT INTO invitation_codes (code, created_by) VALUES (?, ?)',
+        (code, created_by_id)
+    )
+    
+    conn.commit()
+    conn.close()
+    
+    return code
+
+def verify_invitation_code(code: str) -> bool:
+    """Verify if an invitation code is valid and unused"""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    
+    c.execute('SELECT is_used FROM invitation_codes WHERE code = ?', (code,))
+    result = c.fetchone()
+    
+    conn.close()
+    
+    # Code is valid if it exists and is not used
+    return result is not None and not result[0]
+
+def register_user(username: str, invitation_code: str) -> Optional[Dict[str, Any]]:
+    """Register a new user with an invitation code"""
+    if not verify_invitation_code(invitation_code):
+        return None
+    
+    # Hash the invitation code for storage
+    code_hash = hash_invitation_code(invitation_code)
+    
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    
+    try:
+        # Check if username already exists
+        c.execute('SELECT id FROM users WHERE username = ?', (username,))
+        if c.fetchone():
+            conn.close()
+            return None
+        
+        # Create the user
+        c.execute(
+            'INSERT INTO users (username, invitation_code_hash, never_expire_comparisons) VALUES (?, ?, ?)',
+            (username, code_hash, 1)  # All invited users get permanent comparisons
+        )
+        user_id = c.lastrowid
+        
+        # Mark the invitation code as used
+        c.execute(
+            'UPDATE invitation_codes SET is_used = 1, used_by = ? WHERE code = ?',
+            (user_id, invitation_code)
+        )
+        
+        # Get the user data
+        c.execute('SELECT id, username, is_admin, never_expire_comparisons FROM users WHERE id = ?', (user_id,))
+        user = c.fetchone()
+        
+        conn.commit()
+        
+        if user:
+            return {
+                "id": user[0],
+                "username": user[1],
+                "is_admin": bool(user[2]),
+                "never_expire_comparisons": bool(user[3])
+            }
+        return None
+    except Exception as e:
+        print(f"Error registering user: {e}")
+        conn.rollback()
+        return None
+    finally:
+        conn.close()
+
+def authenticate_user(username: str, invitation_code: str) -> Optional[Dict[str, Any]]:
+    """Authenticate a user with their username and invitation code"""
+    # Hash the invitation code for comparison
+    code_hash = hash_invitation_code(invitation_code)
+    
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    
+    c.execute(
+        'SELECT id, username, is_admin, never_expire_comparisons FROM users WHERE username = ? AND invitation_code_hash = ?',
+        (username, code_hash)
+    )
+    user = c.fetchone()
+    
+    conn.close()
+    
+    if user:
+        return {
+            "id": user[0],
+            "username": user[1],
+            "is_admin": bool(user[2]),
+            "never_expire_comparisons": bool(user[3])
+        }
+    return None
+
+def create_access_token(data: dict) -> str:
+    """Create a JWT access token"""
+    to_encode = data.copy()
+    expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+def get_current_user(session: str = Depends(cookie_sec)) -> Optional[Dict[str, Any]]:
+    """Get the current user from the session cookie"""
+    try:
+        payload = jwt.decode(session, SECRET_KEY, algorithms=[ALGORITHM])
+        user_id = payload.get("sub")
+        if user_id is None:
+            return None
+        
+        conn = sqlite3.connect(DB_PATH)
+        c = conn.cursor()
+        
+        c.execute(
+            'SELECT id, username, is_admin, never_expire_comparisons FROM users WHERE id = ?',
+            (user_id,)
+        )
+        user = c.fetchone()
+        
+        conn.close()
+        
+        if user:
+            return {
+                "id": user[0],
+                "username": user[1],
+                "is_admin": bool(user[2]),
+                "never_expire_comparisons": bool(user[3])
+            }
+        return None
+    except JWTError:
+        return None
+    except Exception as e:
+        print(f"Error getting current user: {e}")
+        return None
+
+def get_user_invitation_codes(user_id: int) -> list:
+    """Get all invitation codes created by a user"""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    
+    c.execute('''
+        SELECT ic.code, ic.is_used, u.username, ic.created_at
+        FROM invitation_codes ic
+        LEFT JOIN users u ON ic.used_by = u.id
+        WHERE ic.created_by = ?
+        ORDER BY ic.created_at DESC
+    ''', (user_id,))
+    
+    codes = []
+    for code, is_used, used_by, created_at in c.fetchall():
+        codes.append({
+            "code": code,
+            "is_used": bool(is_used),
+            "used_by": used_by,
+            "created_at": created_at
+        })
+    
+    conn.close()
+    return codes
+
+def is_admin(user: dict) -> bool:
+    """Check if a user is an admin"""
+    return user and user.get("is_admin", False)
+
+async def get_optional_user(request: Request) -> Optional[Dict[str, Any]]:
+    """Get the current user if logged in, otherwise return None"""
+    session = request.cookies.get("session")
+    if not session:
+        return None
+    
+    try:
+        payload = jwt.decode(session, SECRET_KEY, algorithms=[ALGORITHM])
+        user_id = payload.get("sub")
+        if user_id is None:
+            return None
+        
+        conn = sqlite3.connect(DB_PATH)
+        c = conn.cursor()
+        
+        c.execute(
+            'SELECT id, username, is_admin, never_expire_comparisons FROM users WHERE id = ?',
+            (user_id,)
+        )
+        user = c.fetchone()
+        
+        conn.close()
+        
+        if user:
+            return {
+                "id": user[0],
+                "username": user[1],
+                "is_admin": bool(user[2]),
+                "never_expire_comparisons": bool(user[3])
+            }
+        return None
+    except JWTError:
+        return None
+    except Exception as e:
+        print(f"Error getting optional user: {e}")
+        return None

database.py

@@ -1,6 +1,6 @@
 import sqlite3
 import os
-import shutil
+import shutil 
 from datetime import datetime, timedelta
 from typing import List, Optional
 from migrations.manager import MigrationManager
@@ -12,15 +12,31 @@ def init_db():
     migration_manager = MigrationManager()
     migration_manager.migrate(DB_PATH)
 
-def create_comparison(comparison_id: str, name: Optional[str], show_name: Optional[str], tags: Optional[List[str]], metadata: dict):
+def create_comparison(comparison_id: str, name: Optional[str], show_name: Optional[str], tags: Optional[List[str]], metadata: dict, user_id: Optional[int] = None):
     conn = sqlite3.connect(DB_PATH)
     c = conn.cursor()
 
     c.execute(
-        'INSERT INTO comparisons (id, name, show_name, total_rows, total_columns, expiration_type, expiration_days) VALUES (?, ?, ?, ?, ?, ?, ?)',
-        (comparison_id, name, show_name, metadata.get('total_rows', 1), metadata.get('total_columns', 2), metadata.get('expiration_type', 'from_last_access'), int(metadata.get('expiration_days', 7)))
+        'INSERT INTO comparisons (id, name, show_name, total_rows, total_columns, expiration_type, expiration_days, never_expire) VALUES (?, ?, ?, ?, ?, ?, ?, ?)',
+        (comparison_id, name, show_name, metadata.get('total_rows', 1), metadata.get('total_columns', 2), 
+         metadata.get('expiration_type', 'from_last_access'), int(metadata.get('expiration_days', 7)), 0)
     )
-    c.execute('CREATE INDEX IF NOT EXISTS idx_image_positions ON image_positions(comparison_id, row_number, column_position)')
+
+    # Set default expiration based on metadata
+    if metadata.get('never_expire') is not None:
+        c.execute('UPDATE comparisons SET never_expire = ? WHERE id = ?', 
+                 (1 if metadata.get('never_expire') else 0, comparison_id))
+    
+    # If user is authenticated, associate the comparison with them and set never_expire if applicable
+    if user_id is not None:
+        c.execute('UPDATE comparisons SET user_id = ? WHERE id = ?', (user_id, comparison_id))
+        c.execute('SELECT never_expire_comparisons FROM users WHERE id = ?', (user_id,))
+        never_expire = c.fetchone()[0]
+        if never_expire:
+            # Only set never_expire if the user didn't explicitly choose to have it expire
+            if metadata.get('never_expire') is None:
+                c.execute('UPDATE comparisons SET never_expire = 1 WHERE id = ?', (comparison_id,))
+
 
     if tags:
         for tag in tags:
@@ -60,6 +76,11 @@ def get_comparison(comparison_id: str):
         c.execute('SELECT tag FROM tags WHERE comparison_id = ?', (comparison_id,))
         tags = [row[0] for row in c.fetchall()] 
 
+        # Get user information if available
+        c.execute('SELECT user_id, never_expire FROM comparisons WHERE id = ?', (comparison_id,))
+        user_info = c.fetchone()
+        user_id, never_expire = user_info if user_info else (None, 0)
+        
         return {
             'id': comparison[0],
             'name': comparison[1],
@@ -70,7 +91,9 @@ def get_comparison(comparison_id: str):
             'expiration_type': comparison[5] or 'from_last_access',
             'expiration_days': comparison[6] or 7,
             'created_at': comparison[7],
-            'last_accessed': comparison[8]
+            'last_accessed': comparison[8],
+            'user_id': user_id,
+            'never_expire': bool(never_expire)
         }
 
     return None
@@ -180,16 +203,21 @@ def get_expired_comparisons(retention_days: int):
 
         if 'last_accessed' in columns and 'expiration_type' in columns and 'expiration_days' in columns:
             # Get all comparisons with their expiration settings
-            c.execute('SELECT id, expiration_type, expiration_days, created_at, last_accessed FROM comparisons')
+            c.execute('SELECT id, expiration_type, expiration_days, created_at, last_accessed, never_expire FROM comparisons')
             comparisons = c.fetchall()
-            
+
             print(f"Checking for expired comparisons with retention_days={retention_days}")
             print(f"Found {len(comparisons)} comparisons to check for expiration")
             current_time = datetime.now()
-            for comp_id, exp_type, exp_days, created_at, last_accessed in comparisons:
+            for comp_id, exp_type, exp_days, created_at, last_accessed, never_expire in comparisons:
                 # Use comparison's own expiration days if available, otherwise use default
                 days = exp_days if exp_days is not None else retention_days
                 print(f"Checking comparison {comp_id}: type={exp_type}, days={days}, created={created_at}, last_accessed={last_accessed}")
+
+                # Check if this comparison is marked as never expire
+                if never_expire:
+                    print(f"  Comparison {comp_id} is marked as never expire, skipping")
+                    continue
 
                 if exp_type == 'from_creation' and created_at:
                     cutoff_date = datetime.strptime(created_at, '%Y-%m-%d %H:%M:%S') + timedelta(days=days)

docker-compose.yml

@@ -13,6 +13,7 @@ services:
       - PGID=1000
       - DB_PATH=/config/comparisons.db
       - UPLOADS_PATH=/data/uploads
+      - ADMIN_INVITATION_CODE=your-secure-admin-code
     volumes:
       - ./config:/config
       - ./data:/data