From 2ef999fa63f299b4dd946f9dae354c30b10616d6 Mon Sep 17 00:00:00 2001 From: Luc Didry Date: Tue, 10 Dec 2024 13:46:23 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20=E2=80=94=20Allow=20to=20specify=20?= =?UTF-8?q?form=20data=20and=20headers=20for=20checks=20(fix=20#70)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CHANGELOG.md | 1 + argos/agent.py | 65 +++++++++++++++---- argos/config-example.yaml | 11 ++++ argos/schemas/config.py | 21 +++++- argos/schemas/models.py | 1 + .../31255a412d63_add_form_data_to_tasks.py | 28 ++++++++ argos/server/models.py | 8 ++- argos/server/queries.py | 11 +++- tests/config.yaml | 4 +- tests/test_checks.py | 1 + tests/websites.yaml | 3 +- 11 files changed, 137 insertions(+), 17 deletions(-) create mode 100644 argos/server/migrations/versions/31255a412d63_add_form_data_to_tasks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fe361a8..b368c0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - ⚡ — Mutualize check requests (#68) - ✨ — Ability to delay notification after X failures (#71) - 🐛 — Fix bug when changing IP version not removing tasks (#72) +- ✨ — Allow to specify form data and headers for checks (#70) ## 0.6.1 diff --git a/argos/agent.py b/argos/agent.py index 23c62fc..2cca0d6 100644 --- a/argos/agent.py +++ b/argos/agent.py @@ -6,6 +6,7 @@ import asyncio import json import logging import socket +from hashlib import md5 from time import sleep from typing import List @@ -57,7 +58,7 @@ class ArgosAgent: # pylint: disable-msg=too-many-instance-attributes self._http_client = httpx.AsyncClient(headers=auth_header) ua_header = { - "User-Agent": f"Prout Argos Panoptes {VERSION} " + "User-Agent": f"Argos Panoptes {VERSION} " "(about: https://argos-monitoring.framasoft.org/)", } self._http_client_v4 = httpx.AsyncClient( @@ -77,24 +78,59 @@ class ArgosAgent: # pylint: disable-msg=too-many-instance-attributes await asyncio.sleep(self.wait_time) async def _do_request(self, group: str, details: dict): + headers = {} + if details["request_data"] is not None: + request_data = json.loads(details["request_data"]) + if request_data["headers"] is not None: + headers = request_data["headers"] + + if details["ip_version"] == "4": + http_client = self._http_client_v4 + else: + http_client = self._http_client_v6 try: - if details["ip_version"] == "4": - response = await self._http_client_v4.request( # type: ignore[union-attr] - method=details["method"], url=details["url"], timeout=60 + if details["request_data"] is None or request_data["data"] is None: + response = await http_client.request( # type: ignore[union-attr] + method=details["method"], + url=details["url"], + headers=headers, + timeout=60, + ) + elif request_data["json"]: + response = await http_client.request( # type: ignore[union-attr] + method=details["method"], + url=details["url"], + headers=headers, + json=request_data["data"], + timeout=60, ) else: - response = await self._http_client_v6.request( # type: ignore[union-attr] - method=details["method"], url=details["url"], timeout=60 + response = await http_client.request( # type: ignore[union-attr] + method=details["method"], + url=details["url"], + headers=headers, + data=request_data["data"], + timeout=60, ) except httpx.ReadError: sleep(1) - if details["ip_version"] == "4": - response = await self._http_client_v4.request( # type: ignore[union-attr] + if details["request_data"] is None or request_data["data"] is None: + response = await http_client.request( # type: ignore[union-attr] method=details["method"], url=details["url"], timeout=60 ) + elif request_data["json"]: + response = await http_client.request( # type: ignore[union-attr] + method=details["method"], + url=details["url"], + json=request_data["data"], + timeout=60, + ) else: - response = await self._http_client_v6.request( # type: ignore[union-attr] - method=details["method"], url=details["url"], timeout=60 + response = await http_client.request( # type: ignore[union-attr] + method=details["method"], + url=details["url"], + data=request_data["data"], + timeout=60, ) self._res_cache[group] = response @@ -135,14 +171,21 @@ class ArgosAgent: # pylint: disable-msg=too-many-instance-attributes group = task.task_group if task.check == "http-to-https": + data = task.request_data + if data is None: + data = "" url = str(httpx.URL(task.url).copy_with(scheme="http")) - group = f"{task.method}-{task.ip_version}-{url}" + group = ( + f"{task.method}-{task.ip_version}-{url}-" + f"{md5(data.encode()).hexdigest()}" + ) _task["task_group"] = group req_groups[group] = { "url": url, "ip_version": task.ip_version, "method": task.method, + "request_data": task.request_data, } requests = [] diff --git a/argos/config-example.yaml b/argos/config-example.yaml index b957502..5508bff 100644 --- a/argos/config-example.yaml +++ b/argos/config-example.yaml @@ -190,6 +190,17 @@ websites: - 302 - 307 - path: "/admin/" + methode: "POST" + # Send form data in the request + request_data: + data: + login: "admin" + password: "my-password" + # To send data as JSON (optional, default is false): + is_json: true + # To send additional headers + headers: + Authorization: "Bearer foo-bar-baz" checks: # Check that the return HTTP status is one of those # Similar to status-is, verify that you don’t mistyped it! diff --git a/argos/schemas/config.py b/argos/schemas/config.py index caeba95..1baa7ed 100644 --- a/argos/schemas/config.py +++ b/argos/schemas/config.py @@ -5,7 +5,7 @@ For database models, see argos.server.models. import json -from typing import Dict, List, Literal, Tuple +from typing import Any, Dict, List, Literal, Tuple from durations_nlp import Duration from pydantic import ( @@ -18,7 +18,7 @@ from pydantic import ( PositiveInt, field_validator, ) -from pydantic.functional_validators import BeforeValidator +from pydantic.functional_validators import AfterValidator, BeforeValidator from pydantic.networks import UrlConstraints from pydantic_core import Url from typing_extensions import Annotated @@ -104,9 +104,26 @@ def parse_checks(value): return (name, expected) +def parse_request_data(value): + """Turn form or JSON data into JSON string""" + + return json.dumps( + {"data": value.data, "json": value.is_json, "headers": value.headers} + ) + + +class RequestData(BaseModel): + data: Any = None + is_json: bool = False + headers: Dict[str, str] | None = None + + class WebsitePath(BaseModel): path: str method: Method = "GET" + request_data: Annotated[ + RequestData, AfterValidator(parse_request_data) + ] | None = None checks: List[ Annotated[ Tuple[str, str], diff --git a/argos/schemas/models.py b/argos/schemas/models.py index b1eb33a..ed1bc20 100644 --- a/argos/schemas/models.py +++ b/argos/schemas/models.py @@ -22,6 +22,7 @@ class Task(BaseModel): ip_version: IPVersion check: str method: Method + request_data: str | None expected: str task_group: str retry_before_notification: int diff --git a/argos/server/migrations/versions/31255a412d63_add_form_data_to_tasks.py b/argos/server/migrations/versions/31255a412d63_add_form_data_to_tasks.py new file mode 100644 index 0000000..bb4deaa --- /dev/null +++ b/argos/server/migrations/versions/31255a412d63_add_form_data_to_tasks.py @@ -0,0 +1,28 @@ +"""Add request data to tasks + +Revision ID: 31255a412d63 +Revises: 80a29f64f91c +Create Date: 2024-12-09 16:40:20.926138 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = "31255a412d63" +down_revision: Union[str, None] = "80a29f64f91c" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + with op.batch_alter_table("tasks", schema=None) as batch_op: + batch_op.add_column(sa.Column("request_data", sa.String(), nullable=True)) + + +def downgrade() -> None: + with op.batch_alter_table("tasks", schema=None) as batch_op: + batch_op.drop_column("request_data") diff --git a/argos/server/models.py b/argos/server/models.py index c031d7b..6fd7b6c 100644 --- a/argos/server/models.py +++ b/argos/server/models.py @@ -1,6 +1,7 @@ """Database models""" from datetime import datetime, timedelta +from hashlib import md5 from typing import List, Literal from sqlalchemy import ( @@ -17,10 +18,14 @@ from argos.schemas.utils import IPVersion, Method def compute_task_group(context) -> str: + data = context.current_parameters["request_data"] + if data is None: + data = "" return ( f"{context.current_parameters['method']}-" f"{context.current_parameters['ip_version']}-" - f"{context.current_parameters['url']}" + f"{context.current_parameters['url']}-" + f"{md5(data.encode()).hexdigest()}" ) @@ -67,6 +72,7 @@ class Task(Base): ), insert_default="GET", ) + request_data: Mapped[str] = mapped_column(nullable=True) # Orchestration-related selected_by: Mapped[str] = mapped_column(nullable=True) diff --git a/argos/server/queries.py b/argos/server/queries.py index bea4bc9..be9afd7 100644 --- a/argos/server/queries.py +++ b/argos/server/queries.py @@ -256,6 +256,7 @@ async def update_from_config(db: Session, config: schemas.Config): # pylint: di .filter( Task.url == url, Task.method == p.method, + Task.request_data == p.request_data, Task.check == check_key, Task.expected == expected, Task.ip_version == ip_version, @@ -300,7 +301,14 @@ async def update_from_config(db: Session, config: schemas.Config): # pylint: di ) else: - properties = (url, p.method, check_key, expected, ip_version) + properties = ( + url, + p.method, + check_key, + expected, + ip_version, + p.request_data, + ) if properties not in unique_properties: unique_properties.append(properties) task = Task( @@ -308,6 +316,7 @@ async def update_from_config(db: Session, config: schemas.Config): # pylint: di url=url, ip_version=ip_version, method=p.method, + request_data=p.request_data, check=check_key, expected=expected, frequency=frequency, diff --git a/tests/config.yaml b/tests/config.yaml index 82f055b..a1cec52 100644 --- a/tests/config.yaml +++ b/tests/config.yaml @@ -1,6 +1,8 @@ +--- general: db: - # The database URL, as defined in SQLAlchemy docs : https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls + # The database URL, as defined in SQLAlchemy docs: + # https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls url: "sqlite:////tmp/test-argos.db" env: test cookie_secret: "foo-bar-baz" diff --git a/tests/test_checks.py b/tests/test_checks.py index 18f3a37..c3a458d 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -37,6 +37,7 @@ def ssl_task(now): domain="https://example.org", ip_version="6", method="GET", + request_data=None, task_group="GET-6-https://example.org", check="ssl-certificate-expiration", retry_before_notification=0, diff --git a/tests/websites.yaml b/tests/websites.yaml index f2d50dc..da19ae9 100644 --- a/tests/websites.yaml +++ b/tests/websites.yaml @@ -1,6 +1,7 @@ +--- - domain: "https://mypads.framapad.org" paths: - path: "/mypads/" checks: - status-is: 200 - - body-contains: '
' \ No newline at end of file + - body-contains: '
'