#!/usr/bin/python
"""This program is a set of supplemental commands to be used with snapper"""

__version__ = '1.1.5'

from argparse import RawTextHelpFormatter  # used so we can use \n in argparse descriptions
import argparse  # used for argument parsing

from subprocess import CalledProcessError  # used to catch exceptions
from subprocess import check_output as run  # used to run command to interface with btrfs
import os  # used to check for root and listdir

from tabulate import tabulate  # used for snapperS list


def mountSnapshotRW(path, enableWrite):  # set the snapshot at path to rw or ro
    """Remount a snapshot as ro or rw"""
    with open(os.devnull, 'w') as devnull:  # used so we don't have to write stderr to output
        run("btrfs property set -ts " + path + " ro " + str(not enableWrite).lower(),
            shell=True,
            stderr=devnull)


def rm(args):  # delete functionality
    """Delete a file or folder from a range of snapshots"""
    if args['range'] is not None:  # if they specified a range argument
        startPoint = int(args['range'].split('..')[0])  # first number
        endPoint = int(args['range'].split('..')[1])  # second number

    if '/' not in args['filename']:  # if the user tried to specify a relative path, warn them
        print "Specify the absolute path of the file starting with a /"
        exit(2)

    for snapshot in os.listdir(args['directory']):
        if args['range'] is not None:
            # only limit the start/end points if they user specified them
            if not startPoint <= int(snapshot) <= endPoint:
                continue
        try:
            mountSnapshotRW(args['directory'] + snapshot + "/snapshot/", True)
            # path to the actual snapshot is /.snapshots/[num]/snapshot/
        except CalledProcessError as e:
            if '50' in str(e):
                # when btrfs exits with an exit code of 50, there is no valid subvolume
                if args['verbose']:
                    print "No subvolume found in " + args['directory'] + snapshot + ". Skipping. "
                continue
            print e
            # if it isn't an exit code of 50, then there is some other
            # unrecoverable error, so we exit
            exit(3)
        try:
            with open(os.devnull, 'w') as devnull:
                # used so we don't have to write stderr to output
                if args['recursive']:
                    command = "rm -r "
                else:
                    command = "rm "
                out = run(command + args['directory'] + snapshot + "/snapshot" + args['filename'],
                          shell=True,
                          stderr=devnull)
                if args['verbose']:
                    print out
            if args['verbose']:
                print "Deleted from " + args['directory'] + snapshot
        except CalledProcessError as e:
            if "1" in str(e):  # an exit code 1 means the file doesn't exist
                if args['verbose']:
                    print ("Failed to delete the file in " + args['directory'] +
                           snapshot + "/snapshot/. Check that it exists and" +
                           " is not a directory. Skipping. ")
            else:
                print e  # otherwise there is an unrecoverable error
                exit(4)
        mountSnapshotRW(args['directory'] + snapshot + "/snapshot/", False)
        # set it back to ro once we are done
    try:
        mountSnapshotRW("/", True)
    except CalledProcessError as e:
        print "Error! Failed to remount / as rw. This will require maual intervention to fix. "
        exit(5)


def cat(args):  # cat functionality
    """print the contents of a specified file in a specified snapshot"""
    print run("cat " + args['directory'] + args['snapshot'] + "/snapshot" + args['filename'],
              shell=True)


def backup(args):
    """backup the specified snapshot to a specified file"""
    try:
        with open(os.devnull, 'w') as devnull:  # used to surpress stderr
            # redirect output of btrfs send to a file
            out = run(("btrfs send -v " + args['directory'] + args['snapshot'] +
                       "/snapshot/ > " + args['backup']),
                      stderr=devnull,
                      shell=True)
            if args['verbose']:
                print out
        print "Sucessfully backed up the file. "
    except CalledProcessError as e:
        print e
        print "Error! Failed to backup. Try again."


def restore(args):
    """restore the specified file to the specified mount point"""
    try:
        with open(os.devnull, 'w') as devnull:  # used to surpress stderr
            # pipe the backup into btrfs receive
            out = run("cat " + args['restoreLocation'] + " | btrfs receive -v " + args['backup'],
                      stderr=devnull,
                      shell=True)
            if args['verbose']:
                print out
        print "Sucessfully restored the file. "
    except CalledProcessError as e:
        print e
        print "Error! Failed to restore. Try again."


def listSubvolumes():
    """list all of the snapper subvolumes including information on the sizes of the subvolumes"""
    subvols = run("btrfs subvolume list /", shell=True).splitlines()[1:]  # get a list of all of
    btrfsIDToSnapperID = {}  # the subvolumes that we use to generate the dict to map between them
    for line in subvols:
        btrfsID = int(line.split(' ')[1])  # parse the btrfsID
        snapperID = int(line.split(' ')[-1].split('/')[1])  # parse the snapperID
        btrfsIDToSnapperID[btrfsID] = snapperID  # put the information into the dict
    # get the information about the size of the subvolumes
    storageInfo = run("btrfs qgroup show /", shell=True).splitlines()[2:]
    snapperIDToSpaceTuple = {}  # maps a snapperID to a tuple of (dataSize, uniqueDataSize)
    btrfsIDToSpaceTuple = {}  # maps a btrfsID to a tuple of (dataSize, uniqueDataSize)
    for line in storageInfo:
        # parse the data points by splitting on ' ' and ignoring the blanks
        data = [datum for datum in line.split(' ') if datum != '']
        tuple = (data[1], data[2])  # create the spaceTuple
        btrfsID = int(data[0].split('/')[1])  # parse the btrfsID
        try:  # try because sometimes there will be a KeyError
            snapperID = btrfsIDToSnapperID[btrfsID]  # parse the snapperID
            snapperIDToSpaceTuple[snapperID] = tuple  # add to the snapperID dict
        except KeyError as e:
            pass
        btrfsIDToSpaceTuple[btrfsID] = tuple  # add to the btrfsID dict
    # get data from snapper that we will then supplement
    snapperList = run("snapper list", shell=True).splitlines()
    table = []  # a list of lists that tabulate() will use
    headers = [section for section in snapperList[0].replace('|', '').split(' ')
               if section != '']  # parse the headers
    headers.extend(["Size", "Unique Size"])  # add the headers that we will use
    for line in snapperList[2:]:  # for line in list of data points
        lineArr = line.split(' | ')  # split on the ' | ' in order to get the data points
        snapperID = int(lineArr[1])  # parse the snapperID
        try:
            spaceTuple = snapperIDToSpaceTuple[snapperID]
        except:
            spaceTuple = ("N/A", "N/A")  # for snapperID == 0
        lineArr.append(" ")  # for the userdata column
        lineArr.extend(list(spaceTuple))  # add the space data
        table.append(lineArr)  # append our line to the table
    # print a table with the headers and in format of psql
    print tabulate(table, headers=headers, tablefmt="psql")


