Fix API client retries and return the job ID

README.rst

@@ -896,7 +896,24 @@ Dalton API Client
 -----------------
 
 An API Client has been added in ``api/dalton.py`` that performs API calls with Python requests. 
-The client is limited to GET and POST requests.
+The client is limited to GET and POST requests. It needs the ``requests`` package.
+
+``submit_job()`` submits a teapot job and returns its job ID (several, comma-separated, if Dalton
+split the pcap). If Dalton returns an error page instead, it raises ``DaltonAPIError`` with the
+error message. Pcaps can be passed as bytes or open files; a file is read from its current position,
+so pass bytes to submit the same pcap more than once.
+
+Each request is tried up to three times. A GET is retried after a 429 or 5xx response, a
+connection error or a timeout. A job submission is retried only when the request can't have
+queued the job: a 429 response, or a failure to connect. A retry after a 429 or 5xx waits for the
+server's ``Retry-After``, if it gives one in seconds (at most 300), or 60 seconds; a retry after a
+connection error waits 10 seconds. With urllib3 2, a host name that doesn't resolve isn't retried.
+
+The client waits up to 60 seconds to connect and for the reply to a GET (``DaltonAPI(url,
+timeout=...)``). Dalton processes the pcap before it replies to a job submission, so the client
+waits up to an hour for that reply (``submit_timeout=...``; ``None`` waits forever). A submission
+that times out may still have been queued. ``headers={...}`` adds headers, such as an API key, to
+every request.
 
 Submit Job using the API
 ------------------------
@@ -931,10 +948,12 @@ Submit the job as a shortlived ``teapotJob`` if you plan to make multiple calls
 
 Other useful arguments to submit a job are: 
 
-- ``custom_rules`` in which you may include the custom rules you may want to test with your job,
+- ``custom_ruleset`` in which you may include the custom rules you may want to test with your job,
+  along with ``optionCustomRuleset`` set to ``custom`` (without it, the custom rules are ignored),
 - ``optionAlertDetailed``, ``optionEveLog``, ``optionOtherLogs``: this can be set to ``True`` if you want to generate additional logs with your job.
 
-An example script can be found in ``api/examples/submit_job.py``.
+An example script can be found in ``api/examples/job_submission.py``. Run it from the repository
+root with a pcap: ``python -m api.examples.job_submission <pcap>``.
 
 Teapot Jobs
 ===========

api/dalton.py

@@ -1,70 +1,149 @@
 """Dalton API client."""
 
+from __future__ import annotations
+
+import html
+import re
 import time
 
 import requests
-from requests.exceptions import HTTPError
+import urllib3.exceptions
+from urllib3.exceptions import ConnectTimeoutError, ProxyError
 
 RETRIES = 3
 SLEEP_TIME = 60
+# wait before retrying a connection failure: long enough for a restart, short
+# enough that a wrong port fails quickly
+CONNECT_SLEEP_TIME = 10
+TIMEOUT = 60
+# Dalton processes the pcap before it answers a job submission; this only
+# stops the client waiting forever
+SUBMIT_TIMEOUT = 3600
+# longest wait a server's Retry-After can ask for
+MAX_RETRY_AFTER = 300
 
+# GETs are retried on any of these. A job submission (POST) is retried only on
+# a 429, which means the request wasn't processed, so a retry can't create a
+# second job (a proxy's 503 can come after Dalton queued it).
+RETRY_CODES_GET = (429, 500, 502, 503, 504)
+RETRY_CODES_POST = (429,)
 
-class DaltonAPI:
-    def __init__(self, api_link):
-        self.api_link = api_link
+# what a teapot job submission returns: a job ID, or several, comma-separated,
+# when the pcap was split
+TEAPOT_JIDS_RE = re.compile(r"^teapot_[0-9a-f]{16}(,teapot_[0-9a-f]{16})*$")
 
