mirror of
https://github.com/mblanke/ThreatHunt.git
synced 2026-03-01 14:00:20 -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 import APIRouter, Depends, HTTPException, status
|
||||||
from fastapi.security import OAuth2PasswordRequestForm
|
from fastapi.security import OAuth2PasswordRequestForm
|
||||||
from sqlalchemy.orm import Session
|
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.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.core.deps import get_current_active_user
|
||||||
from app.models.user import User
|
from app.models.user import User
|
||||||
from app.models.tenant import Tenant
|
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
|
from app.schemas.user import UserRead, UserUpdate
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
@@ -102,6 +115,17 @@ async def login(
|
|||||||
detail="Inactive user"
|
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
|
# Create access token
|
||||||
access_token = create_access_token(
|
access_token = create_access_token(
|
||||||
data={
|
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)
|
@router.get("/me", response_model=UserRead)
|
||||||
@@ -162,3 +200,233 @@ async def update_current_user_profile(
|
|||||||
db.refresh(current_user)
|
db.refresh(current_user)
|
||||||
|
|
||||||
return 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"
|
database_url: str = "postgresql://postgres:postgres@db:5432/velocicompanion"
|
||||||
secret_key: str = "your-secret-key-change-in-production-min-32-chars-long"
|
secret_key: str = "your-secret-key-change-in-production-min-32-chars-long"
|
||||||
access_token_expire_minutes: int = 30
|
access_token_expire_minutes: int = 30
|
||||||
|
refresh_token_expire_days: int = 30
|
||||||
algorithm: str = "HS256"
|
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:
|
class Config:
|
||||||
env_file = ".env"
|
env_file = ".env"
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
from datetime import datetime, timedelta, timezone
|
from datetime import datetime, timedelta, timezone
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
import secrets
|
||||||
|
import pyotp
|
||||||
from jose import JWTError, jwt
|
from jose import JWTError, jwt
|
||||||
from passlib.context import CryptContext
|
from passlib.context import CryptContext
|
||||||
|
|
||||||
@@ -56,3 +58,63 @@ def verify_token(token: str) -> Optional[dict]:
|
|||||||
return payload
|
return payload
|
||||||
except JWTError:
|
except JWTError:
|
||||||
return None
|
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 import FastAPI
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
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
|
from app.core.config import settings
|
||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title=settings.app_name,
|
title=settings.app_name,
|
||||||
description="Multi-tenant threat hunting companion for Velociraptor",
|
description="Multi-tenant threat hunting companion for Velociraptor",
|
||||||
version="0.1.0"
|
version="0.2.0"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Configure CORS
|
# 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(hosts.router, prefix="/api/hosts", tags=["Hosts"])
|
||||||
app.include_router(ingestion.router, prefix="/api/ingestion", tags=["Ingestion"])
|
app.include_router(ingestion.router, prefix="/api/ingestion", tags=["Ingestion"])
|
||||||
app.include_router(vt.router, prefix="/api/vt", tags=["VirusTotal"])
|
app.include_router(vt.router, prefix="/api/vt", tags=["VirusTotal"])
|
||||||
|
app.include_router(audit.router, prefix="/api/audit", tags=["Audit Logs"])
|
||||||
|
|
||||||
|
|
||||||
@app.get("/")
|
@app.get("/")
|
||||||
@@ -33,7 +34,7 @@ async def root():
|
|||||||
"""Root endpoint"""
|
"""Root endpoint"""
|
||||||
return {
|
return {
|
||||||
"message": f"Welcome to {settings.app_name}",
|
"message": f"Welcome to {settings.app_name}",
|
||||||
"version": "0.1.0",
|
"version": "0.2.0",
|
||||||
"docs": "/docs"
|
"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
|
role = Column(String, default="user", nullable=False) # user, admin
|
||||||
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
tenant_id = Column(Integer, ForeignKey("tenants.id"), nullable=False)
|
||||||
is_active = Column(Boolean, default=True, 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))
|
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
|
||||||
|
|
||||||
# Relationships
|
# Relationships
|
||||||
tenant = relationship("Tenant", back_populates="users")
|
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
|
from typing import Optional
|
||||||
|
|
||||||
|
|
||||||
class Token(BaseModel):
|
class Token(BaseModel):
|
||||||
"""Token response schema"""
|
"""Token response schema"""
|
||||||
access_token: str
|
access_token: str
|
||||||
|
refresh_token: Optional[str] = None
|
||||||
token_type: str = "bearer"
|
token_type: str = "bearer"
|
||||||
|
|
||||||
|
|
||||||
@@ -19,11 +20,40 @@ class UserLogin(BaseModel):
|
|||||||
"""User login request schema"""
|
"""User login request schema"""
|
||||||
username: str
|
username: str
|
||||||
password: str
|
password: str
|
||||||
|
totp_code: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
class UserRegister(BaseModel):
|
class UserRegister(BaseModel):
|
||||||
"""User registration request schema"""
|
"""User registration request schema"""
|
||||||
username: str
|
username: str
|
||||||
password: str
|
password: str
|
||||||
|
email: Optional[EmailStr] = None
|
||||||
tenant_id: Optional[int] = None
|
tenant_id: Optional[int] = None
|
||||||
role: str = "user"
|
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 typing import Optional
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
@@ -6,6 +6,7 @@ from datetime import datetime
|
|||||||
class UserBase(BaseModel):
|
class UserBase(BaseModel):
|
||||||
"""Base user schema"""
|
"""Base user schema"""
|
||||||
username: str
|
username: str
|
||||||
|
email: Optional[EmailStr] = None
|
||||||
role: str = "user"
|
role: str = "user"
|
||||||
tenant_id: int
|
tenant_id: int
|
||||||
|
|
||||||
@@ -18,15 +19,18 @@ class UserCreate(UserBase):
|
|||||||
class UserUpdate(BaseModel):
|
class UserUpdate(BaseModel):
|
||||||
"""Schema for updating a user"""
|
"""Schema for updating a user"""
|
||||||
username: Optional[str] = None
|
username: Optional[str] = None
|
||||||
|
email: Optional[EmailStr] = None
|
||||||
password: Optional[str] = None
|
password: Optional[str] = None
|
||||||
role: Optional[str] = None
|
role: Optional[str] = None
|
||||||
is_active: Optional[bool] = None
|
is_active: Optional[bool] = None
|
||||||
|
|
||||||
|
|
||||||
class UserRead(UserBase):
|
class UserRead(UserBase):
|
||||||
"""Schema for reading user data (excludes password_hash)"""
|
"""Schema for reading user data (excludes password_hash and secrets)"""
|
||||||
id: int
|
id: int
|
||||||
is_active: bool
|
is_active: bool
|
||||||
|
email_verified: bool
|
||||||
|
totp_enabled: bool
|
||||||
created_at: datetime
|
created_at: datetime
|
||||||
|
|
||||||
class Config:
|
class Config:
|
||||||
|
|||||||
@@ -8,3 +8,6 @@ python-multipart==0.0.6
|
|||||||
alembic==1.13.1
|
alembic==1.13.1
|
||||||
pydantic==2.5.3
|
pydantic==2.5.3
|
||||||
pydantic-settings==2.1.0
|
pydantic-settings==2.1.0
|
||||||
|
pyotp==2.9.0
|
||||||
|
qrcode[pil]==7.4.2
|
||||||
|
websockets==12.0
|
||||||
|
|||||||
Reference in New Issue
Block a user