shadowhunter_auth.py

#!/usr/bin/env python3
"""
ShadowHunter - Dark Web Credential Intelligence Platform
Module: Authentication & User Management
Author: Fevra
Version: 0.1.0

Implements JWT-based authentication, user management, and role-based access control.
"""

from fastapi import Depends, HTTPException, status, Security
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm, HTTPBearer, HTTPAuthorizationCredentials
from jose import JWTError, jwt
from passlib.context import CryptContext
from pydantic import BaseModel, EmailStr, validator
from datetime import datetime, timedelta
from typing import Optional, List, Dict, Any
import secrets
import logging
from enum import Enum

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('ShadowHunter.Auth')

# ============================================================================
# CONFIGURATION
# ============================================================================

# Secret key for JWT - In production, use environment variable
SECRET_KEY = secrets.token_urlsafe(32)
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30
REFRESH_TOKEN_EXPIRE_DAYS = 7

# Password hashing
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# OAuth2 scheme
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/login")
security = HTTPBearer()

# ============================================================================
# ENUMS
# ============================================================================

class UserRole(str, Enum):
    """User roles with different permission levels"""
    ADMIN = "admin"              # Full access to everything
    ANALYST = "analyst"          # View and analyze threats
    VIEWER = "viewer"            # Read-only access
    API_USER = "api_user"        # API access only

class PermissionLevel(str, Enum):
    """Permission levels for resources"""
    READ = "read"
    WRITE = "write"
    DELETE = "delete"
    ADMIN = "admin"

# ============================================================================
# DATA MODELS
# ============================================================================

class User(BaseModel):
    """User model"""
    user_id: str
    email: EmailStr
    username: str
    full_name: Optional[str] = None
    role: UserRole = UserRole.VIEWER
    is_active: bool = True
    is_verified: bool = False
    created_at: datetime
    last_login: Optional[datetime] = None
    preferences: Dict[str, Any] = {}

class UserInDB(User):
    """User model with hashed password"""
    hashed_password: str

class UserCreate(BaseModel):
    """Model for user registration"""
    email: EmailStr
    username: str
    password: str
    full_name: Optional[str] = None
    role: UserRole = UserRole.VIEWER
    
    @validator('username')
    def username_alphanumeric(cls, v):
        if not v.isalnum():
            raise ValueError('Username must be alphanumeric')
        if len(v) < 3:
            raise ValueError('Username must be at least 3 characters')
        return v
    
    @validator('password')
    def password_strength(cls, v):
        if len(v) < 8:
            raise ValueError('Password must be at least 8 characters')
        if not any(c.isupper() for c in v):
            raise ValueError('Password must contain uppercase letter')
        if not any(c.islower() for c in v):
            raise ValueError('Password must contain lowercase letter')
        if not any(c.isdigit() for c in v):
            raise ValueError('Password must contain digit')
        return v

class UserUpdate(BaseModel):
    """Model for user updates"""
    full_name: Optional[str] = None
    email: Optional[EmailStr] = None
    role: Optional[UserRole] = None
    is_active: Optional[bool] = None
    preferences: Optional[Dict[str, Any]] = None

class Token(BaseModel):
    """JWT token response"""
    access_token: str
    refresh_token: str
    token_type: str = "bearer"
    expires_in: int

class TokenData(BaseModel):
    """Data extracted from JWT token"""
    username: Optional[str] = None
    user_id: Optional[str] = None
    role: Optional[UserRole] = None
    scopes: List[str] = []

class PasswordReset(BaseModel):
    """Password reset request"""
    email: EmailStr

class PasswordResetConfirm(BaseModel):
    """Password reset confirmation"""
    token: str
    new_password: str

class LoginHistory(BaseModel):
    """Login history entry"""
    user_id: str
    timestamp: datetime
    ip_address: str
    user_agent: str
    success: bool

class APIKey(BaseModel):
    """API key model"""
    key_id: str
    user_id: str
    key_name: str
    api_key: str
    created_at: datetime
    expires_at: Optional[datetime] = None
    last_used: Optional[datetime] = None
    is_active: bool = True