-    def _dalton_get(self, endpoint: str) -> requests.Response:
-        for _ in range(RETRIES):
-            try:
-                response = requests.get(url=f"{self.api_link}/{endpoint}")
-                response.raise_for_status()
-            except HTTPError as exc:
-                code = exc.response.status_code
 
-                if code in [429, 500, 502, 503, 504]:
-                    # retry after n seconds
-                    time.sleep(SLEEP_TIME)
+class DaltonAPIError(Exception):
+    """Dalton rejected a request without an HTTP error status."""
 
-                raise
 
-        return response
+class DaltonAPI:
+    def __init__(
+        self, api_link, timeout=TIMEOUT, submit_timeout=SUBMIT_TIMEOUT, headers=None
+    ):
+        """
+        Args:
+            api_link (str): Dalton base URL. Ex. http://localhost
+            timeout (float): seconds to wait for a connection, and for a reply to a GET.
+                A (connect, read) tuple sets the two separately.
+            submit_timeout (float): seconds to wait for the reply to a job submission.
+                Dalton processes the pcap before it replies, so the default is an hour;
+                None waits forever. A submission that times out may still have been queued.
+            headers (dict): extra headers sent with every request, e.g. an API key
+        """
+        self.api_link = api_link
+        self.timeout = timeout
+        self.submit_timeout = submit_timeout
+        self.headers = headers
 
-    def _dalton_post(self, endpoint: str, data: dict, files: dict) -> requests.Response:
-        for _ in range(RETRIES):
+    def _request(
+        self, method: str, endpoint: str, retry_codes, timeout, **kwargs
+    ) -> requests.Response:
+        """Make a request, trying up to RETRIES times.
+
+        A retryable status waits for the server's Retry-After, if it gives one in
+        seconds (up to MAX_RETRY_AFTER), or SLEEP_TIME. Connection failures are
+        retried after CONNECT_SLEEP_TIME: for a GET any failure, for a POST only
+        one that happened before the request was sent. A host name that doesn't
+        resolve isn't retried. Each attempt uses a new connection, so a POST never
+        goes out on a kept-alive connection the server is closing.
+        """
+        attempts = max(1, RETRIES)
+        for attempt in range(1, attempts + 1):
+            last = attempt == attempts
             try:
