From 96fdf9694f23a28c99d009facf6a0174f616af47 Mon Sep 17 00:00:00 2001
From: Idriss Neumann <idriss.neumann@comwork.io>
Date: Mon, 4 Nov 2024 10:34:18 +0100
Subject: [PATCH] Add header support monitors

---
 README.md                |  3 +++
 VERSION                  |  2 +-
 src/tests/test_common.py | 29 +++++++++++++++++++++++++++++
 src/utils/common.py      |  3 +++
 src/utils/monitor.py     | 12 +++++++++---
 5 files changed, 45 insertions(+), 4 deletions(-)
 create mode 100644 src/tests/test_common.py

diff --git a/README.md b/README.md
index c136a58..8df30c2 100644
--- a/README.md
+++ b/README.md
@@ -253,6 +253,9 @@ monitors:
     timeout: 30 # optional (30 seconds if not present)
     username: changeit # optional (no basic auth if not present)
     password: changerit # optional (no basic auth if not present)
+    headers: # optional (no headers if empty)
+      - name: Accept
+        value: application/json
 ```
 
 ## Development / contributions
diff --git a/VERSION b/VERSION
index 7919852..2d2d681 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-4.0.9
+4.0.10
diff --git a/src/tests/test_common.py b/src/tests/test_common.py
new file mode 100644
index 0000000..a7992f1
--- /dev/null
+++ b/src/tests/test_common.py
@@ -0,0 +1,29 @@
+import re
+
+from unittest import TestCase
+
+from utils.common import sanitize_header_name
+
+class TestCommon(TestCase):
+    def init(self, *args, **kwargs):
+        super(TestCommon, self).__init__(*args, **kwargs)
+
+    def test_sanitize_header_name_accept(self):
+        ## Given
+        header = "accept"
+
+        #When
+        result = sanitize_header_name(header)
+
+        ## Then
+        self.assertEqual("Accept", result)
+
+    def test_sanitize_header_name_contenttype(self):
+        ## Given
+        header = "content-type"
+
+        #When
+        result = sanitize_header_name(header)
+
+        ## Then
+        self.assertEqual("Content-Type", result)
diff --git a/src/utils/common.py b/src/utils/common.py
index 9df0175..5f67114 100644
--- a/src/utils/common.py
+++ b/src/utils/common.py
@@ -93,3 +93,6 @@ def is_not_uuid (var):
 _allowed_chars_metric_pattern = re.compile(r'[^a-zA-Z0-9]')
 def sanitize_metric_name(name: str):
     return re.sub(_allowed_chars_metric_pattern, '_', name)
+
+def sanitize_header_name(name: str) -> str:
+    return '-'.join(word.capitalize() for word in name.split('-'))
diff --git a/src/utils/monitor.py b/src/utils/monitor.py
index 4d5f906..efc1ea7 100644
--- a/src/utils/monitor.py
+++ b/src/utils/monitor.py
@@ -11,7 +11,7 @@
 from time import sleep
 from requests.auth import HTTPBasicAuth
 
-from utils.common import is_empty_key, get_or_else, is_not_empty, is_not_empty_key, del_key_if_exists
+from utils.common import is_empty_key, get_or_else, is_not_empty, is_not_empty_key, del_key_if_exists, sanitize_header_name
 from utils.gauge import create_gauge, set_gauge
 from utils.heartbit import WAIT_TIME
 from utils.logger import log_msg
@@ -53,21 +53,27 @@ def check_http_monitor(monitor, gauges):
     expected_contain = get_or_else(monitor, 'expected_contain', None)
     duration = None
     auth = None
+    headers = {}
 
     if is_not_empty_key(monitor, 'username') and is_not_empty_key(monitor, 'password'): 
         auth = HTTPBasicAuth(monitor['username'], monitor['password'])
 
+    if is_not_empty_key(monitor, 'headers'):
+        for header in monitor['headers']:
+            if is_not_empty_key(header, 'name') and is_not_empty_key(header, 'value'):
+                headers[sanitize_header_name(header['name'])] = header['value']
+
     pmonitor = monitor.copy()
     del_key_if_exists(pmonitor, 'username')
     del_key_if_exists(pmonitor, 'password')
 
     try:
         if method == "GET":
-            response = requests.get(monitor['url'], auth=auth, timeout=timeout)
+            response = requests.get(monitor['url'], auth=auth, headers=headers, timeout=timeout)
             duration = response.elapsed.total_seconds()
             set_gauge(gauges['duration'], duration, {**labels, 'kind': 'duration'})
         elif method == "POST":
-            response = requests.post(monitor['url'], auth=auth, timeout=timeout)
+            response = requests.post(monitor['url'], auth=auth, headers=headers, timeout=timeout)
             duration = response.elapsed.total_seconds()
             set_gauge(gauges['duration'], duration, {**labels, 'kind': 'duration'})
         else: