#!/usr/bin/env python3

# Public domain.
#
# Like any other piece of software (and information generally),
# this script comes with NO WARRANTY.
#
# This is version 20191221 of zfs-auto-snapshot.py3.
#
# See <URL:https://su.bze.ro/software/zfs-auto-snapshot.html>
# for further details and usage instructions.
#
# -- Toby Betts <toby at su dot bze dot ro>

import argparse
import datetime
import subprocess
import sys

dt_parse_fmt    = '%Y %b %d %H:%M'
dt_snapshot_fmt = '%Y-%m-%d_%H%MZ'
now             = datetime.datetime.utcnow()
snapshot_prefix = 'zfs-auto-snap'
zfs_bin         = '/sbin/zfs'
zpool_bin       = '/sbin/zpool'

def die(e, msg):
  out = "fatal: %s: %s" % (sys.argv[0], msg)
  print(out, file=sys.stderr)
  sys.exit(e)

def create_snapshots(snapshots, datasets):
  if (args.verbose):
    print('    start create_snapshots')

  label = '-'.join((snapshot_prefix, args.label, now.strftime(dt_snapshot_fmt)))

  for dataset in datasets:
    snapshot_name            = "%s@%s" % (dataset['name'], label)
    snapshots_existing       = [x for x in snapshots if dataset['name'] == x['dataset']]
    snapshots_existing_names = [x['snapshot'] for x in snapshots_existing]

    if (label in snapshots_existing_names):
      die(100, "snapshot '%s' already exists" % (snapshot_name))
    snapshot_status = syscall([zfs_bin, 'snapshot', snapshot_name])
    snapshot        = { 'dataset': dataset['name'], 'ctime': now, 'snapshot': label }
    snapshots.append(snapshot)

  if (args.verbose):
    print('    end create_snapshots')

def delete_snapshots(snapshots, datasets):
  if (args.verbose):
    print('    start delete_snapshots')

  for dataset in datasets:
    snapshots_dataset = [x for x in snapshots if dataset == x['dataset']]
    snapshots_dataset.sort(key=lambda x: x['ctime'])

    if (int(args.keep) < len(snapshots_dataset)):
      for i in range(len(snapshots_dataset) - int(args.keep)):
        snapshot      = snapshots_dataset[i]
        snapshot_name = "%s@%s" % (snapshot['dataset'], snapshot['snapshot'])
        syscall([zfs_bin, 'destroy', snapshot_name])

  if (args.verbose):
    print('    end delete_snapshots')

def eval_zpool_state(zpools):
  if (args.verbose):
    print('    start eval_zpool_state')

  for zpool in zpools:
    if ('scrub in progress' == zpool['scan'][:17]):
      zpool['snapstate'] = 'scrubbing'
    elif (zpool['state'] not in ['ONLINE', 'DEGRADED']):
      zpool['snapstate'] = 'offline'
    else:
      zpool['snapstate'] = 'ok'

  if (args.verbose):
    print('    end eval_zpool_state')

def parse_datasets(s):
  if (args.verbose):
    print('    start parse_datasets')

  datasets = []
  for line in s.split('\n'):
    if (0 == len(line)):
      continue
    line    = line.strip()
    fields  = line.split()
    slash   = fields[0].find('/')
    if (-1 == slash):
      slash = len(fields[0])
    dataset = {}
    dataset['zpool'] = fields[0][:slash]
    dataset['name']  = fields[0]
    if ('true' == fields[1]):
      dataset['auto_snapshot'] = True
    else:
      dataset['auto_snapshot'] = False
    datasets.append(dataset)

  if (args.verbose):
    print('    end parse_datasets')
  return datasets

