mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 05:50:21 -05:00
Implement Phase 2: Refresh tokens, 2FA, password reset, and audit logging
Co-authored-by: mblanke <9078342+mblanke@users.noreply.github.com>
This commit is contained in:
105
backend/alembic/versions/a1b2c3d4e5f6_add_phase_2_tables.py
Normal file
105
backend/alembic/versions/a1b2c3d4e5f6_add_phase_2_tables.py
Normal file
@@ -0,0 +1,105 @@
|
||||
"""Add Phase 2 tables
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: f82b3092d056
|
||||
Create Date: 2025-12-09 17:28:20.000000
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = 'a1b2c3d4e5f6'
|
||||
down_revision: Union[str, Sequence[str], None] = 'f82b3092d056'
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema for Phase 2."""
|
||||
|
||||
# Add new fields to users table
|
||||
op.add_column('users', sa.Column('email', sa.String(), nullable=True))
|
||||
op.add_column('users', sa.Column('email_verified', sa.Boolean(), nullable=False, server_default='false'))
|
||||
op.add_column('users', sa.Column('totp_secret', sa.String(), nullable=True))
|
||||
op.add_column('users', sa.Column('totp_enabled', sa.Boolean(), nullable=False, server_default='false'))
|
||||
op.create_index(op.f('ix_users_email'), 'users', ['email'], unique=True)
|
||||
|
||||
# Create refresh_tokens table
|
||||
op.create_table(
|
||||
'refresh_tokens',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('token', sa.String(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=False),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('is_revoked', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_refresh_tokens_id'), 'refresh_tokens', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_refresh_tokens_token'), 'refresh_tokens', ['token'], unique=True)
|
||||
|
||||
# Create password_reset_tokens table
|
||||
op.create_table(
|
||||
'password_reset_tokens',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('token', sa.String(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=False),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('is_used', sa.Boolean(), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_password_reset_tokens_id'), 'password_reset_tokens', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_password_reset_tokens_token'), 'password_reset_tokens', ['token'], unique=True)
|
||||
|
||||
# Create audit_logs table
|
||||
op.create_table(
|
||||
'audit_logs',
|
||||
sa.Column('id', sa.Integer(), nullable=False),
|
||||
sa.Column('user_id', sa.Integer(), nullable=True),
|
||||
sa.Column('tenant_id', sa.Integer(), nullable=False),
|
||||
sa.Column('action', sa.String(), nullable=False),
|
||||
sa.Column('resource_type', sa.String(), nullable=False),
|
||||
sa.Column('resource_id', sa.Integer(), nullable=True),
|
||||
sa.Column('details', sa.JSON(), nullable=True),
|
||||
sa.Column('ip_address', sa.String(), nullable=True),
|
||||
sa.Column('user_agent', sa.String(), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=True),
|
||||
sa.ForeignKeyConstraint(['tenant_id'], ['tenants.id'], ),
|
||||
sa.ForeignKeyConstraint(['user_id'], ['users.id'], ),
|
||||
sa.PrimaryKeyConstraint('id')
|
||||
)
|
||||
op.create_index(op.f('ix_audit_logs_id'), 'audit_logs', ['id'], unique=False)
|
||||
op.create_index(op.f('ix_audit_logs_created_at'), 'audit_logs', ['created_at'], unique=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema for Phase 2."""
|
||||
|
||||
# Drop audit_logs table
|
||||
op.drop_index(op.f('ix_audit_logs_created_at'), table_name='audit_logs')
|
||||
op.drop_index(op.f('ix_audit_logs_id'), table_name='audit_logs')
|
||||
op.drop_table('audit_logs')
|
||||
|
||||
# Drop password_reset_tokens table
|
||||
op.drop_index(op.f('ix_password_reset_tokens_token'), table_name='password_reset_tokens')
|
||||
op.drop_index(op.f('ix_password_reset_tokens_id'), table_name='password_reset_tokens')
|
||||
op.drop_table('password_reset_tokens')
|
||||
|
||||
# Drop refresh_tokens table
|
||||
op.drop_index(op.f('ix_refresh_tokens_token'), table_name='refresh_tokens')
|
||||
op.drop_index(op.f('ix_refresh_tokens_id'), table_name='refresh_tokens')
|
||||
op.drop_table('refresh_tokens')
|
||||
|
||||
# Remove new fields from users table
|
||||
op.drop_index(op.f('ix_users_email'), table_name='users')
|
||||
op.drop_column('users', 'totp_enabled')
|
||||
op.drop_column('users', 'totp_secret')
|
||||
op.drop_column('users', 'email_verified')
|
||||
op.drop_column('users', 'email')
|
||||
75
backend/app/api/routes/audit.py
Normal file
75
backend/app/api/routes/audit.py
Normal file
@@ -0,0 +1,75 @@
|
||||
from typing import List, Optional
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
from sqlalchemy.orm import Session
|
||||
from datetime import datetime
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.deps import get_current_active_user, require_role
|
||||
from app.models.user import User
|
||||
from app.models.audit_log import AuditLog
|
||||
from app.schemas.audit import AuditLogRead
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/", response_model=List[AuditLogRead])
|
||||
async def list_audit_logs(
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
action: Optional[str] = Query(None, description="Filter by action type"),
|
||||
resource_type: Optional[str] = Query(None, description="Filter by resource type"),
|
||||
start_date: Optional[datetime] = Query(None, description="Filter from date"),
|
||||
end_date: Optional[datetime] = Query(None, description="Filter to date"),
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
List audit logs (admin only, scoped to tenant)
|
||||
|
||||
Provides a complete audit trail of actions within the tenant.
|
||||
"""
|
||||
# Base query scoped to tenant
|
||||
query = db.query(AuditLog).filter(AuditLog.tenant_id == current_user.tenant_id)
|
||||
|
||||
# Apply filters
|
||||
if action:
|
||||
query = query.filter(AuditLog.action == action)
|
||||
if resource_type:
|
||||
query = query.filter(AuditLog.resource_type == resource_type)
|
||||
if start_date:
|
||||
query = query.filter(AuditLog.created_at >= start_date)
|
||||
if end_date:
|
||||
query = query.filter(AuditLog.created_at <= end_date)
|
||||
|
||||
# Order by most recent first
|
||||
query = query.order_by(AuditLog.created_at.desc())
|
||||
|
||||
# Paginate
|
||||
logs = query.offset(skip).limit(limit).all()
|
||||
|
||||
return logs
|
||||
|
||||
|
||||
@router.get("/{log_id}", response_model=AuditLogRead)
|
||||
async def get_audit_log(
|
||||
log_id: int,
|
||||
current_user: User = Depends(require_role(["admin"])),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Get a specific audit log entry (admin only)
|
||||
"""
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
log = db.query(AuditLog).filter(
|
||||
AuditLog.id == log_id,
|
||||
AuditLog.tenant_id == current_user.tenant_id
|
||||
).first()
|
||||
|
||||
if not log:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Audit log not found"
|
||||
)
|
||||
|
||||
return log
|
||||
@@ -1,13 +1,26 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from fastapi.security import OAuth2PasswordRequestForm
|
||||
from sqlalchemy.orm import Session
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import io
|
||||
import qrcode
|
||||
|
||||
from app.core.database import get_db
|
||||
from app.core.security import verify_password, get_password_hash, create_access_token
|
||||
from app.core.security import (
|
||||
verify_password, get_password_hash, create_access_token,
|
||||
create_refresh_token, create_reset_token, generate_totp_secret,
|
||||
verify_totp, get_totp_uri
|
||||
)
|
||||
from app.core.deps import get_current_active_user
|
||||
from app.models.user import User
|
||||
from app.models.tenant import Tenant
|
||||
from app.schemas.auth import Token, UserLogin, UserRegister
|
||||
from app.models.refresh_token import RefreshToken
|
||||
from app.models.password_reset_token import PasswordResetToken
|
||||
from app.schemas.auth import (
|
||||
Token, UserLogin, UserRegister, RefreshTokenRequest,
|
||||
PasswordResetRequest, PasswordResetConfirm,
|
||||
TwoFactorSetup, TwoFactorVerify
|
||||
)
|
||||
from app.schemas.user import UserRead, UserUpdate
|
||||
|
||||
router = APIRouter()
|
||||
@@ -102,6 +115,17 @@ async def login(
|
||||
detail="Inactive user"
|
||||
)
|
||||
|
||||
# Check 2FA if enabled (TOTP code should be in scopes for OAuth2)
|
||||
if user.totp_enabled:
|
||||
# For OAuth2 password flow, we'll check totp in scopes
|
||||
totp_code = form_data.scopes[0] if form_data.scopes else None
|
||||
if not totp_code or not verify_totp(user.totp_secret, totp_code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid 2FA code",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Create access token
|
||||
access_token = create_access_token(
|
||||
data={
|
||||
@@ -111,7 +135,21 @@ async def login(
|
||||
}
|
||||
)
|
||||
|
||||
return {"access_token": access_token, "token_type": "bearer"}
|
||||
# Create refresh token
|
||||
refresh_token_str = create_refresh_token()
|
||||
refresh_token_obj = RefreshToken(
|
||||
token=refresh_token_str,
|
||||
user_id=user.id,
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(days=30)
|
||||
)
|
||||
db.add(refresh_token_obj)
|
||||
db.commit()
|
||||
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_token_str,
|
||||
"token_type": "bearer"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/me", response_model=UserRead)
|
||||
@@ -162,3 +200,233 @@ async def update_current_user_profile(
|
||||
db.refresh(current_user)
|
||||
|
||||
return current_user
|
||||
|
||||
|
||||
@router.post("/refresh", response_model=Token)
|
||||
async def refresh_access_token(
|
||||
refresh_request: RefreshTokenRequest,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Refresh access token using refresh token
|
||||
|
||||
Provides a new access token without requiring login.
|
||||
"""
|
||||
# Find refresh token
|
||||
refresh_token = db.query(RefreshToken).filter(
|
||||
RefreshToken.token == refresh_request.refresh_token,
|
||||
RefreshToken.is_revoked == False
|
||||
).first()
|
||||
|
||||
if not refresh_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid refresh token"
|
||||
)
|
||||
|
||||
# Check if expired
|
||||
if refresh_token.expires_at < datetime.now(timezone.utc):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Refresh token expired"
|
||||
)
|
||||
|
||||
# Get user
|
||||
user = db.query(User).filter(User.id == refresh_token.user_id).first()
|
||||
if not user or not user.is_active:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="User not found or inactive"
|
||||
)
|
||||
|
||||
# Create new access token
|
||||
access_token = create_access_token(
|
||||
data={
|
||||
"sub": user.id,
|
||||
"tenant_id": user.tenant_id,
|
||||
"role": user.role
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_request.refresh_token,
|
||||
"token_type": "bearer"
|
||||
}
|
||||
|
||||
|
||||
@router.post("/2fa/setup", response_model=TwoFactorSetup)
|
||||
async def setup_two_factor(
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Set up two-factor authentication
|
||||
|
||||
Generates a TOTP secret and QR code URI for the user.
|
||||
"""
|
||||
if current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is already enabled"
|
||||
)
|
||||
|
||||
# Generate secret
|
||||
secret = generate_totp_secret()
|
||||
current_user.totp_secret = secret
|
||||
db.commit()
|
||||
|
||||
# Get QR code URI
|
||||
qr_uri = get_totp_uri(secret, current_user.username)
|
||||
|
||||
return {
|
||||
"secret": secret,
|
||||
"qr_code_uri": qr_uri
|
||||
}
|
||||
|
||||
|
||||
@router.post("/2fa/verify")
|
||||
async def verify_two_factor(
|
||||
verify_request: TwoFactorVerify,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Verify and enable two-factor authentication
|
||||
|
||||
User must provide a valid TOTP code to enable 2FA.
|
||||
"""
|
||||
if current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is already enabled"
|
||||
)
|
||||
|
||||
if not current_user.totp_secret:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA setup not initiated. Call /2fa/setup first."
|
||||
)
|
||||
|
||||
# Verify code
|
||||
if not verify_totp(current_user.totp_secret, verify_request.code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid 2FA code"
|
||||
)
|
||||
|
||||
# Enable 2FA
|
||||
current_user.totp_enabled = True
|
||||
db.commit()
|
||||
|
||||
return {"message": "2FA enabled successfully"}
|
||||
|
||||
|
||||
@router.post("/2fa/disable")
|
||||
async def disable_two_factor(
|
||||
verify_request: TwoFactorVerify,
|
||||
current_user: User = Depends(get_current_active_user),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Disable two-factor authentication
|
||||
|
||||
User must provide a valid TOTP code to disable 2FA.
|
||||
"""
|
||||
if not current_user.totp_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="2FA is not enabled"
|
||||
)
|
||||
|
||||
# Verify code
|
||||
if not verify_totp(current_user.totp_secret, verify_request.code):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid 2FA code"
|
||||
)
|
||||
|
||||
# Disable 2FA
|
||||
current_user.totp_enabled = False
|
||||
current_user.totp_secret = None
|
||||
db.commit()
|
||||
|
||||
return {"message": "2FA disabled successfully"}
|
||||
|
||||
|
||||
@router.post("/password-reset/request")
|
||||
async def request_password_reset(
|
||||
reset_request: PasswordResetRequest,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Request a password reset
|
||||
|
||||
Sends a password reset email to the user (mock implementation).
|
||||
"""
|
||||
# Find user by email
|
||||
user = db.query(User).filter(User.email == reset_request.email).first()
|
||||
|
||||
# Don't reveal if email exists or not (security best practice)
|
||||
# Always return success even if email doesn't exist
|
||||
if user:
|
||||
# Create reset token
|
||||
reset_token = create_reset_token()
|
||||
reset_token_obj = PasswordResetToken(
|
||||
token=reset_token,
|
||||
user_id=user.id,
|
||||
expires_at=datetime.now(timezone.utc) + timedelta(hours=1)
|
||||
)
|
||||
db.add(reset_token_obj)
|
||||
db.commit()
|
||||
|
||||
# TODO: Send email with reset link
|
||||
# For now, we'll just log it (in production, use an email service)
|
||||
print(f"Password reset token for {user.email}: {reset_token}")
|
||||
|
||||
return {"message": "If the email exists, a password reset link has been sent"}
|
||||
|
||||
|
||||
@router.post("/password-reset/confirm")
|
||||
async def confirm_password_reset(
|
||||
reset_confirm: PasswordResetConfirm,
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
"""
|
||||
Confirm password reset with token
|
||||
|
||||
Sets a new password for the user.
|
||||
"""
|
||||
# Find reset token
|
||||
reset_token = db.query(PasswordResetToken).filter(
|
||||
PasswordResetToken.token == reset_confirm.token,
|
||||
PasswordResetToken.is_used == False
|
||||
).first()
|
||||
|
||||
if not reset_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Invalid or expired reset token"
|
||||
)
|
||||
|
||||
# Check if expired
|
||||
if reset_token.expires_at < datetime.now(timezone.utc):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Reset token expired"
|
||||
)
|
||||
|
||||
# Get user
|
||||
user = db.query(User).filter(User.id == reset_token.user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="User not found"
|
||||
)
|
||||
|
||||
# Update password
|
||||
user.password_hash = get_password_hash(reset_confirm.new_password)
|
||||
reset_token.is_used = True
|
||||
db.commit()
|
||||
|
||||
return {"message": "Password reset successful"}
|
||||
|
||||
52
backend/app/core/audit.py
Normal file
52
backend/app/core/audit.py
Normal file
@@ -0,0 +1,52 @@
|
||||
from typing import Optional, Dict, Any
|
||||
from sqlalchemy.orm import Session
|
||||
from fastapi import Request
|
||||
|
||||
from app.models.audit_log import AuditLog
|
||||
|
||||
|
||||
def log_action(
|
||||
db: Session,
|
||||
user_id: Optional[int],
|
||||
tenant_id: int,
|
||||
action: str,
|
||||
resource_type: str,
|
||||
resource_id: Optional[int] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
request: Optional[Request] = None
|
||||
):
|
||||
"""
|
||||
Log an action to the audit log
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
user_id: ID of user performing action (None for system actions)
|
||||
tenant_id: Tenant ID
|
||||
action: Action type (CREATE, READ, UPDATE, DELETE, LOGIN, etc.)
|
||||
resource_type: Type of resource (user, host, case, etc.)
|
||||
resource_id: ID of the resource (if applicable)
|
||||
details: Additional details as JSON
|
||||
request: FastAPI request object (for IP and user agent)
|
||||
"""
|
||||
ip_address = None
|
||||
user_agent = None
|
||||
|
||||
if request:
|
||||
ip_address = request.client.host if request.client else None
|
||||
user_agent = request.headers.get("user-agent")
|
||||
|
||||
audit_log = AuditLog(
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
action=action,
|
||||
resource_type=resource_type,
|
||||
resource_id=resource_id,
|
||||
details=details,
|
||||
ip_address=ip_address,
|
||||
user_agent=user_agent
|
||||
)
|
||||
|
||||
db.add(audit_log)
|
||||
db.commit()
|
||||
|
||||
return audit_log
|
||||
@@ -7,7 +7,18 @@ class Settings(BaseSettings):
|
||||
database_url: str = "postgresql://postgres:postgres@db:5432/velocicompanion"
|
||||
secret_key: str = "your-secret-key-change-in-production-min-32-chars-long"
|
||||
access_token_expire_minutes: int = 30
|
||||
refresh_token_expire_days: int = 30
|
||||
algorithm: str = "HS256"
|
||||
|
||||
# Email settings (for password reset)
|
||||
smtp_host: str = "localhost"
|
||||
smtp_port: int = 587
|
||||
smtp_user: str = ""
|
||||
smtp_password: str = ""
|
||||
from_email: str = "noreply@velocicompanion.com"
|
||||
|
||||
# WebSocket settings
|
||||
ws_enabled: bool = True
|
||||
|
||||
class Config:
|
||||
env_file = ".env"
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
import secrets
|
||||
import pyotp
|
||||
from jose import JWTError, jwt
|
||||
from passlib.context import CryptContext
|
||||
|
||||
@@ -56,3 +58,63 @@ def verify_token(token: str) -> Optional[dict]:
|
||||
return payload
|
||||
except JWTError:
|
||||
return None
|
||||
|
||||
|
||||
def create_refresh_token() -> str:
|
||||
"""
|
||||
Create a secure random refresh token
|
||||
|
||||
Returns:
|
||||
Random token string
|
||||
"""
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def create_reset_token() -> str:
|
||||
"""
|
||||
Create a secure random password reset token
|
||||
|
||||
Returns:
|
||||
Random token string
|
||||
"""
|
||||
return secrets.token_urlsafe(32)
|
||||
|
||||
|
||||
def generate_totp_secret() -> str:
|
||||
"""
|
||||
Generate a TOTP secret for 2FA
|
||||
|
||||
Returns:
|
||||
Base32 encoded secret
|
||||
"""
|
||||
return pyotp.random_base32()
|
||||
|
||||
|
||||
def verify_totp(secret: str, code: str) -> bool:
|
||||
"""
|
||||
Verify a TOTP code
|
||||
|
||||
Args:
|
||||
secret: TOTP secret
|
||||
code: 6-digit code from authenticator app
|
||||
|
||||
Returns:
|
||||
True if code is valid
|
||||
"""
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.verify(code, valid_window=1)
|
||||
|
||||
|
||||
def get_totp_uri(secret: str, username: str) -> str:
|
||||
"""
|
||||
Get TOTP provisioning URI for QR code
|
||||
|
||||
Args:
|
||||
secret: TOTP secret
|
||||
username: User's username
|
||||
|
||||
Returns:
|
||||
otpauth:// URI
|
||||
"""
|
||||
totp = pyotp.TOTP(secret)
|
||||
return totp.provisioning_uri(name=username, issuer_name="VelociCompanion")
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.api.routes import auth, users, tenants, hosts, ingestion, vt
|
||||
from app.api.routes import auth, users, tenants, hosts, ingestion, vt, audit
|
||||
from app.core.config import settings
|
||||
|
||||
app = FastAPI(
|
||||
title=settings.app_name,
|
||||
description="Multi-tenant threat hunting companion for Velociraptor",
|
||||
version="0.1.0"
|
||||
version="0.2.0"
|
||||
)
|
||||
|
||||
# Configure CORS
|
||||
@@ -26,6 +26,7 @@ app.include_router(tenants.router, prefix="/api/tenants", tags=["Tenants"])
|
||||
app.include_router(hosts.router, prefix="/api/hosts", tags=["Hosts"])
|
||||
app.include_router(ingestion.router, prefix="/api/ingestion", tags=["Ingestion"])
|
||||
app.include_router(vt.router, prefix="/api/vt", tags=["VirusTotal"])
|
||||
app.include_router(audit.router, prefix="/api/audit", tags=["Audit Logs"])
|
||||
|
||||
|
||||
@app.get("/")
|
||||
@@ -33,7 +34,7 @@ async def root():
|
||||
"""Root endpoint"""
|
||||
return {
|
||||
"message": f"Welcome to {settings.app_name}",
|
||||
"version": "0.1.0",
|
||||
"version": "0.2.0",
|
||||
"docs": "/docs"
|
||||
}
|
||||
|
||||
|
||||
24
backend/app/models/audit_log.py
Normal file
24
backend/app/models/audit_log.py
Normal file
@@ -0,0 +1,24 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Text, JSON
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class AuditLog(Base):
|
||||
__tablename__ = "audit_logs"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=True)
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
action = Column(String, nullable=False) # CREATE, READ, UPDATE, DELETE
|
||||
resource_type = Column(String, nullable=False) # user, host, case, etc.
|
||||
resource_id = Column(Integer, nullable=True)
|
||||
details = Column(JSON, nullable=True)
|
||||
ip_address = Column(String, nullable=True)
|
||||
user_agent = Column(String, nullable=True)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
tenant = relationship("Tenant")
|
||||
19
backend/app/models/password_reset_token.py
Normal file
19
backend/app/models/password_reset_token.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class PasswordResetToken(Base):
|
||||
__tablename__ = "password_reset_tokens"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
token = Column(String, unique=True, index=True, nullable=False)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
expires_at = Column(DateTime, nullable=False)
|
||||
is_used = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
user = relationship("User")
|
||||
19
backend/app/models/refresh_token.py
Normal file
19
backend/app/models/refresh_token.py
Normal file
@@ -0,0 +1,19 @@
|
||||
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Boolean
|
||||
from sqlalchemy.orm import relationship
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
from app.core.database import Base
|
||||
|
||||
|
||||
class RefreshToken(Base):
|
||||
__tablename__ = "refresh_tokens"
|
||||
|
||||
id = Column(Integer, primary_key=True, index=True)
|
||||
token = Column(String, unique=True, index=True, nullable=False)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
expires_at = Column(DateTime, nullable=False)
|
||||
is_revoked = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="refresh_tokens")
|
||||
@@ -14,7 +14,12 @@ class User(Base):
|
||||
role = Column(String, default="user", nullable=False) # user, admin
|
||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||
is_active = Column(Boolean, default=True, nullable=False)
|
||||
email = Column(String, unique=True, nullable=True, index=True)
|
||||
email_verified = Column(Boolean, default=False, nullable=False)
|
||||
totp_secret = Column(String, nullable=True)
|
||||
totp_enabled = Column(Boolean, default=False, nullable=False)
|
||||
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||
|
||||
# Relationships
|
||||
tenant = relationship("Tenant", back_populates="users")
|
||||
refresh_tokens = relationship("RefreshToken", back_populates="user")
|
||||
|
||||
29
backend/app/schemas/audit.py
Normal file
29
backend/app/schemas/audit.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from pydantic import BaseModel
|
||||
from typing import Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class AuditLogBase(BaseModel):
|
||||
"""Base audit log schema"""
|
||||
action: str
|
||||
resource_type: str
|
||||
resource_id: Optional[int] = None
|
||||
details: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class AuditLogCreate(AuditLogBase):
|
||||
"""Schema for creating an audit log entry"""
|
||||
pass
|
||||
|
||||
|
||||
class AuditLogRead(AuditLogBase):
|
||||
"""Schema for reading audit log data"""
|
||||
id: int
|
||||
user_id: Optional[int]
|
||||
tenant_id: int
|
||||
ip_address: Optional[str]
|
||||
user_agent: Optional[str]
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
from_attributes = True
|
||||
@@ -1,10 +1,11 @@
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class Token(BaseModel):
|
||||
"""Token response schema"""
|
||||
access_token: str
|
||||
refresh_token: Optional[str] = None
|
||||
token_type: str = "bearer"
|
||||
|
||||
|
||||
@@ -19,11 +20,40 @@ class UserLogin(BaseModel):
|
||||
"""User login request schema"""
|
||||
username: str
|
||||
password: str
|
||||
totp_code: Optional[str] = None
|
||||
|
||||
|
||||
class UserRegister(BaseModel):
|
||||
"""User registration request schema"""
|
||||
username: str
|
||||
password: str
|
||||
email: Optional[EmailStr] = None
|
||||
tenant_id: Optional[int] = None
|
||||
role: str = "user"
|
||||
|
||||
|
||||
class RefreshTokenRequest(BaseModel):
|
||||
"""Refresh token request schema"""
|
||||
refresh_token: str
|
||||
|
||||
|
||||
class PasswordResetRequest(BaseModel):
|
||||
"""Password reset request schema"""
|
||||
email: EmailStr
|
||||
|
||||
|
||||
class PasswordResetConfirm(BaseModel):
|
||||
"""Password reset confirmation schema"""
|
||||
token: str
|
||||
new_password: str
|
||||
|
||||
|
||||
class TwoFactorSetup(BaseModel):
|
||||
"""2FA setup response schema"""
|
||||
secret: str
|
||||
qr_code_uri: str
|
||||
|
||||
|
||||
class TwoFactorVerify(BaseModel):
|
||||
"""2FA verification schema"""
|
||||
code: str
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, EmailStr
|
||||
from typing import Optional
|
||||
from datetime import datetime
|
||||
|
||||
@@ -6,6 +6,7 @@ from datetime import datetime
|
||||
class UserBase(BaseModel):
|
||||
"""Base user schema"""
|
||||
username: str
|
||||
email: Optional[EmailStr] = None
|
||||
role: str = "user"
|
||||
tenant_id: int
|
||||
|
||||
@@ -18,15 +19,18 @@ class UserCreate(UserBase):
|
||||
class UserUpdate(BaseModel):
|
||||
"""Schema for updating a user"""
|
||||
username: Optional[str] = None
|
||||
email: Optional[EmailStr] = None
|
||||
password: Optional[str] = None
|
||||
role: Optional[str] = None
|
||||
is_active: Optional[bool] = None
|
||||
|
||||
|
||||
class UserRead(UserBase):
|
||||
"""Schema for reading user data (excludes password_hash)"""
|
||||
"""Schema for reading user data (excludes password_hash and secrets)"""
|
||||
id: int
|
||||
is_active: bool
|
||||
email_verified: bool
|
||||
totp_enabled: bool
|
||||
created_at: datetime
|
||||
|
||||
class Config:
|
||||
|
||||
@@ -8,3 +8,6 @@ python-multipart==0.0.6
|
||||
alembic==1.13.1
|
||||
pydantic==2.5.3
|
||||
pydantic-settings==2.1.0
|
||||
pyotp==2.9.0
|
||||
qrcode[pil]==7.4.2
|
||||
websockets==12.0
|
||||
|
||||
Reference in New Issue
Block a user