Stock_Agent_optimized / auth /auth_manager.py
cryogenic22's picture
Update auth/auth_manager.py
7ac74fa verified
# auth/auth_manager.py
import streamlit as st
import json
import os
from pathlib import Path
import hashlib
from datetime import datetime, timedelta
class AuthManager:
def __init__(self):
"""Initialize AuthManager with storage in HuggingFace's /data directory"""
# Use HuggingFace's /data directory for persistent storage
self.storage_path = Path("/data/users")
self.ensure_storage()
# Initialize demo user if not exists
self._init_demo_user()
# Initialize session state for auth
if 'auth' not in st.session_state:
st.session_state.auth = {
'authenticated': False,
'user': None,
'login_time': None
}
def ensure_storage(self):
"""Ensure storage directories exist and are writable"""
try:
# Create main storage directory
self.storage_path.mkdir(parents=True, exist_ok=True)
# Test write permissions
test_file = self.storage_path / ".write_test"
test_file.touch()
test_file.unlink()
except Exception as e:
st.error(f"Failed to access storage: {str(e)}")
raise RuntimeError("Storage is not accessible")
def _init_demo_user(self):
"""Initialize demo Synaptyx user"""
demo_user = {
"username": "synaptyx",
"password_hash": self._hash_password("demo"),
"display_name": "Synaptyx Demo",
"storage_path": "/data/users/synaptyx",
"created_at": datetime.now().isoformat(),
"last_login": None,
"settings": {
"expertise_level": "Intermediate",
"theme": "default",
"notifications_enabled": True
}
}
user_file = self.storage_path / "synaptyx.json"
if not user_file.exists():
# Create user storage directories
user_storage = Path(demo_user["storage_path"])
for dir_name in ["chats", "images", "context", "exports"]:
(user_storage / dir_name).mkdir(parents=True, exist_ok=True)
# Save user data
with open(user_file, "w") as f:
json.dump(demo_user, f, indent=4)
def _hash_password(self, password):
"""Hash password using SHA-256 with salt"""
salt = "synaptyx_salt" # In production, use a proper salt management system
salted = password + salt
return hashlib.sha256(salted.encode()).hexdigest()
def login(self, username, password):
"""Authenticate user and set up session"""
try:
user_file = self.storage_path / f"{username}.json"
if not user_file.exists():
return False
with open(user_file, "r") as f:
user_data = json.load(f)
if user_data["password_hash"] == self._hash_password(password):
# Update last login
user_data["last_login"] = datetime.now().isoformat()
with open(user_file, "w") as f:
json.dump(user_data, f, indent=4)
# Update session state
st.session_state.auth = {
'authenticated': True,
'user': user_data,
'login_time': datetime.now().isoformat()
}
return True
return False
except Exception as e:
st.error(f"Login error: {str(e)}")
return False
def is_authenticated(self):
"""Check if user is authenticated and session is valid"""
if not st.session_state.auth.get('authenticated'):
return False
# Check session age (optional)
if st.session_state.auth.get('login_time'):
login_time = datetime.fromisoformat(st.session_state.auth['login_time'])
session_age = datetime.now() - login_time
# Session timeout after 12 hours
if session_age > timedelta(hours=12):
self.logout()
return False
return True
def get_user_storage_paths(self):
"""Get user's storage paths"""
if not self.is_authenticated():
return None
base_path = Path(st.session_state.auth['user']["storage_path"])
return {
"chats": base_path / "chats",
"images": base_path / "images",
"context": base_path / "context",
"exports": base_path / "exports"
}
def get_user_settings(self):
"""Get user's settings"""
if not self.is_authenticated():
return None
return st.session_state.auth['user'].get('settings', {})
def update_user_settings(self, settings):
"""Update user's settings"""
if not self.is_authenticated():
return False
try:
user_file = self.storage_path / f"{st.session_state.auth['user']['username']}.json"
with open(user_file, "r") as f:
user_data = json.load(f)
user_data['settings'].update(settings)
with open(user_file, "w") as f:
json.dump(user_data, f, indent=4)
# Update session state
st.session_state.auth['user'] = user_data
return True
except Exception as e:
st.error(f"Failed to update settings: {str(e)}")
return False
def logout(self):
"""Log out user and clear session"""
if 'auth' in st.session_state:
st.session_state.auth = {
'authenticated': False,
'user': None,
'login_time': None
}
def get_user_stats(self):
"""Get user's usage statistics"""
if not self.is_authenticated():
return None
try:
stats = {
'total_analyses': 0,
'total_chats': 0,
'storage_used': 0
}
paths = self.get_user_storage_paths()
if not paths:
return stats
# Count analyses and chats
stats['total_chats'] = len(list(paths['chats'].glob('*.json')))
# Calculate storage used
for path in paths.values():
if path.exists():
stats['storage_used'] += sum(f.stat().st_size
for f in path.rglob('*')
if f.is_file())
# Convert to MB
stats['storage_used'] = round(stats['storage_used'] / (1024 * 1024), 2)
return stats
except Exception as e:
st.error(f"Failed to get user stats: {str(e)}")
return None