# ============================================================================
# USER DATABASE (In-memory - replace with real database)
# ============================================================================

class UserDatabase:
    """In-memory user database"""
    
    def __init__(self):
        self.users: Dict[str, UserInDB] = {}
        self.login_history: List[LoginHistory] = []
        self.api_keys: Dict[str, APIKey] = {}
        self.password_reset_tokens: Dict[str, str] = {}  # token -> email
        
        # Create default admin user
        self._create_default_users()
    
    def _create_default_users(self):
        """Create default users for testing"""
        admin_user = UserInDB(
            user_id="admin_001",
            email="admin@shadowhunter.io",
            username="admin",
            full_name="Admin User",
            role=UserRole.ADMIN,
            is_active=True,
            is_verified=True,
            created_at=datetime.utcnow(),
            hashed_password=pwd_context.hash("Admin123!"),
            preferences={
                "email_alerts": True,
                "alert_severity": ["CRITICAL", "HIGH"],
                "monitored_domains": []
            }
        )
        self.users["admin"] = admin_user
        
        analyst_user = UserInDB(
            user_id="analyst_001",
            email="analyst@shadowhunter.io",
            username="analyst",
            full_name="Security Analyst",
            role=UserRole.ANALYST,
            is_active=True,
            is_verified=True,
            created_at=datetime.utcnow(),
            hashed_password=pwd_context.hash("Analyst123!"),
            preferences={
                "email_alerts": True,
                "alert_severity": ["CRITICAL", "HIGH", "MEDIUM"]
            }
        )
        self.users["analyst"] = analyst_user
        
        logger.info("Default users created: admin, analyst")
    
    def get_user(self, username: str) -> Optional[UserInDB]:
        """Get user by username"""
        return self.users.get(username)
    
    def get_user_by_email(self, email: str) -> Optional[UserInDB]:
        """Get user by email"""
        for user in self.users.values():
            if user.email == email:
                return user
        return None
    
    def create_user(self, user_create: UserCreate) -> UserInDB:
        """Create new user"""
        if user_create.username in self.users:
            raise ValueError("Username already exists")
        
        if self.get_user_by_email(user_create.email):
            raise ValueError("Email already exists")
        
        user = UserInDB(
            user_id=f"user_{secrets.token_hex(8)}",
            email=user_create.email,
            username=user_create.username,
            full_name=user_create.full_name,
            role=user_create.role,
            is_active=True,
            is_verified=False,
            created_at=datetime.utcnow(),
            hashed_password=pwd_context.hash(user_create.password),
            preferences={
                "email_alerts": False,
                "alert_severity": ["CRITICAL"]
            }
        )
        
        self.users[user_create.username] = user
        logger.info(f"User created: {user.username} ({user.role})")
        return user
    
    def update_user(self, username: str, user_update: UserUpdate) -> UserInDB:
        """Update user"""
        user = self.get_user(username)
        if not user:
            raise ValueError("User not found")
        
        update_data = user_update.dict(exclude_unset=True)
        for field, value in update_data.items():
            setattr(user, field, value)
        
        return user
    
    def delete_user(self, username: str) -> bool:
        """Delete user"""
        if username in self.users:
            del self.users[username]
            logger.info(f"User deleted: {username}")
            return True
        return False
    
    def list_users(self) -> List[User]:
        """List all users (without passwords)"""
        return [User(**user.dict()) for user in self.users.values()]
    
    def record_login(self, login: LoginHistory):
        """Record login attempt"""
        self.login_history.append(login)
        
        if login.success:
            user = self.get_user_by_id(login.user_id)
            if user:
                user.last_login = login.timestamp
    
    def get_user_by_id(self, user_id: str) -> Optional[UserInDB]:
        """Get user by ID"""
        for user in self.users.values():
            if user.user_id == user_id:
                return user
        return None
    
    def create_api_key(self, user_id: str, key_name: str, expires_days: Optional[int] = None) -> APIKey:
        """Create API key for user"""
        api_key = APIKey(
            key_id=f"key_{secrets.token_hex(8)}",
            user_id=user_id,
            key_name=key_name,
            api_key=f"shk_{secrets.token_urlsafe(32)}",
            created_at=datetime.utcnow(),
            expires_at=datetime.utcnow() + timedelta(days=expires_days) if expires_days else None,
            is_active=True
        )
        
        self.api_keys[api_key.api_key] = api_key
        logger.info(f"API key created: {key_name} for user {user_id}")
        return api_key
    
    def get_api_key(self, api_key: str) -> Optional[APIKey]:
        """Get API key"""
        return self.api_keys.get(api_key)
    
    def revoke_api_key(self, api_key: str) -> bool:
        """Revoke API key"""
        if api_key in self.api_keys:
            self.api_keys[api_key].is_active = False
            logger.info(f"API key revoked: {api_key}")
            return True
        return False

