Mock API Server for Testing Lumos Connectors

import time
import uuid
from datetime import datetime, timedelta
from enum import Enum
from typing import Dict, List, Optional, Set, Any
from fastapi import FastAPI, Form, HTTPException, Query, Depends, Header, Response, status
from pydantic import BaseModel, Field, EmailStr

app = FastAPI(title="Mock API for Lumos Connector Testing")

class UserStatus(str, Enum):
    ACTIVE = "active"
    INACTIVE = "inactive"
    DELETED = "deleted"

class ResourceType(str, Enum):
    PROJECT = "project"
    TEAM = "team"
    GLOBAL = "global"

class EntitlementType(str, Enum):
    PROJECT_ROLE = "project_role"
    TEAM_MEMBER = "team_member"
    ADMIN_ROLE = "admin_role"

class User(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    email: EmailStr
    name: str
    status: UserStatus = UserStatus.ACTIVE
    created_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    last_login: Optional[str] = None

class Resource(BaseModel):
    id: str
    name: str
    type: ResourceType
    description: Optional[str] = None

class Entitlement(BaseModel):
    id: str
    name: str
    type: EntitlementType
    resource_id: str
    description: Optional[str] = None

class Association(BaseModel):
    user_id: str
    entitlement_id: str
    resource_id: str
    assigned_at: str = Field(default_factory=lambda: datetime.now().isoformat())

class UserCreate(BaseModel):
    email: EmailStr
    name: str
    status: UserStatus = UserStatus.ACTIVE

class UserUpdate(BaseModel):
    email: Optional[EmailStr] = None
    name: Optional[str] = None
    status: Optional[UserStatus] = None

class EntitlementAssign(BaseModel):
    entitlement_id: str
    resource_id: str

class PaginatedResponse(BaseModel):
    items: List[Any]
    total: int
    has_more: bool

db_users: Dict[str, User] = {}

db_resources: Dict[str, Resource] = {
    "project-1": Resource(id="project-1", name="Marketing Project", type=ResourceType.PROJECT,
                          description="Marketing campaigns and assets"),
    "project-2": Resource(id="project-2", name="Engineering Project", type=ResourceType.PROJECT,
                          description="Product development"),
    "team-1": Resource(id="team-1", name="Marketing Team", type=ResourceType.TEAM,
                       description="Marketing department"),
    "team-2": Resource(id="team-2", name="Engineering Team", type=ResourceType.TEAM,
                       description="Engineering department"),
    "global": Resource(id="global", name="Global", type=ResourceType.GLOBAL,
                      description="Global resources")
}

db_entitlements: Dict[str, Entitlement] = {
    "project-admin": Entitlement(id="project-admin", name="Project Admin",
                               type=EntitlementType.PROJECT_ROLE,
                               resource_id="project-1",
                               description="Admin access to projects"),
    "project-viewer": Entitlement(id="project-viewer", name="Project Viewer",
                                type=EntitlementType.PROJECT_ROLE,
                                resource_id="project-1",
                                description="View-only access to projects"),
    "team-member": Entitlement(id="team-member", name="Team Member",
                             type=EntitlementType.TEAM_MEMBER,
                             resource_id="team-1",
                             description="Member of a team"),
    "team-lead": Entitlement(id="team-lead", name="Team Lead",
                           type=EntitlementType.TEAM_MEMBER,
                           resource_id="team-1",
                           description="Team leadership role"),
    "system-admin": Entitlement(id="system-admin", name="System Administrator",
                              type=EntitlementType.ADMIN_ROLE,
                              resource_id="global",
                              description="Global administrator"),
}

db_associations: List[Association] = []

api_tokens = {
    "valid-token": "mock-tenant-id"
}

def verify_token(authorization: str = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing token",
        )

    token = authorization.replace("Bearer ", "")
    if token not in api_tokens:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid token",
        )

    return token

