Refactor codebase to by DRY
This commit is contained in:
parent
c925079e8b
commit
83a526c533
13 changed files with 320 additions and 131 deletions
192
ops_bot/main.py
192
ops_bot/main.py
|
|
@ -1,36 +1,28 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Literal, Optional, Tuple, cast
|
||||
from typing import Any, Dict, List, Optional, Protocol, Tuple, cast
|
||||
|
||||
import uvicorn
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from pydantic import BaseSettings
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request, status
|
||||
from fastapi.security import (
|
||||
HTTPAuthorizationCredentials,
|
||||
HTTPBasic,
|
||||
HTTPBasicCredentials,
|
||||
HTTPBearer,
|
||||
)
|
||||
|
||||
from ops_bot import aws, pagerduty
|
||||
from ops_bot.config import BotSettings, RoutingKey, load_config
|
||||
from ops_bot.gitlab import hook as gitlab_hook
|
||||
from ops_bot.matrix import MatrixClient, MatrixClientSettings
|
||||
from ops_bot.matrix import MatrixClient
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class BotSettings(BaseSettings):
|
||||
bearer_token: str
|
||||
routing_keys: Dict[str, str]
|
||||
log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO"
|
||||
matrix: MatrixClientSettings
|
||||
|
||||
class Config:
|
||||
env_prefix = "BOT_"
|
||||
secrets_dir = "/run/secrets"
|
||||
case_sensitive = False
|
||||
|
||||
|
||||
app = FastAPI()
|
||||
security = HTTPBearer()
|
||||
bearer_security = HTTPBearer(auto_error=False)
|
||||
basic_security = HTTPBasic(auto_error=False)
|
||||
|
||||
|
||||
async def get_matrix_service(request: Request) -> MatrixClient:
|
||||
|
|
@ -46,14 +38,8 @@ async def matrix_main(matrix_client: MatrixClient) -> None:
|
|||
|
||||
@app.on_event("startup")
|
||||
async def startup_event() -> None:
|
||||
# if "config.json" exists read it
|
||||
if Path("config.json").exists():
|
||||
bot_settings = BotSettings.parse_file("config.json")
|
||||
else:
|
||||
bot_settings = BotSettings(_env_file=".env", _env_file_encoding="utf-8")
|
||||
logging.getLogger().setLevel(bot_settings.log_level)
|
||||
bot_settings.matrix.join_rooms = list(bot_settings.routing_keys.values())
|
||||
c = MatrixClient(settings=bot_settings.matrix)
|
||||
bot_settings = load_config()
|
||||
c = MatrixClient(settings=bot_settings.matrix, join_rooms=bot_settings.get_rooms())
|
||||
app.state.matrix_client = c
|
||||
app.state.bot_settings = bot_settings
|
||||
asyncio.create_task(matrix_main(c))
|
||||
|
|
@ -69,87 +55,111 @@ async def root() -> Dict[str, str]:
|
|||
return {"message": "Hello World"}
|
||||
|
||||
|
||||
def authorize(
|
||||
request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)
|
||||
async def bearer_token_authorizer(
|
||||
route: RoutingKey,
|
||||
request: Request,
|
||||
basic_credentials: Optional[HTTPBasicCredentials],
|
||||
bearer_credentials: Optional[HTTPAuthorizationCredentials],
|
||||
) -> bool:
|
||||
|
||||
bearer_token: Optional[str] = route.secret_token
|
||||
return (
|
||||
bearer_credentials is not None
|
||||
and bearer_credentials.credentials == bearer_token
|
||||
)
|
||||
|
||||
|
||||
async def nop_authorizer(
|
||||
route: RoutingKey,
|
||||
request: Request,
|
||||
basic_credentials: Optional[HTTPBasicCredentials],
|
||||
bearer_credentials: Optional[HTTPAuthorizationCredentials],
|
||||
) -> bool:
|
||||
bearer_token = request.app.state.bot_settings.bearer_token
|
||||
if credentials.credentials != bearer_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Incorrect bearer token",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
def get_destination(bot_settings: BotSettings, routing_key: str) -> Optional[str]:
|
||||
return bot_settings.routing_keys.get(routing_key, None)
|
||||
def get_route(bot_settings: BotSettings, path_key: str) -> Optional[RoutingKey]:
|
||||
# find path_key in bot_settings.routing_keys
|
||||
for route in bot_settings.routing_keys:
|
||||
if route.path_key == path_key:
|
||||
return route
|
||||
return None
|
||||
|
||||
|
||||
async def receive_helper(request: Request) -> Tuple[str, Any]:
|
||||
payload: Any = await request.json()
|
||||
routing_key = request.path_params["routing_key"]
|
||||
room_id = get_destination(request.app.state.bot_settings, routing_key)
|
||||
if room_id is None:
|
||||
class Authorizer(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
route: RoutingKey,
|
||||
request: Request,
|
||||
basic_credentials: Optional[HTTPBasicCredentials],
|
||||
bearer_credentials: Optional[HTTPAuthorizationCredentials],
|
||||
) -> bool:
|
||||
...
|
||||
|
||||
|
||||
class ParseHandler(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
route: RoutingKey,
|
||||
payload: Any,
|
||||
request: Request,
|
||||
) -> List[Tuple[str, str]]:
|
||||
...
|
||||
|
||||
|
||||
handlers: Dict[str, Tuple[Authorizer, ParseHandler]] = {
|
||||
"gitlab": (gitlab_hook.authorize, gitlab_hook.parse_event),
|
||||
"pagerduty": (bearer_token_authorizer, pagerduty.parse_pagerduty_event),
|
||||
"aws-sns": (nop_authorizer, aws.parse_sns_event),
|
||||
}
|
||||
|
||||
|
||||
@app.post("/hook/{routing_key}")
|
||||
async def webhook_handler(
|
||||
request: Request,
|
||||
routing_key: str,
|
||||
basic_credentials: Optional[HTTPBasicCredentials] = Depends(basic_security),
|
||||
bearer_credentials: Optional[HTTPAuthorizationCredentials] = Depends(
|
||||
bearer_security
|
||||
),
|
||||
matrix_client: MatrixClient = Depends(get_matrix_service),
|
||||
) -> Dict[str, str]:
|
||||
route = get_route(request.app.state.bot_settings, routing_key)
|
||||
|
||||
if not route:
|
||||
logging.error(f"unknown routing key {routing_key}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail="Unknown routing key"
|
||||
)
|
||||
payload_str = json.dumps(payload, sort_keys=True, indent=2)
|
||||
logging.info(f"received payload: \n {payload_str}")
|
||||
return room_id, payload
|
||||
|
||||
|
||||
@app.post("/hook/pagerduty/{routing_key}")
|
||||
async def pagerduty_hook(
|
||||
request: Request,
|
||||
matrix_client: MatrixClient = Depends(get_matrix_service),
|
||||
auth: bool = Depends(authorize),
|
||||
) -> Dict[str, str]:
|
||||
room_id, payload = await receive_helper(request)
|
||||
msg_plain, msg_formatted = pagerduty.parse_pagerduty_event(payload)
|
||||
await matrix_client.room_send(
|
||||
room_id,
|
||||
msg_plain,
|
||||
message_formatted=msg_formatted,
|
||||
)
|
||||
return {"message": msg_plain, "message_formatted": msg_formatted}
|
||||
|
||||
|
||||
@app.post("/hook/aws-sns/{routing_key}")
|
||||
async def aws_sns_hook(
|
||||
request: Request, matrix_client: MatrixClient = Depends(get_matrix_service)
|
||||
) -> Dict[str, str]:
|
||||
room_id, payload = await receive_helper(request)
|
||||
msg_plain, msg_formatted = aws.parse_sns_event(payload)
|
||||
await matrix_client.room_send(
|
||||
room_id,
|
||||
msg_plain,
|
||||
message_formatted=msg_formatted,
|
||||
)
|
||||
return {"message": msg_plain, "message_formatted": msg_formatted}
|
||||
|
||||
|
||||
@app.post("/hook/gitlab/{routing_key}")
|
||||
async def gitlab_webhook(
|
||||
request: Request,
|
||||
x_gitlab_token: str = Header(default=""),
|
||||
x_gitlab_event: str = Header(default=""),
|
||||
matrix_client: MatrixClient = Depends(get_matrix_service),
|
||||
) -> Dict[str, str]:
|
||||
bearer_token = request.app.state.bot_settings.bearer_token
|
||||
if x_gitlab_token != bearer_token:
|
||||
handler: Optional[Tuple[Authorizer, ParseHandler]] = handlers.get(route.hook_type)
|
||||
if not handler:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Incorrect X-Gitlab-Token"
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail="Unknown hook type"
|
||||
)
|
||||
room_id, payload = await receive_helper(request)
|
||||
messages = await gitlab_hook.parse_event(x_gitlab_event, payload)
|
||||
|
||||
authorizer, parse_handler = handler
|
||||
|
||||
if not await authorizer(
|
||||
route,
|
||||
request=request,
|
||||
bearer_credentials=bearer_credentials,
|
||||
basic_credentials=basic_credentials,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid credentials"
|
||||
)
|
||||
|
||||
payload: Any = await request.json()
|
||||
|
||||
messages = await parse_handler(route, payload, request=request)
|
||||
for msg_plain, msg_formatted in messages:
|
||||
await matrix_client.room_send(
|
||||
room_id,
|
||||
route.room_id,
|
||||
msg_plain,
|
||||
message_formatted=msg_formatted,
|
||||
)
|
||||
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue