implement patch manager plus data pipeline for its dashboard

setup_test.sh

@@ -0,0 +1,49 @@
+#!/bin/bash
+
+set -eu
+
+pushd "$(dirname "$0")" > /dev/null
+
+# set env variables from .env file
+set -a && source .env && set +a
+
+golang_migrate_dir="./golang-migrate/"
+
+if [ ! -d "$golang_migrate_dir" ]; then
+	echo "golang-migrate not found. downloading..."
+
+	mkdir -p $golang_migrate_dir
+	pushd ./golang-migrate > /dev/null
+	curl -L https://github.com/golang-migrate/migrate/releases/download/v${GOLANG_MIGRATE_VERSION}/migrate.linux-amd64.tar.gz | tar xvz
+	popd > /dev/null
+
+	echo "golang-migrate installed at $golang_migrate_dir"
+fi
+
+pm_tools=(
+	"tenable"
+	"pmp"
+)
+
+for i in "${pm_tools[@]}"; do
+	echo "Setting up test database for $i..."
+
+	sql=$(cat <<-END
+CREATE DATABASE IF NOT EXISTS ${i}_test;
+GRANT ALL PRIVILEGES ON ${i}_test.* TO '$MARIADB_USER'@'$MARIADB_HOST' IDENTIFIED BY '$MARIADB_PWD';
+FLUSH PRIVILEGES;
+END
+)
+
+
+	echo "Migrating database for $i..."
+
+	echo -e "[client]\nuser=root\npassword=$MARIADB_ROOT_PWD" | sudo mariadb --defaults-extra-file=/dev/stdin -e "$sql"
+	./golang-migrate/migrate -path "src/$i/migrations" -database "mysql://$MARIADB_USER:$MARIADB_PWD@tcp($MARIADB_HOST:$MARIADB_PORT)/${i}_test" up
+
+
+	echo "Test database created and migrated for $i!"
+done
+
+popd >> /dev/null;
+

src/pmp/auth.py

@@ -0,0 +1,24 @@
+import requests
+
+from src.pmp.constants import (
+    PMP_API_URL,
+    PMP_CLIENT_ID,
+    PMP_CLIENT_SECRET,
+    PMP_REFRESH_TOKEN,
+)
+
+
+def request_access_token() -> str:
+    res = requests.post(
+        f"https://accounts.zoho.eu/oauth/v2/token",
+        params={
+            "client_id": PMP_CLIENT_ID,
+            "client_secret": PMP_CLIENT_SECRET,
+            "refresh_token": PMP_REFRESH_TOKEN,
+            "grant_type": "refresh_token",
+        },
+    )
+
+    j = res.json()
+
+    return j["access_token"]

src/pmp/constants.py

@@ -0,0 +1,8 @@
+import os
+
+PMP_API_URL = "https://patch.manageengine.eu"
+PMP_CLIENT_ID = os.getenv("PMP_CLIENT_ID")
+PMP_CLIENT_SECRET = os.getenv("PMP_CLIENT_SECRET")
+PMP_REFRESH_TOKEN = os.getenv("PMP_REFRESH_TOKEN")
+
+PMP_API_ABSENT_VALUE = "--"

src/pmp/context.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+from logging import Logger
+
+from mariadb import ConnectionPool
+
+
+@dataclass
+class PmpPipelineContext:
+    pool: ConnectionPool
+    access_token: str
+    logger: Logger | None

src/pmp/exceptions.py

@@ -0,0 +1 @@
+class UnexpectedApiResponse(Exception): ...

src/pmp/load_patches.py

@@ -0,0 +1,80 @@
+from typing import Any
+from mariadb import mariadb
+
+from src.pmp.constants import PMP_API_URL
+from src.pmp.context import PmpPipelineContext
+from src.pmp.paginate import paginate
+from src.util.database import wait_for_pool_connection
+
+__INSERT_PATCH_SQL = """
+INSERT INTO patches(patch_id, installed_count, missing_count, severity, patch_name, patch_description, release_date)
+VALUES (?, ?, ?, ?, ?, ?, FROM_UNIXTIME(? / 1000))
+ON DUPLICATE KEY UPDATE installed_count   = ?,
+                        missing_count     = ?,
+                        severity          = ?,
+                        patch_name        = ?,
+                        patch_description = ?,
+                        release_date = FROM_UNIXTIME(? / 1000);
+"""
+
+__severity_enum_values = ("UNRATED", "LOW", "MODERATE", "IMPORTANT", "CRITICAL")
+
+
+def load_patches(ctx: PmpPipelineContext):
+    if ctx.logger:
+        ctx.logger.info("loading patch data to database")
+
+    paginate(
+        f"{PMP_API_URL}/api/1.4/patch/allpatches",
+        headers={"Authorization": f"Bearer {ctx.access_token}"},
+        on_page_fetched=__load_page_to_db,
+        args=(ctx,),
+        max_workers=2,
+    )
+
+    if ctx.logger:
+        ctx.logger.info("all patch data loaded into database.")
+
+
+def __load_page_to_db(
+    page: dict[str, Any],
+    page_number: int,
+    ctx: PmpPipelineContext,
+):
+    conn = wait_for_pool_connection(ctx.pool)
+    cursor = conn.cursor()
+
+    try:
+        conn.begin()
+
+        data = [
+            (
+                int(patch.get("patch_id")),
+                int(patch.get("installed")),
+                int(patch.get("missing")),
+                __severity_enum_values[int(patch.get("severity"))],
+                patch.get("patch_name"),
+                patch.get("patch_description"),
+                int(patch.get("patch_released_time")),
+                int(patch.get("installed")),
+                int(patch.get("missing")),
+                __severity_enum_values[int(patch.get("severity"))],
+                patch.get("patch_name"),
+                patch.get("patch_description"),
+                int(patch.get("patch_released_time")),
+            )
+            for patch in page["message_response"]["allpatches"]
+        ]
+
+        cursor.executemany(__INSERT_PATCH_SQL, data)
+
+        conn.commit()
+    except Exception as e:
+        if ctx.logger:
+            ctx.logger.error(
+                f"an error occurred when loading patch data to database at page {page_number}",
+                e,
+            )
+        conn.rollback()
+    finally:
+        conn.close()

src/pmp/load_systems.py

@@ -0,0 +1,91 @@
+from typing import Any
+from mariadb import mariadb
+
+from src.pmp.constants import PMP_API_ABSENT_VALUE, PMP_API_URL
+from src.pmp.context import PmpPipelineContext
+from src.pmp.paginate import paginate
+from src.util.database import wait_for_pool_connection
+
+__INSERT_SYSTEM_SQL = """
+INSERT INTO systems(resource_id, service_pack, health_status, os_platform, os_name, requires_restart)
+VALUES (?, ?, ?, ?, ?, ?)
+ON DUPLICATE KEY UPDATE service_pack  = ?,
+                        health_status = ?,
+                        os_platform   = ?,
+                        os_name       = ?,
+                        requires_restart = ?;
+"""
+
+__resource_health_status_enum_values = [
+    "UNKNOWN",
+    "HEALTHY",
+    "VULNERABLE",
+    "HIGHLY_VULNERABLE",
+]
+
+
+def load_systems(ctx: PmpPipelineContext):
+    if ctx.logger:
+        ctx.logger.info("loading all system data into database using 2 workers.")
+
+    paginate(
+        f"{PMP_API_URL}/api/1.4/patch/allsystems",
+        headers={"Authorization": f"Bearer {ctx.access_token}"},
+        on_page_fetched=__load_page_to_db,
+        args=(ctx,),
+        max_workers=2,
+    )
+
+    if ctx.logger:
+        ctx.logger.info("all system data loaded into database")
+
+
+def __load_page_to_db(page: dict[str, Any], page_number: int, ctx: PmpPipelineContext):
+    conn = wait_for_pool_connection(ctx.pool)
+    try:
+        cursor = conn.cursor()
+
+        conn.begin()
+
+        data = []
+        for system_info in page["message_response"]["allsystems"]:
+            health_status_value = system_info.get("resource_health_status")
+            if not health_status_value or health_status_value == PMP_API_ABSENT_VALUE:
+                health_status_value = 0
+
+            reboot_req_status = system_info.get(
+                "resourcetorebootdetails.reboot_req_status"
+            )
+
+            requires_restart = (
+                reboot_req_status and reboot_req_status != PMP_API_ABSENT_VALUE
+            )
+
+            data.append(
+                (
+                    system_info.get("resource_id"),
+                    system_info.get("service_pack"),
+                    __resource_health_status_enum_values[health_status_value],
+                    system_info.get("os_platform_name"),
+                    system_info.get("os_name"),
+                    requires_restart,
+                    system_info.get("service_pack"),
+                    __resource_health_status_enum_values[health_status_value],
+                    system_info.get("os_platform_name"),
+                    system_info.get("os_name"),
+                    requires_restart,
+                )
+            )
+
+        cursor.executemany(__INSERT_SYSTEM_SQL, data)
+
+        conn.commit()
+    except Exception as e:
+        if ctx.logger:
+            ctx.logger.error(
+                f"an error occurred when loading system data to database at page {page_number}",
+                e,
+            )
+        conn.rollback()
+    finally:
+        conn.close()