# Initialize database
user_db = UserDatabase()

# ============================================================================
# AUTHENTICATION FUNCTIONS
# ============================================================================

def verify_password(plain_password: str, hashed_password: str) -> bool:
    """Verify password against hash"""
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password: str) -> str:
    """Hash password"""
    return pwd_context.hash(password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None) -> str:
    """Create JWT access token"""
    to_encode = data.copy()
    
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    
    to_encode.update({"exp": expire, "type": "access"})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

def create_refresh_token(data: dict) -> str:
    """Create JWT refresh token"""
    to_encode = data.copy()
    expire = datetime.utcnow() + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
    to_encode.update({"exp": expire, "type": "refresh"})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

def authenticate_user(username: str, password: str) -> Optional[UserInDB]:
    """Authenticate user with username and password"""
    user = user_db.get_user(username)
    if not user:
        return None
    if not verify_password(password, user.hashed_password):
        return None
    if not user.is_active:
        return None
    return user

async def get_current_user(token: str = Depends(oauth2_scheme)) -> User:
    """Get current user from JWT token"""
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        user_id: str = payload.get("user_id")
        token_type: str = payload.get("type")
        
        if username is None or token_type != "access":
            raise credentials_exception
            
        token_data = TokenData(username=username, user_id=user_id)
        
    except JWTError:
        raise credentials_exception
    
    user = user_db.get_user(username=token_data.username)
    if user is None:
        raise credentials_exception
    
    return User(**user.dict())

async def get_current_active_user(current_user: User = Depends(get_current_user)) -> User:
    """Get current active user"""
    if not current_user.is_active:
        raise HTTPException(status_code=400, detail="Inactive user")
    return current_user

# ============================================================================
# ROLE-BASED ACCESS CONTROL
# ============================================================================

class RoleChecker:
    """Check if user has required role"""
    
    def __init__(self, allowed_roles: List[UserRole]):
        self.allowed_roles = allowed_roles
    
    def __call__(self, user: User = Depends(get_current_active_user)) -> User:
        if user.role not in self.allowed_roles:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=f"User role '{user.role}' not authorized. Required: {self.allowed_roles}"
            )
        return user

# Role dependencies
require_admin = RoleChecker([UserRole.ADMIN])
require_analyst = RoleChecker([UserRole.ADMIN, UserRole.ANALYST])
require_viewer = RoleChecker([UserRole.ADMIN, UserRole.ANALYST, UserRole.VIEWER])

# ============================================================================
# API KEY AUTHENTICATION
# ============================================================================

async def get_api_key_user(credentials: HTTPAuthorizationCredentials = Security(security)) -> User:
    """Authenticate user via API key"""
    api_key = credentials.credentials
    
    if not api_key.startswith("shk_"):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key format"
        )
    
    key_obj = user_db.get_api_key(api_key)
    
    if not key_obj:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key"
        )
    
    if not key_obj.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="API key has been revoked"
        )
    
    if key_obj.expires_at and datetime.utcnow() > key_obj.expires_at:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="API key has expired"
        )
    
    # Update last used
    key_obj.last_used = datetime.utcnow()
    
    user = user_db.get_user_by_id(key_obj.user_id)
    if not user or not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="User account is inactive"
        )
    
    return User(**user.dict())

