Add authorization headers

- Service endpoints now require an authentication
- Changed the location of the commands, which are now in `commands.py`
- There is now only one "argos" command which can start the server or an agent.

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 "<auth-token>"
 ```
 
 ## Configuration

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()}")
+            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()

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()

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])
+@app.get("/tasks", response_model=list[Task], dependencies=[Depends(verify_token)])
-async def read_tasks(request: Request, limit: int = 20, db: Session = Depends(get_db)):
+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

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"

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