Python script to monitor a disk partition and send email alerts when it is almost full.

disk-monitor-v2.py

''' Python script to monitor a disk partition and send email alerts when it is
    almost full. The alert is sent from a Gmail account.

    First install the following dependencies:
    $ pip3 install google-api-python-client 
    $ pip3 install google-auth-httplib2 
    $ pip3 install google-auth-oauthlib

    Then schedule a regular execution of this script with a cron job, e.g.,
    */15 * * * * /usr/bin/python3 /home/pi/disk-alert.py
'''
import os
import re
import pickle
import syslog
import subprocess
from base64 import urlsafe_b64encode
from email.message import EmailMessage
from googleapiclient.discovery import build
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request

# --- START CONFIG --- #
ALERT_RECIPIENT = 'recipient@example.com'

GMAIL_USERNAME = 'gmail_username@gmail.com'
GMAIL_TOKEN_FILE = 'gmail_token.pickle'
GMAIL_CREDENTIALS_FILE = 'gmail_credentials.json'

MONITORED_PARTITION = '/mnt/storage'
DISK_USAGE_THRESHOLD = 95
# --- END CONFIG --- #


class DiskAlertError(Exception):
    pass


class DiskMonitor:
    def __init__(self, partition):
        self.partition = partition

    def _convert_bytes(self, size):
        ''' Converts the input size (in KB) into a more readable format '''
        for x in ['KB', 'MB', 'GB', 'TB']:
            if size < 1024.0:
                return f'{size:3.1f} {x}'
            size /= 1024.0
        return f'{size:,}'

    def get_disk_space(self):
        ''' Returns the available disk space (in KB) and its percentage utilization '''

        # Get the disk utilization data from the command 'df'.
        try:
            out = subprocess.run(
                ['df'], check=True, capture_output=True, text=True
            )
        except subprocess.SubprocessError as e:
            raise DiskAlertError(f'Failed to measure disk space: {e}')

        # Parse the output to get the available space and the percentage of disk
        # utilization of the specified disk partition.
        parsed = re.search(
            f'.* (\d+) .* (\d+)%.*{self.partition}\\n', out.stdout
        )
        if parsed is None:
            raise DiskAlertError(
                'Failed to read disk space: Cannot parse "df" output'
            )

        available = self._convert_bytes(int(parsed.group(1)))
        use = int(parsed.group(2))
        return available, use


class EmailSender:
    def __init__(self, username, token_file, credentials_file):
        self.username = username
        try:
            self.service = self._gmail_authenticate(
                token_file, credentials_file
            )
        except Exception as e:
            raise DiskAlertError(f'Failed to authenticate with Gmail: {e}')

    def _gmail_authenticate(self, token_file, credentials_file):
        ''' Authenticate with GMail through Oauth '''
        credentials = None

        # The file `GMAIL_TOKEN_FILE` stores the user's access and refresh
        # tokens, and is created automatically when the authorization flow
        # completes for the first time.
        if os.path.exists(token_file):
            with open(token_file, 'rb') as token:
                credentials = pickle.load(token)

        # If there are no (valid) credentials available, let the user log in.
        if not credentials or not credentials.valid:
            if credentials and credentials.expired and credentials.refresh_token:
                credentials.refresh(Request())
            else:
                flow = InstalledAppFlow.from_client_secrets_file(
                    credentials_file, ['https://mail.google.com/']
                )
                credentials = flow.run_local_server(port=0)

            # Save the credentials for the next run.
            with open(token_file, 'wb') as token:
                pickle.dump(credentials, token)

        return build('gmail', 'v1', credentials=credentials)

    def _build_message(self, destination, object, body):
        ''' Build an email message '''
        msg = EmailMessage()
        msg['To'] = destination
        msg['From'] = self.username
        msg['subject'] = object
        msg.set_content(body)
        return {'raw': urlsafe_b64encode(msg.as_bytes()).decode()}

    def send_message(self, destination, obj, body):
        ''' Send the email '''
        try:
            return self.service.users().messages().send(
                userId=self.username,
                body=self._build_message(destination, obj, body)
            ).execute()
        except Exception as e:
            raise DiskAlertError(f'Failed to send email alert: {e}')


class DiskAlert:
    def __init__(self, username, token_file, credentials_file, partition, alert_recipient, threshold):
        self.email_sender = EmailSender(username, token_file, credentials_file)
        self.disk_monitor = DiskMonitor(partition)
        self.alert_recipient = alert_recipient
        self.threshold = threshold

    def get_hostname(self):
        ''' Returns the hostname of the machine '''
        out = subprocess.run(['hostname'], capture_output=True, text=True)
        if out.stdout is not None:
            return out.stdout.capitalize().strip()
        else:
            return '[Unknown server]'

    def monitor(self):
        ''' Send an email alert if the disk is almost full'''
        available, use = self.disk_monitor.get_disk_space()

        subject = f'{self.get_hostname()} disk space alert'
        content = (
            f'Disk space utilization over {use}%.\n'
            f'Only {available} left.'
        )
        print(subject)
        print(content)

        if use >= self.threshold:
            self.email_sender.send_message(
                self.alert_recipient, subject, content
            )


if __name__ == '__main__':
    try:
        DiskAlert(
            GMAIL_USERNAME,
            GMAIL_TOKEN_FILE,
            GMAIL_CREDENTIALS_FILE,
            MONITORED_PARTITION,
            ALERT_RECIPIENT,
            DISK_USAGE_THRESHOLD
        ).monitor()
    except DiskAlertError as e:
        print(e)
        syslog.syslog(syslog.LOG_ALERT, str(e))