argos/argos/schemas/config.py

72 lines
1.3 KiB
Python

from pydantic import BaseModel
from enum import StrEnum
from typing import List, Optional, Tuple
import yaml
from pydantic import BaseModel, Field, HttpUrl, validator
from datetime import datetime
from argos.checks import get_names as get_check_names
# This file contains the pydantic schemas. For the database models, check in argos.model.
class Thresholds(BaseModel):
critical: str = Field(alias="critical")
warning: str = Field(alias="warning")
class SSL(BaseModel):
thresholds: Thresholds
WebsiteCheck = dict[StrEnum("Check", get_check_names()), str | int]
class WebsitePath(BaseModel):
path: str
checks: List[WebsiteCheck]
class Website(BaseModel):
domain: HttpUrl
paths: List[WebsitePath]
class Service(BaseModel):
port: int
secrets: List[str]
class Alert(BaseModel):
error: List[str]
warning: List[str]
alert: List[str]
class General(BaseModel):
frequency: str
alerts: Alert
class Config(BaseModel):
general: General
service: Service
ssl: SSL
websites: List[Website]
def validate_config(config: dict):
return Config(**config)
# Method to load YAML file
def from_yaml(file_name):
parsed = load_yaml(file_name)
return validate_config(parsed)
def load_yaml(file_name):
with open(file_name, "r") as stream:
return yaml.safe_load(stream)