src/pmp/migrations/1_initial.up.sql

@@ -0,0 +1,28 @@
+BEGIN;
+
+CREATE TABLE IF NOT EXISTS systems
+(
+    resource_id      BIGINT                                                         NOT NULL,
+    service_pack     TEXT,
+    health_status    ENUM ('UNKNOWN', 'HEALTHY', 'VULNERABLE', 'HIGHLY_VULNERABLE') NOT NULL,
+    os_platform      TEXT                                                           NOT NULL,
+    os_name          TEXT                                                           NOT NULL,
+    requires_restart BOOL                                                           NOT NULL,
+
+    CONSTRAINT pk_systems PRIMARY KEY (resource_id)
+);
+
+CREATE TABLE IF NOT EXISTS patches
+(
+    installed_count   INTEGER                                                      NOT NULL,
+    missing_count     INTEGER                                                      NOT NULL,
+    patch_id          INTEGER                                                      NOT NULL,
+    severity          ENUM ('UNRATED', 'LOW', 'MODERATE', 'IMPORTANT', 'CRITICAL') NOT NULL,
+    patch_name        TEXT                                                         NOT NULL,
+    patch_description TEXT                                                         NOT NULL,
+    release_date      DATETIME                                                     NOT NULL,
+
+    CONSTRAINT pk_patches PRIMARY KEY (patch_id)
+);
+
+COMMIT

src/pmp/paginate.py

@@ -0,0 +1,55 @@
+import math
+from collections.abc import Callable
+from concurrent.futures.thread import ThreadPoolExecutor
+from typing import Any
+
+import requests
+
+
+def paginate(
+    url: str,
+    headers: dict[str, str],
+    on_page_fetched: Callable[..., None],
+    args: tuple[Any] = ((),),
+    max_workers: int = 30,
+):
+    first_page = __fetch_page(url, 1, headers)
+    assert first_page is not None
+
+    total_pages = math.ceil(
+        first_page["message_response"]["total"]
+        / first_page["message_response"]["limit"]
+    )
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        executor.submit(on_page_fetched, first_page, 1, *args)
+
+        # since the first page is already fetched, we need to fetch one less page
+        for i in range(total_pages - 1):
+            # i is zero-based, so +1 to make it one-based, and another +1 to skip the first page
+            executor.submit(
+                __fetch_page,
+                url,
+                i + 2,
+                headers,
+                on_page_fetched,
+                args,
+            )
+
+
+def __fetch_page(
+    url: str,
+    page: int,
+    headers: dict[str, str],
+    cb: Callable[..., None] | None = None,
+    args: tuple[Any] = ((),),
+):
+    try:
+        res = requests.get(f"{url}", params={"page": page}, headers=headers)
+        page_json = res.json()
+        if cb:
+            cb(page_json, page, *args)
+        else:
+            return page_json
+    except Exception as e:
+        print(e)