def create_sample_users(count: int = 10):
    """Create sample users if none exist"""
    if not db_users:
        domains = ["example.com", "test.org", "mockdata.net"]
        departments = ["Marketing", "Engineering", "Sales", "HR", "Finance"]

        for i in range(1, count + 1):
            domain = domains[i % len(domains)]
            dept = departments[i % len(departments)]

            user_id = f"user-{i}"
            user = User(
                id=user_id,
                email=f"user{i}@{domain}",
                name=f"{dept} User {i}",
                status=UserStatus.ACTIVE,
                created_at=(datetime.now() - timedelta(days=i*10)).isoformat()
            )

            if i % 3 != 0:
                days_ago = i % 30
                user.last_login = (datetime.now() - timedelta(days=days_ago)).isoformat()

            db_users[user_id] = user

            if i <= 3:
                db_associations.append(Association(
                    user_id=user_id,
                    entitlement_id="system-admin",
                    resource_id="global"
                ))

            if dept == "Marketing":
                db_associations.append(Association(
                    user_id=user_id,
                    entitlement_id="project-admin" if i % 5 == 0 else "project-viewer",
                    resource_id="project-1"
                ))

            team_id = "team-1" if dept in ["Marketing", "Sales"] else "team-2"
            entitlement_id = "team-lead" if i % 7 == 0 else "team-member"
            db_associations.append(Association(
                user_id=user_id,
                entitlement_id=entitlement_id,
                resource_id=team_id
            ))

@app.get("/")
async def root():
    return {"status": "healthy", "service": "Mock API Server for Lumos Connector Testing"}

@app.get("/api/users", response_model=PaginatedResponse)
async def list_users(
    limit: int = Query(10, ge=1, le=100),
    offset: int = Query(0, ge=0),
    token: str = Depends(verify_token)
):
    create_sample_users()
    users = list(db_users.values())
    total = len(users)
    paginated_users = users[offset:offset+limit]
    time.sleep(0.1)
    return {
        "items": paginated_users,
        "total": total,
        "has_more": offset + limit < total
    }

@app.get("/api/users/USER_ID", response_model=User)
async def get_user(user_id: str, token: str = Depends(verify_token)):
    create_sample_users()
    if user_id not in db_users:
        raise HTTPException(status_code=404, detail="User not found")
    return db_users[user_id]

@app.post("/api/users", response_model=User, status_code=status.HTTP_201_CREATED)
async def create_user(user: UserCreate, token: str = Depends(verify_token)):
    user_id = str(uuid.uuid4())
    new_user = User(id=user_id, email=user.email, name=user.name, status=user.status)
    db_users[user_id] = new_user
    return new_user

@app.put("/api/users/USER_ID", response_model=User)
async def update_user(user_id: str, user_update: UserUpdate, token: str = Depends(verify_token)):
    create_sample_users()
    if user_id not in db_users:
        raise HTTPException(status_code=404, detail="User not found")
    current_user = db_users[user_id]
    if user_update.email is not None:
        current_user.email = user_update.email
    if user_update.name is not None:
        current_user.name = user_update.name
    if user_update.status is not None:
        current_user.status = user_update.status
    db_users[user_id] = current_user
    return current_user

@app.delete("/api/users/USER_ID", status_code=status.HTTP_204_NO_CONTENT)
async def delete_user(user_id: str, token: str = Depends(verify_token)):
    create_sample_users()
    if user_id not in db_users:
        raise HTTPException(status_code=404, detail="User not found")
    db_users[user_id].status = UserStatus.DELETED
    return Response(status_code=status.HTTP_204_NO_CONTENT)

@app.get("/api/resources", response_model=List[Resource])
async def list_resources(token: str = Depends(verify_token)):
    return list(db_resources.values())

@app.get("/api/entitlements", response_model=List[Entitlement])
async def list_entitlements(token: str = Depends(verify_token)):
    return list(db_entitlements.values())