def parse_snapshots(s):
  if (args.verbose):
    print('    start parse_snapshots')

  snapshots = []
  prefix    = '-'.join((snapshot_prefix, args.label))

  for line in s.split('\n'):
    if (0 == len(line)):
      continue
    snapshot       = {}
    line           = line.strip()
    fields         = line.split()
    name_segments  = fields[0].split('@')
    snapshot_str   = ' '.join((fields[5], fields[2], fields[3], fields[4]))
    snapshot_ctime = datetime.datetime.strptime(snapshot_str, dt_parse_fmt)

    snapshot['dataset']  = name_segments[0]
    snapshot['snapshot'] = name_segments[1]
    snapshot['ctime']    = snapshot_ctime
    if (snapshot['snapshot'].startswith(prefix)):
      snapshots.append(snapshot)

  if (args.verbose):
    print('    end parse_snapshots')
  return snapshots

def parse_zpool_status(s):
  if (args.verbose):
    print('    start parse_zpool_status')

  zpools = []
  zpool  = {}
  for line in s.split('\n'):
    line = line.strip()
    if ('pool: ' == line[:6]):
      if ('scan' in zpool):
        zpools.append(zpool)
        zpool = {}
      zpool['name'] = line[6:]
    if ('state: ' == line[:7]):
      zpool['state'] = line[7:]
    if ('scan: ' == line[:6]):
      zpool['scan'] = line[6:]

  zpools.append(zpool)

  if (args.verbose):
    print('    end parse_zpool_status')
  return zpools

def syscall(list):
  if (args.verbose):
    print('    start syscall (' + '_'.join(list) + ')')
  status = subprocess.check_output(list)

  if (args.verbose):
    print('    end syscall')
  return status

def doit():
  if (args.verbose):
    print('  start doit')

  zpool_status = syscall([zpool_bin, 'status'])
  zpools       = parse_zpool_status(zpool_status)
  eval_zpool_state(zpools)

  datasets_list   = syscall([zfs_bin, 'list', '-H', '-t', 'filesystem,volume', '-o', 'name,com.sun:auto-snapshot'])
  datasets        = parse_datasets(datasets_list)
  zpools_ok       = [x for x in zpools if 'ok' == x['snapstate']]
  zpools_ok_names = [x['name'] for x in zpools_ok]
  datasets_ok     = [x for x in datasets if x['zpool'] in zpools_ok_names]
  datasets_auto   = [x for x in datasets_ok if x['auto_snapshot']]
  datasets_names  = [x['name'] for x in datasets_auto]

  snapshots_list  = syscall([zfs_bin, 'list', '-H', '-t', 'snapshot', '-o', 'name,creation'])
  snapshots       = parse_snapshots(snapshots_list)
  snapshots_auto  = [x for x in snapshots if x['dataset'] in datasets_names]

  if ((0 < int(args.keep)) and (not args.delete_only)):
    create_snapshots(snapshots_auto, datasets_auto) # appends new snapshots to snapshots_auto

  if (not args.no_delete):
    delete_snapshots(snapshots_auto, datasets_names)

  if (args.verbose):
    print('  end doit')

def main(argv):
  if (args.verbose):
    print('start main')

  if (not str.isalnum(args.label)):
    die(90, 'label not alphanumeric')

  if (not args.no_delete):
    if (None == args.keep):
      die(91, '--keep not set')
    elif ((not args.force) and (1 > int(args.keep))):
      die(92, 'invalid keep value')

  doit()

  if (args.verbose):
    print('end main')

if __name__ =='__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--verbose',     '-v', action='count',      help='set the verbose flag')
  parser.add_argument('--force',       '-f', action='store_true', help='set the force flag')
  parser.add_argument('--delete-only', '-D', action='store_true', help='do not create new snapshots, only delete existing snapshots')
  parser.add_argument('--no-delete',   '-n', action='store_true', help='do not delete any snapshots; ignores the value of --keep')
  parser.add_argument('--keep',        '-k', type=int,            help='the number of snapshots of a given label to retain')
  parser.add_argument('label',               type=str,            help="the name to use for the snapshots (e.g., 'hourly', 'weekly', etc.)")
  args = parser.parse_args()
  main(args)