Spaces:
Paused
Paused
| import os | |
| import secrets | |
| import jwt | |
| from datetime import datetime, timedelta | |
| from argon2 import PasswordHasher | |
| from fastapi import Request, HTTPException | |
| from cryptography.fernet import Fernet | |
| from dotenv import load_dotenv | |
| import os | |
| load_dotenv() | |
| ph = PasswordHasher() | |
| JWT_SECRET = os.getenv("JWT_SECRET", "super-secret-wolf-key") | |
| ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY") | |
| ALGORITHM = "HS256" | |
| def decrypt_token(encrypted_token: str) -> str: | |
| if not ENCRYPTION_KEY: | |
| return encrypted_token | |
| try: | |
| f = Fernet(ENCRYPTION_KEY.encode()) | |
| return f.decrypt(encrypted_token.encode()).decode() | |
| except: | |
| return encrypted_token | |
| def hash_password(password: str) -> str: | |
| return ph.hash(password) | |
| def verify_password(hash: str, password: str) -> bool: | |
| try: | |
| return ph.verify(hash, password) | |
| except: | |
| return False | |
| class SecurityLayer: | |
| def __init__(self): | |
| self.allowed_ips = os.getenv("ALLOWED_IPS", "127.0.0.1").split(",") | |
| self.admin_hash = os.getenv("ADMIN_PASSWORD_HASH") | |
| self.pending_2fa = {} # {code: timestamp} | |
| async def check_ip(self, request: Request): | |
| client_ip = request.client.host | |
| if client_ip not in self.allowed_ips and "127.0.0.1" not in self.allowed_ips: | |
| if os.getenv("ENFORCE_IP_WHITELIST", "false").lower() == "true": | |
| raise HTTPException(status_code=403, detail="IP NOT WHITELISTED") | |
| def validate_admin(self, password: str) -> bool: | |
| if not self.admin_hash: | |
| return False | |
| return verify_password(self.admin_hash, password) | |
| def generate_2fa(self): | |
| code = "".join([str(secrets.randbelow(10)) for _ in range(6)]) | |
| return code | |
| def create_access_token(self, data: dict): | |
| to_encode = data.copy() | |
| expire = datetime.utcnow() + timedelta(hours=12) | |
| to_encode.update({"exp": expire}) | |
| return jwt.encode(to_encode, JWT_SECRET, algorithm=ALGORITHM) | |
| def decode_token(self, token: str): | |
| try: | |
| return jwt.decode(token, JWT_SECRET, algorithms=[ALGORITHM]) | |
| except: | |
| return None | |