Database

#!/usr/bin/env python3
"""
--------------------------------------------------------------------------------
  @author         : Luiggi Vercotti
  @date           : 2026-10-08
  @version        : 1.4.0
  @description    : Monitor filesystem usage

--------------------------------------------------------------------------------
  Update:
  0.1.0  2026-10-08 - Luiggi Vercotti - cre: Script creation
  0.2.0  2026-10-08 - Luiggi Vercotti - evo: Several mount points, log levels
  0.3.0  2026-10-08 - Luiggi Vercotti - evo: Monitoring exit status
  0.4.0  2026-10-08 - Luiggi Vercotti - evo: Configuration file
  0.5.0  2026-10-08 - Luiggi Vercotti - evo: Command line options
  0.6.0  2026-10-08 - Luiggi Vercotti - evo: Alerting API token
  0.7.0  2026-10-08 - Luiggi Vercotti - evo: Report file
  0.8.0  2026-10-08 - Luiggi Vercotti - evo: Error handling
  0.9.0  2026-10-08 - Luiggi Vercotti - evo: Access control
  1.0.0  2026-10-08 - Luiggi Vercotti - evo: Concurrent executions
  1.1.0  2026-10-08 - Luiggi Vercotti - evo: Report templates
  1.2.0  2026-10-08 - Luiggi Vercotti - evo: Remote hosts, reports synchronization
  1.3.0  2026-10-08 - Luiggi Vercotti - evo: SNMP appliances
  1.4.0  2026-10-08 - Luiggi Vercotti - evo: Measures database

--------------------------------------------------------------------------------
  List of authorized users or group:
    @group:ops

--------------------------------------------------------------------------------
  Concurrent executions:
    @max_instance: 1
    @exit_on_wait: true

--------------------------------------------------------------------------------
  List of mandatory configuration parameters:
    @conf:filesystems|mount_points|str|false
    @conf:thresholds|warning|int|false
    @conf:thresholds|critical|int|false
    @conf:alerting|token|str|true
    @conf:remote|user|str|false
    @conf:remote|key|str|false
    @conf:snmp|community|str|true
    @conf:snmp|version|str|false

--------------------------------------------------------------------------------
  List of execution options and argument parameters:
    @args:host|str|Remote hosts to check over SSH (default: local host)|+||snmp-host||false|false
    @args:snmp-host|str|Appliances to check over SNMP (default: local host)|+||host||false|false
    @args:mount-point|str|Mount points to check (default: from configuration)|+||exclude||false|false
    @args:exclude|str|Mount points to skip|+||||false|false
    @args:warning|int|Warning threshold in percent (default: from configuration)|1||||false|false
    @args:critical|int|Critical threshold in percent (default: from configuration)|1||||false|false
    @args:token|str|Alerting API token (default: from configuration)|1||||false|true
    @args:report|bool|Write a report into the report directory||false|||false|false
    @args:format|choice|Report format (default: text)|1|text,csv,html||report|false|false

--------------------------------------------------------------------------------
  Functioning:
  ---------------
    Checks the usage of each mount point against a warning and a critical
    threshold and exits with a monitoring compatible exit status:
      0: OK, 1: WARNING, 2: CRITICAL, 3: UNKNOWN (a mount point can not be checked)

    Configuration:
      - filesystems::mount_points: Comma separated list of mount points
      - thresholds::warning: Warning threshold in percent
      - thresholds::critical: Critical threshold in percent
      - alerting::token: Alerting API token (secret)
      - remote::user: SSH user name on the remote hosts
      - remote::key: SSH private key file
      - snmp::community: SNMP community (secret)
      - snmp::version: SNMP version (1 or 2c)
      - publish::host, publish::dir: Optional. Remote host and directory where
        the report directory is synchronized over SFTP (same user and key as
        remote hosts), in a subdirectory named after the local host
      - publish::retention: Optional. Number of days the reports are kept
        locally and on the publication host (default: 30)
      - database::db_type, database::database: Optional. Database type
        (sqlite3, pgsql or mysql) and name where the measures are stored
      - database::username, database::password, database::host,
        database::port: Optional. Database connection settings

    Templates:
      - <templatedir>/disk_watch/report.<format>.j2
"""
# ------------------------------------------------------------------------------
#  Initialization of the environment
# ------------------------------------------------------------------------------
import os
import sys
import shutil
import datetime
import shlex
import socket
import scrippy_core
from scrippy_core import logger, ScrippyCoreError
from scrippy_template import template
from scrippy_remote import ScrippyRemoteError
from scrippy_remote.ssh import Ssh
from scrippy_remote.local import list_local_dir, delete_local_file
from scrippy_snmp import ScrippySnmpError
from scrippy_snmp.snmp import Snmp
from scrippy_db.db import Database

