api/endpoints/xc.py

import logging
from typing import Optional
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from aiohttp import ClientSession, ClientTimeout
from .constructors import ValidXCRequest

class XC(FastAPI):
    """XC (CrossComm) Endpoints"""   
    def __init__(self, app: FastAPI, util, constants) -> None:
        self.app: FastAPI = app
        self.util = util
        self.constants = constants

        self.endpoints: dict = {
            "xc": self.xc_handler,
            }
        
        for endpoint, handler in self.endpoints.items():
            app.add_api_route(f"/{endpoint}", handler, methods=["POST"],
                              include_in_schema=False)

    async def xc_handler(self, data: ValidXCRequest,
                         request: Request) -> JSONResponse:
        """Handle XC Commands"""
        
        try:
                key: str = data.key
                bid: int = int(data.bid)
                cmd: str = data.cmd
                cmd_data: Optional[dict] = data.data
                if not self.util.check_key(path=request.url.path, req_type=0, key=key):
                    raise HTTPException(status_code=403, detail="Unauthorized")
                
                BID_ADDR_MAP: dict = {
                    # TODO: add Havoc? 
                }

                if not bid in BID_ADDR_MAP:
                    return JSONResponse(status_code=500, content={
                        'err': True,
                        'errorText': 'Invalid bot id'
                    })

                bot_api_url: str = f'http://{BID_ADDR_MAP[bid]}/'
                async with ClientSession() as  session:
                    async with await session.post(f"{bot_api_url}{cmd}", json=cmd_data, headers={
                        'Content-Type': 'application/json; charset=utf-8'
                    }, timeout=ClientTimeout(connect=5, sock_read=5)) as aiohttp_request:
                        response: dict = await aiohttp_request.json()
                        return JSONResponse(content={
                            'success': True,
                            'response': response
                        })
        except Exception as e:
            logging.debug("Error: %s", str(e))
            return JSONResponse(status_code=500, content={
                'err': True,
                'errorText': 'General error.',
            })