def main():
    """main function containing argument parsing code"""
    subcommandsHelp = {}  # dict that contains the help information for the subcommands
    subcommandsHelp['rm'] = ("Delete a specified file from either a range of snapshots " +
                             "or from all snapshots. ")
    subcommandsHelp['cat'] = "Read a specified file from a specified snapshot. "
    subcommandsHelp['backup'] = "Backup a specified snapshot to a file via btrfs send. "
    subcommandsHelp['restore'] = "Restore a snapshot from a file generated with snapperS backup. "
    subcommandsHelp['list'] = ("A more comprehensive version of snapper list that includes "
                               "information on space usage. ")

    # create the parser
    parser = argparse.ArgumentParser(description=("snapperS: A variety of supplemental "
                                                  "snapper subcommands"),
                                     # so we can use \n and \t
                                     formatter_class=RawTextHelpFormatter)
    parser.add_argument('-d',  # for if the user has changed the snapper mount point
                        '--directory',
                        help='Directory containing the snapshots',
                        default='/.snapshots/')
    parser.add_argument('-v',  # enable debug logging
                        '--verbose',
                        action='store_true',
                        default=False,
                        help=("Enable verbose logging. If you are experiencing difficulties " +
                              "with this program, try with -v for debugging. "))
    parser.add_argument('-V',
                        '--version',
                        action="version",
                        version="%(prog)s " + __version__)

    subparsers = parser.add_subparsers(title='Subcommands',  # list of subcommands
                                       description='\n'.join([subcommandsHelp[key]
                                                              for key in subcommandsHelp]))

    # rm subcommand
    parser_rm = subparsers.add_parser('rm',
                                      description=subcommandsHelp['rm'])
    parser_rm.set_defaults(which='rm')
    parser_rm.add_argument('-f',
                           '--filename',
                           required=True,
                           metavar="~/largeFile.img",
                           help="Delete a file from all past snapshots. ")
    parser_rm.add_argument('-r',
                           '--range',
                           metavar="1..42",
                           help=("The range of snapshots to delete the file from in the " +
                                 "form of startPoint..endPoint (e.g. 2..5)"))
    parser_rm.add_argument('--recursive',
                           action='store_true',
                           default=False,
                           help="Delete recursively (i.e. a folder)")

    # cat subcommand
    parser_cat = subparsers.add_parser('cat',
                                       description=subcommandsHelp['cat'])
    parser_cat.set_defaults(which='cat')
    parser_cat.add_argument('-f',
                            '--filename',
                            required=True,
                            metavar="~/file.txt",
                            help="The file to cat")
    parser_cat.add_argument('-s',
                            '--snapshot',
                            required=True,
                            metavar="42",
                            help="The snapshot to view")

    # backup subcommand
    parser_backup = subparsers.add_parser('backup',
                                          description=(subcommandsHelp['backup'] +
                                                       "It is recommended to compress this file." +
                                                       "\n\t-In order to restore this file, run " +
                                                       "`cat backup | btrfs receive " +
                                                       "/mnt/subvol`"
                                                       "\n\t-If you want to sync your backups to "
                                                       "another BTRFS filesystem, ButterSink is "
                                                       "better suited for that purpose. "),
                                          formatter_class=RawTextHelpFormatter)
    parser_backup.set_defaults(which='backup')
    parser_backup.add_argument('-b',
                               '--backup',
                               required=True,
                               metavar="~/BTRFS_Backup.send",
                               help="The location to store the backup")
    parser_backup.add_argument('-s',
                               '--snapshot',
                               required=True,
                               metavar="42",
                               help="The number of the snapshot you want to backup")

    # restore subcommand
    parser_restore = subparsers.add_parser('restore',
                                           description=subcommandsHelp['restore'])
    parser_restore.set_defaults(which='restore')
    parser_restore.add_argument('-b',
                                '--backup',
                                required=True,
                                metavar="~/BTRFS_Backup.send",
                                help="The location of the backup. ")
    parser_restore.add_argument('-r',
                                '--restoreLocation',
                                required=True,
                                metavar="~/newRestoredSubvolume/",
                                help="The path to where you want to restore the backup. ")

    # list subcommand
    parser_list = subparsers.add_parser('list',
                                        description=subcommandsHelp['list'])
    parser_list.set_defaults(which='list')

    # parse the arguments
    args = vars(parser.parse_args())

    if not os.geteuid() == 0:  # must be root in order to manipulate snapshots
        print "You must run this script as root"
        exit(1)

    # parsing the subcommands
    if args['which'] == 'rm':
        rm(args)
    if args['which'] == 'cat':
        cat(args)
    if args['which'] == 'backup':
        backup(args)
    if args['which'] == 'restore':
        restore(args)
    if args['which'] == 'list':
        listSubvolumes()


if __name__ == '__main__':
    main()