OK, WARNING, CRITICAL, UNKNOWN = 0, 1, 2, 3
STATUS = {OK: "OK", WARNING: "WARNING", CRITICAL: "CRITICAL", UNKNOWN: "UNKNOWN"}
EXTENSIONS = {"text": "txt", "csv": "csv", "html": "html"}
REPORT_PATTERN = r"^disk_watch_\d{8}_\d{6}\.(txt|csv|html)$"
# HOST-RESOURCES-MIB::hrStorageTable columns
HR_STORAGE_DESCR = "1.3.6.1.2.1.25.2.3.1.3"
HR_STORAGE_SIZE = "1.3.6.1.2.1.25.2.3.1.5"
HR_STORAGE_USED = "1.3.6.1.2.1.25.2.3.1.6"
# Measures database queries: Placeholders depend on the database type
CREATE_TABLE = "CREATE TABLE IF NOT EXISTS measures (date VARCHAR(19), name VARCHAR(255), percent REAL, status VARCHAR(8))"
SELECT_LAST = {"sqlite3": "SELECT percent FROM measures WHERE name = :name ORDER BY date DESC LIMIT 1",
               "pgsql": "SELECT percent FROM measures WHERE name = %(name)s ORDER BY date DESC LIMIT 1",
               "mysql": "SELECT percent FROM measures WHERE name = %(name)s ORDER BY date DESC LIMIT 1"}
INSERT = {"sqlite3": "INSERT INTO measures VALUES(:date, :name, :usage, :status)",
          "pgsql": "INSERT INTO measures VALUES(%(date)s, %(name)s, %(usage)s, %(status)s)",
          "mysql": "INSERT INTO measures VALUES(%(date)s, %(name)s, %(usage)s, %(status)s)"}

# ------------------------------------------------------------------------------
#  Definition of functions and classes
# ------------------------------------------------------------------------------
def get_mount_points(config, args):
  if args.mount_point:
    return args.mount_point
  mount_points = [mp.strip() for mp in config.get("filesystems", "mount_points").split(",")]
  return [mp for mp in mount_points if mp not in (args.exclude or [])]

def get_thresholds(config, args):
  warning = args.warning if args.warning is not None else config.get("thresholds", "warning", "int")
  critical = args.critical if args.critical is not None else config.get("thresholds", "critical", "int")
  if not 0 &lt; warning &lt; critical &lt;= 100:
    raise ScrippyCoreError(f"Invalid thresholds: warning={warning}, critical={critical}")
  return warning, critical

def get_local_usage(mount_point):
  logger.debug(f"[+] Checking {mount_point}")
  usage = shutil.disk_usage(mount_point)
  return round(usage.used * 100 / usage.total, 1)

def get_remote_usage(ssh, mount_point):
  logger.debug(f"[+] Checking {ssh.host}:{mount_point}")
  result = ssh.exec(command=f"df -P {shlex.quote(mount_point)}", exit_on_error=False)
  if result["exit_code"] != 0:
    raise OSError(" ".join(result["stderr"]) or f"df exit code {result['exit_code']}")
  # Filesystem 1024-blocks Used Available Capacity Mounted on
  fields = result["stdout"][-1].split()
  return round(int(fields[2]) * 100 / int(fields[1]), 1)

def get_snmp_usages(snmp):
  """Returns the usage of each storage of the appliance, indexed by its description."""
  logger.debug(f"[+] Reading storage table of {snmp.host}")
  columns = {}
  for column in (HR_STORAGE_DESCR, HR_STORAGE_SIZE, HR_STORAGE_USED):
    # OID suffix (storage index) =&gt; value
    columns[column] = {oid.rsplit(".", 1)[1]: value for oid, value in snmp.walk(column)}
  usages = {}
  for index, descr in columns[HR_STORAGE_DESCR].items():
    size = int(columns[HR_STORAGE_SIZE].get(index, 0))
    if size &gt; 0:
      usages[descr] = round(int(columns[HR_STORAGE_USED][index]) * 100 / size, 1)
  return usages

