#!/usr/bin/python3
#
# Sleepdisk (Idle Disk Sleeper) Version 1.0
#
# Copyright 2023 Brandyn Webb
# 
# This file is part of Sleepdisk.
# 
# Sleepdisk 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.
# 
# Sleepdisk 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 
# Sleepdisk. If not, see <https://www.gnu.org/licenses/>. 
# 


#-------- Config ------------

#
# Idle times maps full /dev path or final dev name to idle time in minutes.
#
# List only full drives here, not partitions.
#
# Sym links will be resolved to final dev path, then last element used
#  to look up device in /proc/diskstats.
#
# Disks will be slept with "hdparm -y".  Unfortunately that means this
#   needs to be run as root.
#
idle_times = {
        #
        # List the disks you want idled here, with their idle time in minutes.
        #
        # Specify any full /dev paths here, to the full drive (not partition).
        #
        # It's probably best to use a /dev/disk/by-id path in case the device
        #   ordering changes; or a by-path ID if, say, you want hot swapping
        #   to follow the physical location as opposed to the disk itself.
        #
        # Beware if you have smartd running in the default config, it touches
        #   the disks every 30 minutes (in Ubuntu Server anyway), so you need
        #   to pick a time lower than that (with safety margin for the cycle_time)
        #   or smartd will keep the disks active indefinitely.  (Even so, expect
        #   smartd to extend the active time by 15+/-15 minutes...)
        #

        # THESE ARE EXAMPLES -- replace them with yours:
        "/dev/disk/by-id/scsi-SATA_WDC_WD80EFAX-68K_VGG0H3EG"   : 25,   # 25 minutes
        "/dev/disk/by-id/scsi-SATA_WDC_WD80EFAX-68K_VDKZUA6P"   : 25,
        "/dev/disk/by-id/scsi-SATA_WDC_WD80EAZZ-00B_WD-CA0X2TZK": 25,
    }

cycle_time    = 121                         # Sleep time, in seconds, between updates/checks.
sleep_command = ("/usr/sbin/hdparm", "-y")  # Path from idle_times will be last parameter to this command.

#-------- End Config / Begin arg parsing -------

from sys      import argv, exit
from datetime import datetime

debug   = 0
verbose = 0

def log(msg):
    print(f"{datetime.now()}: {msg}", flush=True)

log("--------- Starting up -----------")

args = argv[1:]
while args:
    arg = args.pop(0)
    if arg == '-v':
        verbose += 1
    elif arg == '-D':
        debug += 1
    else:
        log(f"Unknown arg: {arg}")
        exit(1)

#-------- End Arg parsing -------

if debug:
    log("WARNING: Running in debug mode.")
    cycle_time = 2

from time       import time, sleep
from os.path    import realpath, basename
from subprocess import call

class Disk(object):

    def __init__(self, path, idle_time):

        assert path.startswith('/dev/')

        self.path      = path           # The original disk path specifier
        self.idle_time = idle_time*60   # idle_time in seconds.
        self.dev       = None           # The canonical name of the disk device, once known.
        self.active    = True           # Assume disk is running by default. (No harm sleeping it if already is.)
        self.last_log  = None           # Last log message, just to avoid repetitions.

        self.clear_stats()
        self.find_device()              # Sets self.dev

        if verbose:
            self.log(f"idle_time={self.idle_time}")

    def clear_stats(self):
        """Resets all "last seen" stats to unknown/default state.
        """
        self.events      = 0
        self.last_active = time()

    def set_active(self, active):
        """This is called to declare whether the disk appears to be
            recently-enough active or not.  If it changes, it will
            be logged.  If it goes from True to False, the disk will
            be put into standby.
        """
        if active == self.active:
            return

        self.active = active

        if active:
            self.log("Disk awake.")
        else:
            self.log("Sleeping disk.")
            try:
                call(sleep_command + (self.path,))  # This is where we actually sleep the disk.
            except Exception as e:
                self.log(f"call() failed ({e})")

    def update(self, stats, now):
        """This just examines the current state of this disk, updates
            last-observed state, and calls set_active accordingly.

        stats is just /proc/diskstats pre-loaded as a dict of Stats objects.

        now is time()
        """
        # We could examine the stats to see if any of the counters have
        #  gone down, which would indicate the device has been moved, but
        #  re-finding the device each time should be pretty fast since
        #  it's just hitting a virtual file system (/dev)...
        self.find_device()

        if self.dev not in stats:
            self.log(f"{self.dev} not present.")
            return

        s = stats[self.dev]

        try:
            #
            # It turns out that hdparm -y (and even -C) increment the read
            #  count, so that makes it look like the disk just went active
            #  again.  Hopefully the sector count will be representative of
            #  actual activity...
            #
            #events = s.reads + s.writes     # This doesn't work on my machine...
            #
            events = s.reads_sectors + s.writes_sectors

            if events != self.events:    # More events means active now-ish.
                self.events      = events
                self.last_active = now
                if verbose > 1:
                    self.log("Activity detected.", force=True)

            self.set_active(now - self.last_active < self.idle_time)

        except Exception as e:

            self.log(f"Update failed: {e}")

    def find_device(self):
        """Sets self.dev to, e.g., 'sda' or similar device name as
        it would appear in /proc/diskstats.  This is called any time
        we suspect the disk might have been moved (hot-swapping disks;
        USB disks, etc).
        """
        try:
            dev = basename(realpath(self.path))

            if dev != self.dev:
                if self.dev:
                    self.log(f"{self.dev} -> {dev}")
                    self.clear_stats()
                else:
                    self.log(f"using {dev}")
                self.dev = dev
        except Exception as e:
            self.log(f"Couldn't find device ({e})")
            self.dev = None

    def log(self, msg, force=False):
        if force or msg != self.last_log:
            log(f"{self.path} - {msg}")
            self.last_log = msg

class Stats(object):
    """This is just a container/data object for one line of /proc/diskstats
    """
    props = (   # Properties specified in /proc/diskstats, in order:
        "major",
        "minor",
        "dev",
        "reads",
        "reads_merged",
        "reads_sectors",
        "reads_time",
        "writes",
        "writes_merged",
        "writes_sectors",
        "writes_time",
        "io_in_progress",
        "io_time",
        "io_time_weighted",
    )
    str_props = {"dev"} # All other properties are converted to ints.

    def __init__(self, line):
        try:
            for prop, v in zip(self.props, line.strip().split()):
                if prop in self.str_props:
                    setattr(self, prop, v)
                else:
                    setattr(self, prop, int(v))
        except:
            log(f"Warning: Bad diskstats line ({line!r})")
            self.dev = None
            return

    def __repr__(self):
        return ', '.join(f'{prop}={getattr(self,prop)}' for prop in self.props)

    @classmethod
    def load(cls):
        with open("/proc/diskstats") as fl:
            return {s.dev: s for s in [Stats(line) for line in fl] if s.dev}

disks = [Disk(path, idle_time) for path, idle_time in idle_times.items()]

log(f"Cycle time: {cycle_time} seconds")

while True:

    stats = Stats.load()
    when  = time()

    for disk in disks:
        disk.update(stats, when)

    sleep(cycle_time)