# ============================================================================
# FASTAPI INTEGRATION EXAMPLE
# ============================================================================

def get_auth_router():
    """Get authentication router for FastAPI integration"""
    from fastapi import APIRouter
    
    router = APIRouter(prefix="/api/auth", tags=["authentication"])
    
    @router.post("/register", response_model=User, status_code=status.HTTP_201_CREATED)
    async def register(user_create: UserCreate):
        """Register new user"""
        try:
            user = user_db.create_user(user_create)
            return User(**user.dict())
        except ValueError as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=str(e)
            )
    
    @router.post("/login", response_model=Token)
    async def login(form_data: OAuth2PasswordRequestForm = Depends()):
        """Login and get JWT tokens"""
        user = authenticate_user(form_data.username, form_data.password)
        
        if not user:
            # Record failed login
            user_db.record_login(LoginHistory(
                user_id="unknown",
                timestamp=datetime.utcnow(),
                ip_address="0.0.0.0",
                user_agent="",
                success=False
            ))
            
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Incorrect username or password",
                headers={"WWW-Authenticate": "Bearer"},
            )
        
        # Record successful login
        user_db.record_login(LoginHistory(
            user_id=user.user_id,
            timestamp=datetime.utcnow(),
            ip_address="0.0.0.0",
            user_agent="",
            success=True
        ))
        
        # Create tokens
        access_token = create_access_token(
            data={"sub": user.username, "user_id": user.user_id, "role": user.role}
        )
        refresh_token = create_refresh_token(
            data={"sub": user.username, "user_id": user.user_id}
        )
        
        return Token(
            access_token=access_token,
            refresh_token=refresh_token,
            expires_in=ACCESS_TOKEN_EXPIRE_MINUTES * 60
        )
    
    @router.post("/refresh", response_model=Token)
    async def refresh_token(refresh_token: str):
        """Refresh access token using refresh token"""
        try:
            payload = jwt.decode(refresh_token, SECRET_KEY, algorithms=[ALGORITHM])
            username: str = payload.get("sub")
            token_type: str = payload.get("type")
            
            if token_type != "refresh":
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="Invalid token type"
                )
            
            user = user_db.get_user(username)
            if not user:
                raise HTTPException(
                    status_code=status.HTTP_401_UNAUTHORIZED,
                    detail="User not found"
                )
            
            # Create new access token
            new_access_token = create_access_token(
                data={"sub": user.username, "user_id": user.user_id, "role": user.role}
            )
            
            return Token(
                access_token=new_access_token,
                refresh_token=refresh_token,
                expires_in=ACCESS_TOKEN_EXPIRE_MINUTES * 60
            )
            
        except JWTError:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid refresh token"
            )
    
    @router.get("/me", response_model=User)
    async def get_me(current_user: User = Depends(get_current_active_user)):
        """Get current user information"""
        return current_user
    
    @router.put("/me", response_model=User)
    async def update_me(
        user_update: UserUpdate,
        current_user: User = Depends(get_current_active_user)
    ):
        """Update current user information"""
        # Users can't change their own role
        if user_update.role:
            user_update.role = None
        
        updated_user = user_db.update_user(current_user.username, user_update)
        return User(**updated_user.dict())
    
    @router.post("/api-keys", response_model=APIKey)
    async def create_api_key(
        key_name: str,
        expires_days: Optional[int] = None,
        current_user: User = Depends(get_current_active_user)
    ):
        """Create new API key"""
        api_key = user_db.create_api_key(current_user.user_id, key_name, expires_days)
        return api_key
    
    @router.get("/api-keys", response_model=List[APIKey])
    async def list_api_keys(current_user: User = Depends(get_current_active_user)):
        """List user's API keys"""
        return [key for key in user_db.api_keys.values() if key.user_id == current_user.user_id]
    
    @router.delete("/api-keys/{api_key}")
    async def revoke_api_key(
        api_key: str,
        current_user: User = Depends(get_current_active_user)
    ):
        """Revoke API key"""
        key_obj = user_db.get_api_key(api_key)
        if not key_obj or key_obj.user_id != current_user.user_id:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="API key not found"
            )
        
        user_db.revoke_api_key(api_key)
        return {"status": "revoked"}
    
    @router.post("/password-reset")
    async def request_password_reset(reset_request: PasswordReset):
        """Request password reset"""
        user = user_db.get_user_by_email(reset_request.email)
        if user:
            # Generate reset token
            reset_token = secrets.token_urlsafe(32)
            user_db.password_reset_tokens[reset_token] = user.email
            
            # In production, send email here
            logger.info(f"Password reset requested for {user.email}")
            # send_password_reset_email(user.email, reset_token)
        
        # Always return success to prevent email enumeration
        return {"message": "If the email exists, a reset link has been sent"}
    
    @router.post("/password-reset/confirm")
    async def confirm_password_reset(reset_confirm: PasswordResetConfirm):
        """Confirm password reset"""
        email = user_db.password_reset_tokens.get(reset_confirm.token)
        if not email:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Invalid or expired reset token"
            )
        
        user = user_db.get_user_by_email(email)
        if not user:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="User not found"
            )
        
        # Update password
        user.hashed_password = get_password_hash(reset_confirm.new_password)
        
        # Remove used token
        del user_db.password_reset_tokens[reset_confirm.token]
        
        logger.info(f"Password reset completed for {email}")
        return {"message": "Password has been reset successfully"}
    
    # Admin endpoints
    @router.get("/users", response_model=List[User])
    async def list_users(admin_user: User = Depends(require_admin)):
        """List all users (admin only)"""
        return user_db.list_users()
    
    @router.put("/users/{username}", response_model=User)
    async def update_user(
        username: str,
        user_update: UserUpdate,
        admin_user: User = Depends(require_admin)
    ):
        """Update user (admin only)"""
        try:
            updated_user = user_db.update_user(username, user_update)
            return User(**updated_user.dict())
        except ValueError as e:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=str(e)
            )
    
    @router.delete("/users/{username}")
    async def delete_user(
        username: str,
        admin_user: User = Depends(require_admin)
    ):
        """Delete user (admin only)"""
        if username == admin_user.username:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Cannot delete your own account"
            )
        
        if user_db.delete_user(username):
            return {"status": "deleted"}
        else:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="User not found"
            )
    
    @router.get("/login-history")
    async def get_login_history(
        limit: int = 50,
        admin_user: User = Depends(require_admin)
    ):
        """Get login history (admin only)"""
        return user_db.login_history[-limit:]
    
    return router