def get_snmp_usage(usages, mount_point):
  try:
    return usages[mount_point]
  except KeyError as err:
    raise OSError(f"No such storage: {mount_point}") from err

def check(mount_points, warning, critical, get_usage=get_local_usage, prefix=""):
  results = []
  for mount_point in mount_points:
    name = f"{prefix}{mount_point}"
    try:
      usage = get_usage(mount_point)
    except OSError as err:
      logger.error(f" '-&gt; {name}: {err.strerror or err}")
      results.append((name, None, UNKNOWN))
      continue
    if usage &gt;= critical:
      logger.error(f" '-&gt; {name}: {usage}% [CRITICAL]")
      status = CRITICAL
    elif usage &gt;= warning:
      logger.warning(f" '-&gt; {name}: {usage}% [WARNING]")
      status = WARNING
    else:
      logger.info(f" '-&gt; {name}: {usage}% [OK]")
      status = OK
    results.append((name, usage, status))
  return results

def check_remote_hosts(hosts, mount_points, warning, critical, config):
  results = []
  for host in hosts:
    logger.info(f"[+] Checking remote host: {host}")
    try:
      with Ssh(username=config.get("remote", "user"),
               host=host,
               key=config.get("remote", "key"),
               missing_host_key_policy="reject") as ssh:
        results += check(mount_points, warning, critical,
                         get_usage=lambda mount_point: get_remote_usage(ssh, mount_point),
                         prefix=f"{host}:")
    except ScrippyRemoteError as err:
      logger.error(f" '-&gt; {host}: {err}")
      results += [(f"{host}:{mount_point}", None, UNKNOWN) for mount_point in mount_points]
  return results

def check_snmp_hosts(hosts, mount_points, warning, critical, config):
  results = []
  for host in hosts:
    logger.info(f"[+] Checking SNMP appliance: {host}")
    try:
      snmp = Snmp(host=host,
                  port=161,
                  community=config.get("snmp", "community"),
                  version=config.get("snmp", "version"))
      usages = get_snmp_usages(snmp)
    except ScrippySnmpError as err:
      logger.error(f" '-&gt; {host}: {err}")
      results += [(f"{host}:{mount_point}", None, UNKNOWN) for mount_point in mount_points]
      continue
    results += check(mount_points, warning, critical,
                     get_usage=lambda mount_point: get_snmp_usage(usages, mount_point),
                     prefix=f"{host}:")
  return results

def purge_reports(report_dir, retention):
  """Deletes the local reports older than ``retention`` days (the report date is part of its name)."""
  limit = f"disk_watch_{datetime.datetime.now() - datetime.timedelta(days=retention):%Y%m%d_%H%M%S}"
  logger.info(f"[+] Purging reports older than {retention} days")
  for name in list_local_dir(local_dir=report_dir, file_type="f", pattern=REPORT_PATTERN):
    if name &lt; limit:
      logger.info(f" '-&gt; Deleting {name}")
      delete_local_file(local_file=os.path.join(report_dir, name))

def sync_reports(report_dir, config):
  """Synchronizes the local reports to the publication host: Missing reports are sent, deleted reports are deleted."""
  host = config.get("publish", "host")
  remote_dir = f"{config.get('publish', 'dir')}/{socket.gethostname()}"
  logger.info(f"[+] Synchronizing reports to {host}:{remote_dir}")
  with Ssh(username=config.get("remote", "user"),
           host=host,
           key=config.get("remote", "key"),
           missing_host_key_policy="reject") as ssh:
    ssh.create_remote_dir(remote_dir=remote_dir, create_dirs=True)
    local_reports = set(list_local_dir(local_dir=report_dir, file_type="f", pattern=REPORT_PATTERN))
    remote_reports = set(ssh.list_remote_dir(remote_dir=remote_dir, file_type="f", pattern=REPORT_PATTERN))
    for name in sorted(local_reports - remote_reports):
      logger.info(f" '-&gt; Sending {name}")
      ssh.put_file(local_file=os.path.join(report_dir, name), remote_dir=remote_dir)
    for name in sorted(remote_reports - local_reports):
      logger.info(f" '-&gt; Deleting {name}")
      ssh.delete_remote_file(remote_file=f"{remote_dir}/{name}")
    logger.info(f" '-&gt; {len(local_reports - remote_reports)} sent, {len(remote_reports - local_reports)} deleted, {len(local_reports)} published")

