Refactor codebase to by DRY

This commit is contained in:
Abel Luck 2022-12-01 16:31:04 +00:00
parent c925079e8b
commit 83a526c533
13 changed files with 320 additions and 131 deletions

View file

@ -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"}