@app.get("/api/associations", response_model=List[Association])
async def list_associations(
    user_id: Optional[str] = None,
    entitlement_id: Optional[str] = None,
    resource_id: Optional[str] = None,
    token: str = Depends(verify_token)
):
    create_sample_users()
    result = db_associations
    if user_id:
        result = [a for a in result if a.user_id == user_id]
    if entitlement_id:
        result = [a for a in result if a.entitlement_id == entitlement_id]
    if resource_id:
        result = [a for a in result if a.resource_id == resource_id]
    return result

@app.post("/api/users/USER_ID/entitlements", status_code=status.HTTP_201_CREATED)
async def assign_entitlement(
    user_id: str,
    assignment: EntitlementAssign,
    token: str = Depends(verify_token)
):
    create_sample_users()
    if user_id not in db_users:
        raise HTTPException(status_code=404, detail="User not found")
    if assignment.entitlement_id not in db_entitlements:
        raise HTTPException(status_code=404, detail="Entitlement not found")
    if assignment.resource_id not in db_resources:
        raise HTTPException(status_code=404, detail="Resource not found")
    for assoc in db_associations:
        if (assoc.user_id == user_id and
            assoc.entitlement_id == assignment.entitlement_id and
            assoc.resource_id == assignment.resource_id):
            return {"message": "Assignment already exists"}
    new_assoc = Association(
        user_id=user_id,
        entitlement_id=assignment.entitlement_id,
        resource_id=assignment.resource_id
    )
    db_associations.append(new_assoc)
    return {"message": "Entitlement assigned successfully"}

@app.delete("/api/users/USER_ID/entitlements/{entitlement_id}", status_code=status.HTTP_204_NO_CONTENT)
async def unassign_entitlement(
    user_id: str,
    entitlement_id: str,
    resource_id: str = Query(...),
    token: str = Depends(verify_token)
):
    create_sample_users()
    if user_id not in db_users:
        raise HTTPException(status_code=404, detail="User not found")
    for i, assoc in enumerate(db_associations):
        if (assoc.user_id == user_id and
            assoc.entitlement_id == entitlement_id and
            assoc.resource_id == resource_id):
            db_associations.pop(i)
            return Response(status_code=status.HTTP_204_NO_CONTENT)
    raise HTTPException(status_code=404, detail="Association not found")

@app.get("/api/health")
async def health_check():
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "users_count": len(db_users),
        "resources_count": len(db_resources),
        "entitlements_count": len(db_entitlements),
        "associations_count": len(db_associations)
    }

class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "Bearer"
    expires_in: int = 3600
    refresh_token: Optional[str] = None


@app.post("/oauth/token", response_model=TokenResponse)
async def get_token(
    grant_type: str = Form(...),
    client_id: str = Form(...),
    client_secret: Optional[str] = Form(None),
    code: Optional[str] = Form(None),
    redirect_uri: Optional[str] = Form(None),
    refresh_token: Optional[str] = Form(None),
):
    """Simulate an OAuth token endpoint.

    Fields are read as form data, which is what RFC 6749 specifies and what the
    Lumos SDK's OAuth module sends by default.
    """
    if client_id != "test-client-id":
        raise HTTPException(status_code=400, detail="Invalid client_id")

    if grant_type == "authorization_code":
        if code != "valid-auth-code":
            raise HTTPException(status_code=400, detail="Invalid authorization code")
        return TokenResponse(access_token="valid-token", refresh_token="valid-refresh-token")

    if grant_type == "refresh_token":
        if refresh_token != "valid-refresh-token":
            raise HTTPException(status_code=400, detail="Invalid refresh token")
        return TokenResponse(
            access_token="refreshed-valid-token",
            refresh_token="refreshed-valid-refresh-token",
        )

    if grant_type == "client_credentials":
        if client_secret != "test-client-secret":
            raise HTTPException(status_code=400, detail="Invalid client credentials")
        # Refresh tokens are typically not issued for client credentials.
        return TokenResponse(access_token="valid-token")

    raise HTTPException(status_code=400, detail=f"Unsupported grant type: {grant_type}")