xmpp-api/src/xmpp_api/api/routers/v1/bot.py

import uuid

from fastapi import APIRouter, HTTPException, Request, Response
import sqlalchemy
from sqlmodel import select
from slixmpp.jid import JID

from xmpp_api.api.types.v1.bot import (
    AddJidRequest,
    AddJidResponse,
    AllowedJidInformation,
    BotCreationResponse,
    CreateBotRequest,
    GetBotInformation,
    SendMessageRequest,
    BotConstraint,
    BotDomainConstraint,
)
from xmpp_api.db import BotDep, SessionDep, UserDep
from xmpp_api.db.bot import AllowedJid, Bot, JIDType
import xmpp_api.db.bot as db_bot
from xmpp_api.util.token import generate_token
from xmpp_api.db import get_bot_by_id, get_jids_by_bot_id, get_jid_by_jid_token
from xmpp_api.util.constraints import bot_constraint_to_db, bot_constraint_from_db
from xmpp_api.xmpp.component import XmppApiComponentDep


router = APIRouter(prefix="/bot")


@router.post("/")
def post_create_bot(
    bot_request: CreateBotRequest, user: UserDep, session: SessionDep
) -> BotCreationResponse:
    # Convert the provided constraints into DB-type bot constraints
    try:
        constraints: list[db_bot.BotConstraint] = [
            bot_constraint_to_db(c) for c in bot_request.constraints
        ]
    except Exception:
        raise HTTPException(status_code=500)

    bot = Bot(
        name=bot_request.name,
        description=bot_request.description,
        localpart=bot_request.localpart,
        token=generate_token(64),
        constraints=constraints,
        owner_id=user.id,
    )
    session.add(bot)

    # Handle errors caused by the unique constraint
    try:
        session.commit()
    except sqlalchemy.exc.IntegrityError:
        raise HTTPException(
            status_code=400,
            detail="The name is already in use",
        )

    return BotCreationResponse(
        id=bot.id,
        name=bot.name,
        description=bot.description,
        localpart=bot.localpart,
        token=bot.token,
        constraints=bot_request.constraints,
    )


@router.post("/{bot_id}/jid")
async def post_create_bot_jid(
    bot_id: str,
    creation_request: AddJidRequest,
    user: UserDep,
    session: SessionDep,
    component: XmppApiComponentDep,
) -> AddJidResponse:
    # Check if the bot exists and we own it
    bot = get_bot_by_id(bot_id, user.id, session)
    if bot is None:
        raise HTTPException(
            status_code=404,
            detail="Unknown bot",
        )

    # Validate the domain constraint
    parsed_jid = JID(creation_request.jid)
    for constraint in bot.constraints:
        if isinstance(constraint, BotDomainConstraint):
            if parsed_jid.domain not in constraint.domains:
                raise HTTPException(
                    status_code=400,
                    detail=f'Domain "{parsed_jid.domain}" is not allowed',
                )

    # Query the JID for its type
    allowed_jid_type: JIDType
    if creation_request.type is None:
        jid_type = await component.get_entity_type(parsed_jid)
        if jid_type is None:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to query entity at {creation_request.jid}",
            )

        match jid_type:
            case "account":
                allowed_jid_type = JIDType.DIRECT
            case "groupchat":
                allowed_jid_type = JIDType.GC
    else:
        allowed_jid_type = creation_request.type

    # Deal with groupchat shenanigans
    if allowed_jid_type == JIDType.GC:
        # TODO: Join the groupchat
        raise HTTPException(
            status_code=500,
            detail="Groupchats are not yet handled",
        )

    # Add the JID
    jid = AllowedJid(
        jid=creation_request.jid,
        type=allowed_jid_type,
        # This token is only for identification
        token=uuid.uuid4().hex,
        bot_id=bot.id,
    )
    session.add(jid)
    try:
        session.commit()
    except sqlalchemy.exc.IntegrityError:
        raise HTTPException(
            status_code=400,
            detail="The JID already exists",
        )

    return AddJidResponse(
        token=jid.token,
    )


@router.delete("/{bot_id}/jid/{jid_token}")
def delete_bot_jid_token(
    bot_id: str, jid_token: str, user: UserDep, session: SessionDep
):
    jid = get_jid_by_jid_token(jid_token, session)
    if jid is None or jid.bot_id != bot_id:
        raise HTTPException(
            status_code=404,
        )

    session.delete(jid)
    session.commit()


@router.post("/message/{jid_token}")
def post_bot_message(
    jid_token: str,
    message: SendMessageRequest,
    request: Request,
    bot: BotDep,
    session: SessionDep,
    component: XmppApiComponentDep,
):
    # Is the bot allowed to access this JID?
    jid = session.exec(
        select(AllowedJid).where(
            AllowedJid.token == jid_token, AllowedJid.bot_id == bot.id
        )
    ).first()
    if jid is None:
        raise HTTPException(
            detail="Unknown JID token",
            status_code=404,
        )

    # Validate the domain constraint
    parsed_jid = JID(jid.jid)
    for constraint in bot.constraints:
        if isinstance(constraint, BotDomainConstraint):
            if parsed_jid.domain not in constraint.domains:
                raise HTTPException(status_code=400)

    match jid.type:
        case JIDType.DIRECT:
            component.send_direct_message(
                body=message.body,
                localpart=bot.localpart,
                nick=bot.name,
                recipient=jid.jid,
            )
        case _:
            raise HTTPException(status_code=500)

    return Response(status_code=200)


@router.get("/{bot_id}")
def get_bot_information(
    bot_id: str, user: UserDep, session: SessionDep
) -> GetBotInformation:
    # Find the bot
    bot = get_bot_by_id(bot_id, user.id, session)
    if bot is None:
        raise HTTPException(status_code=404)

    try:
        constraints: list[BotConstraint] = [
            bot_constraint_from_db(c) for c in bot.constraints
        ]
    except Exception:
        raise HTTPException(status_code=500)

    return GetBotInformation(
        id=bot.id,
        name=bot.name,
        description=bot.description,
        localpart=bot.localpart,
        jids=[
            AllowedJidInformation(
                jid=jid.jid,
                token=jid.token,
            )
            for jid in get_jids_by_bot_id(bot.id, session)
        ],
        constraints=constraints,
    )


@router.delete("/{bot_id}")
def delete_bot(bot_id: str, user: UserDep, session: SessionDep):
    bot = get_bot_by_id(bot_id, user.id, session)
    if bot is None:
        raise HTTPException(status_code=404)

    for jid in get_jids_by_bot_id(bot_id, session):
        session.delete(jid)
    session.delete(bot)
    session.commit()
    return Response(status_code=200)