# ============================================================================
# USAGE EXAMPLE
# ============================================================================

if __name__ == "__main__":
    print("="*80)
    print("ShadowHunter Authentication System")
    print("="*80)
    
    print("\n📝 Default Users:")
    print("  Admin:")
    print("    Username: admin")
    print("    Password: Admin123!")
    print("    Role: ADMIN")
    print("\n  Analyst:")
    print("    Username: analyst")
    print("    Password: Analyst123!")
    print("    Role: ANALYST")
    
    print("\n🔑 JWT Configuration:")
    print(f"  Access Token Expiry: {ACCESS_TOKEN_EXPIRE_MINUTES} minutes")
    print(f"  Refresh Token Expiry: {REFRESH_TOKEN_EXPIRE_DAYS} days")
    print(f"  Algorithm: {ALGORITHM}")
    
    print("\n👥 User Roles:")
    for role in UserRole:
        print(f"  - {role.value.upper()}")
    
    print("\n" + "="*80)
    print("Integration: Add to your FastAPI app")
    print("="*80)
    print("""
from shadowhunter_auth import get_auth_router, get_current_active_user, require_admin

app = FastAPI()
app.include_router(get_auth_router())

@app.get("/api/protected")
async def protected_route(user: User = Depends(get_current_active_user)):
    return {"message": f"Hello {user.username}!"}

@app.get("/api/admin-only")
async def admin_route(user: User = Depends(require_admin)):
    return {"message": "Admin access granted"}
    """)