#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# smartd_test.py
#
# Copyright (C) 2019-2021 Franco Masotti (franco \D\o\T masotti {-A-T-} tutanota \D\o\T com)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
r"""Run S.M.A.R.T tests on hard drives."""

import json
import os
import pathlib
import re
import shlex
import subprocess
import sys

import fpyutils
import yaml


class UserNotRoot(Exception):
    """The user running the script is not root."""


def get_disks() -> list:
    r"""Scan all the disks."""
    disks = list()
    for d in pathlib.Path('/dev/disk/by-id').iterdir():
        # Ignore disks ending with part-${integer} to avoid duplicates (names
        # corresponding to partitions of the same disk).
        disk = str(d)
        if re.match('.+-part[0-9]+$', disk) is None:
            try:
                ddict = json.loads(
                    subprocess.run(
                        shlex.split('smartctl --capabilities --json ' +
                                    shlex.quote(disk)),
                        capture_output=True,
                        check=False,
                        shell=False,
                        timeout=30).stdout)
                try:
                    # Check for smart test support.
                    if ddict['ata_smart_data']['capabilities'][
                            'self_tests_supported']:
                        disks.append(disk)
                except KeyError:
                    pass
            except subprocess.TimeoutExpired:
                print('timeout for ' + disk)
            except subprocess.CalledProcessError:
                print('device ' + disk +
                      ' does not support S.M.A.R.T. commands, skipping...')

    return disks


def disk_ready(disk: str, busy_status: int = 249) -> bool:
    r"""Check if the disk is ready."""
    # Raises a KeyError if disk has not S.M.A.R.T. status capabilities.
    ddict = json.loads(
        subprocess.run(shlex.split('smartctl --capabilities --json ' +
                                   shlex.quote(disk)),
                       capture_output=True,
                       check=True,
                       shell=False,
                       timeout=30).stdout)
    if ddict['ata_smart_data']['self_test']['status']['value'] != busy_status:
        return True
    else:
        return False


def run_test(disk: str, test_length: str = 'long') -> str:
    r"""Run the smartd test."""
    return subprocess.run(
        shlex.split('smartctl --test=' + shlex.quote(test_length) + ' ' +
                    shlex.quote(disk)),
        capture_output=True,
        check=True,
        shell=False,
        timeout=30).stdout


if __name__ == '__main__':
    if os.getuid() != 0:
        raise UserNotRoot

    configuration_file = shlex.quote(sys.argv[1])
    config = yaml.load(open(configuration_file, 'r'), Loader=yaml.SafeLoader)

    # Do not prepend '/dev/disk/by-id/'.
    disks_to_check = shlex.quote(sys.argv[2])
    disks_available = get_disks()

    for d in config['devices']:
        dev = '/dev/disk/by-id/' + d
        if config['devices'][d]['enabled'] and dev in disks_available:
            if disks_to_check == 'all' or disks_to_check == d:
                if disk_ready(dev, config['devices'][d]['busy_status']):
                    print('attempting ' + d + ' ...')
                    message = run_test(
                        dev, config['devices'][d]['test']).decode('utf-8')
                    print(message)
                    if config['devices'][d]['log']:
                        if config['notify']['gotify']['enabled']:
                            m = config['notify']['gotify'][
                                'message'] + ' ' + d + '\n' + message
                            fpyutils.notify.send_gotify_message(
                                config['notify']['gotify']['url'],
                                config['notify']['gotify']['token'], m,
                                config['notify']['gotify']['title'],
                                config['notify']['gotify']['priority'])
                        if config['notify']['email']['enabled']:
                            fpyutils.notify.send_email(
                                message,
                                config['notify']['email']['smtp_server'],
                                config['notify']['email']['port'],
                                config['notify']['email']['sender'],
                                config['notify']['email']['user'],
                                config['notify']['email']['password'],
                                config['notify']['email']['receiver'],
                                config['notify']['email']['subject'])
                else:
                    # Drop test requests if a disk is running a test in a particular moment.
                    # This avoid putting the disks under too much stress.
                    print('disk ' + d + ' not ready, checking the next...')