-                response = requests.post(
+                response = requests.request(
+                    method,
                     url=f"{self.api_link}/{endpoint}",
-                    data=data,
-                    files=files,
+                    headers=self.headers,
+                    timeout=timeout,
+                    **kwargs,
                 )
-                response.raise_for_status()
-            except HTTPError as exc:
-                code = exc.response.status_code
-
-                if code in [429, 500, 502, 503, 504]:
-                    # retry after n seconds
-                    time.sleep(SLEEP_TIME)
-
+            except requests.exceptions.SSLError:
                 raise
+            except (
+                requests.ConnectionError,
+                requests.Timeout,
+                requests.exceptions.ChunkedEncodingError,
+                requests.exceptions.ContentDecodingError,
+            ) as exc:
+                if (
+                    last
+                    or _name_not_resolved(exc)
+                    or not (method == "GET" or _never_sent(exc))
+                ):
+                    raise
+                time.sleep(CONNECT_SLEEP_TIME)
+                continue
+            if response.status_code in retry_codes and not last:
+                time.sleep(_retry_after(response))
+                continue
+            response.raise_for_status()
+            return response
 
-        return response
+    def _dalton_get(self, endpoint: str) -> requests.Response:
+        return self._request("GET", endpoint, RETRY_CODES_GET, self.timeout)
+
+    def _dalton_post(self, endpoint: str, data: dict, files) -> requests.Response:
+        # read the files once, from where the caller left them, so that every
+        # attempt sends the same bytes (requests reads them into memory anyway)
+        pairs = files.items() if hasattr(files, "items") else (files or [])
+        files = [(name, _read_file(name, value)) for name, value in pairs]
+        connect = self.timeout[0] if isinstance(self.timeout, tuple) else self.timeout
+        return self._request(
+            "POST",
+            endpoint,
+            RETRY_CODES_POST,
+            (connect, self.submit_timeout),
+            data=data,
+            files=files,
+        )
 
     def submit_job(
         self,
         sensor_tech: str,
         prod_ruleset: str,
         configuration: str,
         custom_rules: str,
-        files: dict,
-    ) -> None:
+        files: dict | list,
+    ) -> str:
         """Submits a job for a specific sensor tech, ruleset, configuration and pcap.
 
         Args:
             sensor_tech (str): sensor tech and version. Ex. suricata/6.0.10
             prod_ruleset (str): path to prod ruleset. Ex. /opt/dalton/rulesets/suricata/ET-20230301-all-suricata.rules
             configuration (str): sensor configuration. Can be read from a yaml file. See "examples/job_submission.py"
             custom_rules (str): rules to be tested. Can be read from a rules file. See "examples/job_submission.py"
-            files (dict): dictionary that includes tuple with pcap file. Ex. {"coverage-pcap0": ("test_job.pcap", pcap)}
+            files (dict | list): the pcap files, as a dict or a list of (name, value) pairs.
+                Ex. {"coverage-pcap0": ("test_job.pcap", pcap)}. An open file is read from its current
+                position, as requests would; pass bytes to submit the same pcap more than once.
+
+        Returns:
+            str: the job ID, or a comma-separated list of job IDs if Dalton split the pcap.
+            Ex. 'teapot_0123456789abcdef'
+
+        Raises:
+            DaltonAPIError: Dalton returned an error page instead of a job ID.
         """
         data = {
             "sensor_tech": sensor_tech,
@@ -74,7 +153,17 @@ def submit_job(
             "custom_ruleset": custom_rules,
             "teapotJob": 1,
         }
-        self._dalton_post("dalton/coverage/summary", data=data, files=files)
+        if custom_rules:
+            # Dalton ignores custom_ruleset unless this is set
+            data["optionCustomRuleset"] = "custom"
+        response = self._dalton_post("dalton/coverage/summary", data=data, files=files)
+        # a teapot job's response is its job ID; anything else is Dalton's error page
+        jids = response.text.strip()
+        if not TEAPOT_JIDS_RE.match(jids):
+            raise DaltonAPIError(
+                f"Dalton did not accept the job: {_error_text(jids)[:500]}"
+            )
+        return jids
 
     def get_sensor_tech(self, sensor: str) -> list:
         """Get sensor type and version.
@@ -146,3 +235,52 @@ def get_current_sensors(self) -> dict:
             "dalton/controller_api/get-current-sensors-json-full"
         )
         return response.json()
+
+
+def _error_text(page: str) -> str:
+    """The message on Dalton's error page, or the page's text if it isn't one."""
+    match = re.search(r'(?is)<div class="span9">(.*?)</div>', page)
+    text = match.group(1) if match else page
+    text = re.sub(r"(?is)<(script|style)\b.*?</\1>", " ", text)
+    text = re.sub(r"<[^>]+>", " ", text)
+    return " ".join(html.unescape(text).split())
+
+
+def _read_file(name, value):
+    """A requests files value with any open file replaced by its contents."""
+    if isinstance(value, (tuple, list)):
+        if len(value) > 1 and hasattr(value[1], "read"):
+            return (value[0], value[1].read(), *value[2:])
+        return value
+    if hasattr(value, "read"):
+        # requests names the upload after the file; keep that name
+        return (requests.utils.guess_filename(value) or name, value.read())
+    return value
+
+
+def _reason(exc: Exception):
+    """The urllib3 error behind a requests connection error, if there is one."""
+    reason = getattr(exc.args[0], "reason", None) if exc.args else None
+    # a ProxyError can wrap a failure from before or after the request went out
+    if isinstance(reason, ProxyError):
+        reason = getattr(reason, "original_error", None)
+    return reason
+
+
+def _never_sent(exc: Exception) -> bool:
+    """Whether a request failed while connecting, before the server (or proxy) saw it."""
+    return isinstance(_reason(exc), ConnectTimeoutError)
+
+
+def _name_not_resolved(exc: Exception) -> bool:
+    """Whether the host name didn't resolve (urllib3 2 only), which a retry won't fix."""
+    error = getattr(urllib3.exceptions, "NameResolutionError", None)
+    return error is not None and isinstance(_reason(exc), error)
+
+
+def _retry_after(response: requests.Response) -> float:
+    """Seconds to wait before retrying: the server's Retry-After, or SLEEP_TIME."""
+    value = response.headers.get("Retry-After", "").strip()
+    if value.isascii() and value.isdigit():
+        return min(int(value), MAX_RETRY_AFTER)
+    return SLEEP_TIME

api/examples/job_submission.py

@@ -1,38 +1,50 @@
-"""Example on how to submit a job using the Dalton Client API. Mock data is in mocks directory."""
+"""Example on how to submit a job using the Dalton Client API.
+
+The engine config and custom rules come from the mocks directory, the pcap from the command line.
+
+Run it from the repository root: python -m api.examples.job_submission <pcap>
+"""
 
 import os
+import sys
 
 from api.dalton import DaltonAPI
 
 DALTON_URL = os.getenv("DALTON_URL", "http://localhost")
 SENSOR_NAME = "suricata"
 CURRENT_PATH = os.getcwd()
 
+if len(sys.argv) != 2:
+    sys.exit("usage: python -m api.examples.job_submission <pcap>")
+PCAP = sys.argv[1]
+
 # create API client
 dalton_client = DaltonAPI(DALTON_URL)
 
 # this file needs to include the sensor configuration in yaml format. Make sure that the files are in the proper directory.
-conf = open(f"{CURRENT_PATH}/api/examples/mocks/engine_conf.yaml", "r")
-configuration = conf.read()
+with open(f"{CURRENT_PATH}/api/examples/mocks/engine_conf.yaml", "r") as conf:
+    configuration = conf.read()
 
 # this file needs to include any custom rules you want to test
-cust = open(f"{CURRENT_PATH}/api/examples/mocks/custom_rules.rules", "r")
-custom_rules = cust.read()
+with open(f"{CURRENT_PATH}/api/examples/mocks/custom_rules.rules", "r") as cust:
+    custom_rules = cust.read()
 
 # get rulesets and sensor techs to use in job submission parameters
 dalton_rulesets = dalton_client.get_prod_rulesets(SENSOR_NAME)
 dalton_sensor_techs = dalton_client.get_sensor_tech(SENSOR_NAME)
 
 # since techs and rulesets are a list we use a loop to get the corresponding tech and ruleset
 # you will probably use this if you want to test with multiple versions of the sensor
+with open(PCAP, "rb") as pcap:
+    pcap_bytes = pcap.read()
 for tech in dalton_sensor_techs:
-    # have to reread the file to send it to multiple instances of the sensor
-    pcap = open(f"{CURRENT_PATH}/api/examples/mocks/test_job.pcap", "rb")
-    files = {"coverage-pcap0": ("test_job.pcap", pcap)}
-    dalton_client.submit_job(
+    # bytes can be sent any number of times; an open file would be read to its end by the first job
+    files = {"coverage-pcap0": (os.path.basename(PCAP), pcap_bytes)}
+    jid = dalton_client.submit_job(
         sensor_tech=tech,
         prod_ruleset=dalton_rulesets[0],
         configuration=configuration,
         custom_rules=custom_rules,
         files=files,
     )
+    print(f"{tech}: {jid}")

pyproject.toml

@@ -37,6 +37,8 @@ Homepage = "https://github.com/secureworks/dalton"
 
 [project.optional-dependencies]
 testing = [
+    # api/dalton.py, the API client
+    "requests",
     "pytest",
 ]
 devtools = [