def get_database(config):
  connection = {"db_type": config.get("database", "db_type"),
                "database": config.get("database", "database")}
  for key in ("username", "password", "host"):
    if config.has("database", key):
      connection[key] = config.get("database", key)
  if config.has("database", "port"):
    connection["port"] = config.get("database", "port", "int")
  if "password" in connection:
    # Undeclared configuration key: Its value has to be registered as secret explicitly
    logger.add_secret(connection["password"])
  return Database(**connection)

def store_measures(results, config):
  db_type = config.get("database", "db_type")
  date = f"{datetime.datetime.now():%Y-%m-%d %H:%M:%S}"
  logger.info("[+] Storing measures")
  with get_database(config) as db:
    db.query(query=CREATE_TABLE, commit=True)
    queries = []
    for name, usage, status in results:
      if usage is None:
        continue
      previous = db.query(query=SELECT_LAST[db_type], params={"name": name})
      if previous:
        logger.info(f" '-&gt; {name}: {usage - previous[0][0]:+.1f}% since last check")
      queries.append({"query": INSERT[db_type],
                      "params": {"date": date, "name": name, "usage": usage, "status": STATUS[status]}})
    db.transaction(queries=queries, exit_on_error=True)

def render_report(results, report_format, warning, critical, templatedir):
  params = {"hostname": socket.gethostname(),
            "date": f"{datetime.datetime.now():%Y-%m-%d %H:%M:%S}",
            "warning": warning,
            "critical": critical,
            "status": STATUS[max(status for _, _, status in results)],
            "results": [{"mount_point": mount_point, "usage": usage, "status": STATUS[status]}
                        for mount_point, usage, status in results]}
  renderer = template.Renderer(os.path.join(templatedir, "disk_watch"),
                               f"report.{EXTENSIONS[report_format]}.j2")
  return renderer.render(params)

def write_report(content, report_format, workspace_path, reportdir):
  name = f"disk_watch_{datetime.datetime.now():%Y%m%d_%H%M%S}.{EXTENSIONS[report_format]}"
  tmp_report = os.path.join(workspace_path, name)
  logger.info(f"[+] Writing report: {tmp_report}")
  with open(tmp_report, mode="w", encoding="utf-8") as report:
    report.write(content)
  report_dir = os.path.join(reportdir, "disk_watch")
  os.makedirs(report_dir, exist_ok=True)
  logger.info(f"[+] Publishing report into {report_dir}")
  return shutil.move(tmp_report, os.path.join(report_dir, name))

# ------------------------------------------------------------------------------
#  Main processing
# ------------------------------------------------------------------------------
def main():
  with scrippy_core.ScriptContext(retention=100) as _context:
    config = _context.config
    args = _context.args
    token = args.token or config.get("alerting", "token")
    logger.debug(f"[+] Alerting API token: {token}")
    warning, critical = get_thresholds(config, args)
    logger.log_section("disk_watch", f"Thresholds: warning={warning}% critical={critical}%")
    mount_points = get_mount_points(config, args)
    if args.host:
      results = check_remote_hosts(args.host, mount_points, warning, critical, config)
    elif args.snmp_host:
      results = check_snmp_hosts(args.snmp_host, mount_points, warning, critical, config)
    else:
      results = check(mount_points, warning, critical)
    if config.has("database", "db_type"):
      store_measures(results, config)
    if args.report:
      report_format = args.format or "text"
      content = render_report(results, report_format, warning, critical, _context.env["templatedir"])
      report = write_report(content, report_format, _context.workspace_path, _context.env["reportdir"])
      if config.has("publish", "host"):
        retention = config.get("publish", "retention", "int") if config.has("publish", "retention") else 30
        purge_reports(os.path.dirname(report), retention)
        sync_reports(os.path.dirname(report), config)
    status = max(status for _, _, status in results)
    logger.info(f"[+] Status: {STATUS[status]}")
    if status != OK:
      sys.exit(status)

# ------------------------------------------------------------------------------
#  Entry point
# ------------------------------------------------------------------------------
if __name__ == '__main__':
  main()