diff --git a/README.md b/README.md index d0b6214..67d0a52 100644 --- a/README.md +++ b/README.md @@ -15,15 +15,18 @@ Features : - [x] Checks can be distributed on the network thanks to a job queue ; - [x] Change the naming and use service/agent. - [x] Packaging (and `argos agent` / `argos service` commands) +- [x] Endpoints are protected by an authentication token - [ ] Local task for database cleanup (to run periodically) - [ ] Handles multiple alerting backends (email, sms, gotify) ; - [ ] Exposes a simple read-only website. +- [ ] Add a way to specify the severity of the alerts in the config +- [ ] Do not return `selected_at` and `selected_by` in the `/tasks` endpoint (?) Implemented checks : - [x] Returned status code matches what you expect ; - [x] Returned body matches what you expect ; -- [ ] SSL certificate expires in more than X days ; +- [x] SSL certificate expires in more than X days ; ## How to run ? @@ -43,7 +46,7 @@ pipenv sync Once all the dependencies are in place, here is how to run the server: ```bash -pipenv run uvicorn argos.server:app --reload +pipenv run argos server ``` The server will read a `config.yaml` file at startup, and will populate the tasks specified in it. See the configuration section below for more information on how to configure the checks you want to run. @@ -51,7 +54,7 @@ The server will read a `config.yaml` file at startup, and will populate the task And here is how to run the agent: ```bash -pipenv run argos-agent --server http://localhost:8000 +pipenv run argos agent --server http://localhost:8000 --auth "" ``` ## Configuration diff --git a/argos/agent/cli.py b/argos/agent.py similarity index 72% rename from argos/agent/cli.py rename to argos/agent.py index 3d74cad..2094b82 100644 --- a/argos/agent/cli.py +++ b/argos/agent.py @@ -1,11 +1,8 @@ import httpx import asyncio -import click from typing import List - -from argos import logging from argos.logging import logger -from argos.checks import CheckNotFound, get_check_by_name +from argos.checks import CheckNotFound, get_registered_check from argos.schemas import Task, AgentResult, SerializableException @@ -13,7 +10,7 @@ from argos.schemas import Task, AgentResult, SerializableException async def complete_task(http_client: httpx.AsyncClient, task: dict) -> dict: try: task = Task(**task) - check_class = get_check_by_name(task.check) + check_class = get_registered_check(task.check) check = check_class(http_client, task) result = await check.run() status = result.status @@ -39,10 +36,11 @@ async def post_results( logger.error(f"Failed to post results: {response.read()}") -async def run(server: str, max_tasks: int): +async def run_agent(server: str, auth: str, max_tasks: int): tasks = [] - async with httpx.AsyncClient() as http_client: + headers = {"Authorization": f"Bearer {auth}"} + async with httpx.AsyncClient(headers=headers) as http_client: # Fetch the list of tasks response = await http_client.get(f"{server}/tasks") @@ -60,21 +58,4 @@ async def run(server: str, max_tasks: int): # Post the results await post_results(http_client, server, results) else: - logger.error(f"Failed to fetch tasks: {response.read()}") - - -@click.command() -@click.option("--server", required=True, help="Server URL") -@click.option("--max-tasks", default=10, help="Maximum number of concurrent tasks") -@click.option( - "--log-level", - default="INFO", - type=click.Choice(logging.LOG_LEVELS, case_sensitive=False), -) -def main(server, max_tasks, log_level): - logging.set_log_level(log_level) - asyncio.run(run(server, max_tasks)) - - -if __name__ == "__main__": - main() + logger.error(f"Failed to fetch tasks: {response.read()}") \ No newline at end of file diff --git a/argos/agent/__init__.py b/argos/agent/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/argos/commands.py b/argos/commands.py new file mode 100644 index 0000000..75c5539 --- /dev/null +++ b/argos/commands.py @@ -0,0 +1,40 @@ +import click +import subprocess +import asyncio + +from argos.agent import run_agent +from argos import logging +from argos.logging import logger + +@click.group() +def cli(): + pass + +@cli.command() +@click.option("--server", required=True, help="Server URL") +@click.option("--auth", required=True, help="The authentication token") +@click.option("--max-tasks", default=10, help="Maximum number of concurrent tasks") +@click.option( + "--log-level", + default="INFO", + type=click.Choice(logging.LOG_LEVELS, case_sensitive=False), +) +def agent(server, auth, max_tasks, log_level): + """Runs an agent""" + logging.set_log_level(log_level) + asyncio.run(run_agent(server, auth, max_tasks)) + + +@cli.command() +@click.option("--host", default="127.0.0.1", help="Host to bind") +@click.option("--port", default=8000, type=int, help="Port to bind") +@click.option("--reload", is_flag=True, help="Enable hot reloading") +def server(host, port, reload): + """Starts the server.""" + command = ["uvicorn", "argos.server:app", "--host", host, "--port", str(port)] + if reload: + command.append("--reload") + subprocess.run(command) + +if __name__ == "__main__": + cli() diff --git a/argos/server/api.py b/argos/server/api.py index 2d9ec50..9171fe0 100644 --- a/argos/server/api.py +++ b/argos/server/api.py @@ -1,19 +1,22 @@ -from fastapi import Depends, FastAPI, HTTPException, Request -from sqlalchemy.orm import Session -from pydantic import BaseModel, ValidationError import sys +from typing import Annotated, List, Optional + +from fastapi import Depends, FastAPI, HTTPException, Request, Header +from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials +from pydantic import BaseModel, ValidationError +from sqlalchemy.orm import Session -from argos.server import queries, models -from argos.schemas import AgentResult, Task -from argos.schemas.config import from_yaml as get_schemas_from_yaml -from argos.server.database import SessionLocal, engine from argos.checks import get_registered_check from argos.logging import logger -from typing import List +from argos.schemas import AgentResult, Task +from argos.schemas.config import from_yaml as get_schemas_from_yaml +from argos.server import models, queries +from argos.server.database import SessionLocal, engine models.Base.metadata.create_all(bind=engine) app = FastAPI() +auth_scheme = HTTPBearer() def get_db(): @@ -24,6 +27,12 @@ def get_db(): db.close() +async def verify_token(token: HTTPAuthorizationCredentials = Depends(auth_scheme)): + if token.credentials not in app.config.service.secrets: + raise HTTPException(status_code=401, detail="Unauthorized") + return token + + @app.on_event("startup") async def read_config_and_populate_db(): # XXX Get filename from environment. @@ -44,8 +53,8 @@ async def read_config_and_populate_db(): # XXX Get the default limit from the config -@app.get("/tasks", response_model=list[Task]) -async def read_tasks(request: Request, limit: int = 20, db: Session = Depends(get_db)): +@app.get("/tasks", response_model=list[Task], dependencies=[Depends(verify_token)]) +async def read_tasks(request: Request, db: Session = Depends(get_db), limit: int = 20): # XXX Let the agents specifify their names (and use hostnames) tasks = await queries.list_tasks(db, agent_id=request.client.host, limit=limit) return tasks diff --git a/pyproject.toml b/pyproject.toml index daf30e4..96a8155 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ dev = [ ] [project.scripts] -argos-agent = "argos.agent.cli:main" +argos = "argos.commands:cli" [tool.pytest.ini_options] minversion = "6.0" diff --git a/tests/test_schemas_config.py b/tests/test_schemas_config.py index f7c6a43..0e9676e 100644 --- a/tests/test_schemas_config.py +++ b/tests/test_schemas_config.py @@ -16,7 +16,7 @@ def test_ssl_duration_parsing(): SSL(**erroneous_data) -def test_path_parsing(): +def test_path_parsing_transforms_to_tuples(): data = {"path": "/", "checks": [{"body-contains": "youpi"}, {"status-is": 200}]} path = WebsitePath(**data) assert len(path.checks) == 2