#!/usr/bin/env python3

import os
import re
import time

HOSTNAME = os.getenv('COLLECTD_HOSTNAME', "localhost")
INTERVAL = float(os.getenv('COLLECTD_INTERVAL', '60.0'))
FILENAME = "/tmp/spitfire_smbspy_status.txt"

SMB_STATUS_REGEX = re.compile(r'<b>Spitfire SMB</b> has <b>(?P<files>[0-9]+)</b> file\(s\) opened by <b>(?P<users>[0-9]+)</b> user\(s\) \(<b>(?P<active>[0-@9]+)</b> active users\)\.')
SMB_NOOPEN_REGEX = re.compile(r'<b>Spitfire SMB</b> has no open files or command failure\.')
HTTP_STATUS_REGEX = re.compile(r'<b>Spitfire HTTP</b> has <b>(?P<transfers>[0-9]+)</b> active transfer\(s\) by <b>(?P<users>[0-9]+)</b> user\(s\)\.')
HTTP_NOOPEN_REGEX = re.compile(r'<b>Spitfire HTTP</b> has no open files or command failure\.')

smb_current_files, smb_current_users, smb_active_users = -1, -1, -1
http_current_transfers, http_current_users = -1, -1

while True:
    # Read values from status file
    if os.path.isfile(FILENAME):
        with open(FILENAME, 'r', encoding='utf-8') as f:
            status_str = f.readlines()
    else:
        status_str = ""

    for line in status_str:
        # Try to match SMB status
        smb_match = SMB_STATUS_REGEX.match(line)
        if smb_match:
            smb_current_files = int(smb_match.group('files'))
            smb_current_users = int(smb_match.group('users'))
            smb_active_users = int(smb_match.group('active'))

        # Try to match SMB failure
        smb_noopen_match = SMB_NOOPEN_REGEX.match(line)
        if smb_noopen_match:
            smb_current_files, smb_current_users, smb_active_users = 0, 0, 0

        # Try to match HTTP status
        http_match = HTTP_STATUS_REGEX.match(line)
        if http_match:
            http_current_transfers = int(http_match.group('transfers'))
            http_current_users = int(http_match.group('users'))

        # Try to match HTTP failure
        http_noopen_match = HTTP_NOOPEN_REGEX.match(line)
        if http_noopen_match:
            http_current_transfers, http_current_users = 0, 0

    print("PUTVAL \"{}/shares/gauge-smb_open_connections\" interval={} N:{}".format(HOSTNAME, INTERVAL, smb_current_users), flush=True)
    print("PUTVAL \"{}/shares/gauge-smb_active_users\" interval={} N:{}".format(HOSTNAME, INTERVAL, smb_active_users), flush=True)
    print("PUTVAL \"{}/shares/gauge-smb_open_files\" interval={} N:{}".format(HOSTNAME, INTERVAL, smb_current_files), flush=True)
    print("PUTVAL \"{}/shares/gauge-http_active_users\" interval={} N:{}".format(HOSTNAME, INTERVAL, http_current_users), flush=True)
    print("PUTVAL \"{}/shares/gauge-http_open_transfers\" interval={} N:{}".format(HOSTNAME, INTERVAL, http_current_transfers), flush=True)

    time.sleep(INTERVAL)