schemas.py

from pydantic import BaseModel, EmailStr, validator
from typing import Optional, List, Dict, Any
from datetime import datetime
from enum import Enum

class UserRole(str, Enum):
    ADMIN = "admin"
    USER = "user"
    VIEWER = "viewer"

class ScanStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"

# Auth schemas
class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None
    user_id: Optional[int] = None
    org_uuid: Optional[str] = None
    role: Optional[str] = None

class LoginRequest(BaseModel):
    username: str
    password: str

class UserRegistrationRequest(BaseModel):
    email: EmailStr
    username: str
    password: str
    confirm_password: str
    first_name: Optional[str] = None
    last_name: Optional[str] = None
    organization_name: Optional[str] = None
    
    @validator('username')
    def validate_username(cls, v):
        if len(v) < 3:
            raise ValueError('Username must be at least 3 characters long')
        if len(v) > 50:
            raise ValueError('Username must be less than 50 characters')
        if not v.replace('_', '').replace('-', '').isalnum():
            raise ValueError('Username can only contain letters, numbers, hyphens, and underscores')
        return v.lower()
    
    @validator('password')
    def validate_password(cls, v):
        if len(v) < 12:
            raise ValueError('Password must be at least 12 characters long')
        if len(v) > 128:
            raise ValueError('Password must be less than 128 characters')
        
        # Check for complexity
        has_upper = any(c.isupper() for c in v)
        has_lower = any(c.islower() for c in v)
        has_digit = any(c.isdigit() for c in v)
        has_special = any(c in '!@#$%^&*()_+-=[]{}|;:,.<>?' for c in v)
        
        if not all([has_upper, has_lower, has_digit, has_special]):
            raise ValueError('Password must contain at least one uppercase letter, one lowercase letter, one digit, and one special character')
        
        # Check for common patterns
        common_patterns = ['password', '123456', 'qwerty', 'admin', 'letmein']
        if any(pattern in v.lower() for pattern in common_patterns):
            raise ValueError('Password contains common patterns and is not secure')
            
        return v
    
    @validator('confirm_password')
    def passwords_match(cls, v, values, **kwargs):
        if 'password' in values and v != values['password']:
            raise ValueError('Passwords do not match')
        return v

class UserRegistrationResponse(BaseModel):
    id: int
    uuid: str
    email: EmailStr
    username: str
    first_name: Optional[str] = None
    last_name: Optional[str] = None
    is_active: bool
    email_verified: bool
    created_at: datetime
    message: str

# User schemas
class UserBase(BaseModel):
    email: EmailStr
    username: Optional[str] = None
    first_name: Optional[str] = None
    last_name: Optional[str] = None

class User(UserBase):
    id: int
    uuid: str
    is_active: bool
    is_superuser: bool
    email_verified: bool
    last_login: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True

class UserInDB(User):
    password_hash: str

# Organization schemas
class OrganizationBase(BaseModel):
    name: str
    slug: str
    description: Optional[str] = None
    website: Optional[str] = None
    contact_email: Optional[EmailStr] = None

class Organization(OrganizationBase):
    id: int
    uuid: str
    is_active: bool
    created_at: datetime
    updated_at: datetime
    created_by: Optional[int] = None

    class Config:
        from_attributes = True

class UserOrganization(BaseModel):
    user_id: int
    organization_id: int
    role_id: int
    is_active: bool
    invited_by: Optional[int] = None
    joined_at: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True

# Scan schemas
class ScanRequest(BaseModel):
    target: str
    protocol_filter: Optional[str] = 'sui'
    scan_options: Optional[Dict[str, Any]] = {}

class ScanSession(BaseModel):
    session_id: str
    organization_uuid: str
    target: str
    protocol_filter: str
    status: ScanStatus
    started_at: datetime
    completed_at: Optional[datetime] = None
    results: Optional[Dict[str, Any]] = None
    error_message: Optional[str] = None

class Report(BaseModel):
    id: Optional[int] = None
    session_id: str
    organization_uuid: str
    report_type: str
    title: str
    summary: str
    findings: Dict[str, Any]
    risk_score: Optional[float] = None
    created_at: Optional[datetime] = None

    class Config:
        from_attributes = True

# Pipeline schemas
class PipelineRequest(BaseModel):
    organization_uuid: str
    protocol_filter: Optional[str] = 'sui'
    scan_options: Optional[Dict[str, Any]] = {}

class PipelineStatus(BaseModel):
    pipeline_id: str
    organization_uuid: str
    status: ScanStatus
    started_at: datetime
    progress: Optional[Dict[str, Any]] = None
    estimated_completion: Optional[datetime] = None

# User schemas
class UserBase(BaseModel):
    email: EmailStr
    username: Optional[str] = None
    first_name: Optional[str] = None
    last_name: Optional[str] = None

class User(UserBase):
    id: int
    uuid: str
    is_active: bool
    is_superuser: bool
    email_verified: bool
    last_login: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True

class UserInDB(User):
    password_hash: str

# Organization schemas
class OrganizationBase(BaseModel):
    name: str
    slug: str
    description: Optional[str] = None
    website: Optional[str] = None
    contact_email: Optional[EmailStr] = None

class Organization(OrganizationBase):
    id: int
    uuid: str
    is_active: bool
    created_at: datetime
    updated_at: datetime
    created_by: Optional[int] = None

    class Config:
        from_attributes = True

class UserOrganization(BaseModel):
    user_id: int
    organization_id: int
    role_id: int
    is_active: bool
    invited_by: Optional[int] = None
    joined_at: Optional[datetime] = None
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True

# Scan schemas
class ScanRequest(BaseModel):
    target: str
    protocol_filter: Optional[str] = 'sui'
    scan_options: Optional[Dict[str, Any]] = {}

class ScanSession(BaseModel):
    session_id: str
    organization_uuid: str
    target: str
    protocol_filter: str
    status: ScanStatus
    started_at: datetime
    completed_at: Optional[datetime] = None
    results: Optional[Dict[str, Any]] = None
    error_message: Optional[str] = None

class Report(BaseModel):
    id: Optional[int] = None
    session_id: str
    organization_uuid: str
    report_type: str
    title: str
    summary: str
    findings: Dict[str, Any]
    risk_score: Optional[float] = None
    created_at: Optional[datetime] = None

    class Config:
        from_attributes = True

# Pipeline schemas
class PipelineRequest(BaseModel):
    organization_uuid: str
    protocol_filter: Optional[str] = 'sui'
    scan_options: Optional[Dict[str, Any]] = {}

class PipelineStatus(BaseModel):
    pipeline_id: str
    organization_uuid: str
    status: ScanStatus
    started_at: datetime
    progress: Optional[Dict[str, Any]] = None
    estimated_completion: Optional[datetime] = None