Spaces:
Sleeping
Sleeping
File size: 3,843 Bytes
2978bba 0457a2b 2978bba 0457a2b 2978bba 56af25f 0457a2b 2978bba 0457a2b 2978bba 0457a2b 2978bba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | from functools import wraps
from flask import request, jsonify, g, session
import hashlib
try:
import psycopg2
HAS_POSTGRES = True
except ImportError:
psycopg2 = None
HAS_POSTGRES = False
import os
IS_HF_SPACE = os.environ.get('HF_SPACE', '0') == '1'
def get_db_connection():
if IS_HF_SPACE:
raise RuntimeError("PostgreSQL database is disabled on HuggingFace Spaces.")
if not HAS_POSTGRES:
raise RuntimeError("PostgreSQL / psycopg2 is not installed or available.")
# Helper to get DB connection (duplicated from telemetry for standalone usage)
# In a real app config should be shared properly
try:
import config
DB_PARAMS = {
"dbname": getattr(config, 'DB_NAME', 'morphguard'),
"user": getattr(config, 'DB_USER', 'morphguard'),
"password": getattr(config, 'DB_PASSWORD', 'morphguard'),
"host": getattr(config, 'DB_HOST', 'localhost'),
"port": getattr(config, 'DB_PORT', 5432)
}
except (ImportError, AttributeError):
DB_PARAMS = {
"dbname": os.environ.get('MORPHGUARD_DB_NAME', 'morphguard'),
"user": os.environ.get('MORPHGUARD_DB_USER', 'morphguard'),
"password": os.environ.get('MORPHGUARD_DB_PASS', 'morphguard'),
"host": os.environ.get('MORPHGUARD_DB_HOST', 'localhost'),
"port": int(os.environ.get('MORPHGUARD_DB_PORT', 5432))
}
return psycopg2.connect(**DB_PARAMS)
def hash_key(key):
return hashlib.sha256(key.encode()).hexdigest()
def validate_api_key(key):
"""
Validate API key against database.
Returns (is_valid, key_info_dict)
"""
if IS_HF_SPACE or not HAS_POSTGRES:
return False, None
hashed = hash_key(key)
try:
conn = get_db_connection()
cur = conn.cursor()
cur.execute(
"SELECT id, name, permissions FROM api_keys WHERE key_hash = %s AND is_active = TRUE",
(hashed,)
)
row = cur.fetchone()
cur.close()
conn.close()
if row:
return True, {'id': row[0], 'name': row[1], 'permissions': row[2]}
return False, None
except Exception as e:
import logging
logging.getLogger("APIAuth").error(f"API Key validation error: {e}")
return False, None
def require_api_1(require_perm=None):
"""
Decorator to require API key or Session login.
If session is active, proceeds.
If not, checks X-API-Key header.
"""
def decorator(f):
@wraps(f)
def decorated_function(*args, **kwargs):
# 1. Check Session
if session.get('logged_in'):
# Admin/Mod check if perms needed (simplified)
# In robust system, check specific user permissions
return f(*args, **kwargs)
# 2. Check API Key
api_key = request.headers.get('X-API-Key')
if api_key:
is_valid, info = validate_api_key(api_key)
if is_valid:
# Check permissions if required
if require_perm:
perms = info.get('permissions', [])
# if 'admin' in perms, allow all. validation logic here.
if require_perm not in perms and 'admin' not in perms:
return jsonify({'error': 'Insufficient API Permissions'}), 403
g.api_user = info
return f(*args, **kwargs)
return jsonify({'error': 'Unauthorized. Session or Valid API Key required.'}), 401
return decorated_function
return decorator
# Alias just primarily for simple usage
require_api_key = require_api_1
|