tar-home.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import (
    absolute_import, division, print_function, unicode_literals
)
from nine import (
    IS_PYTHON2, basestring, chr, class_types, filter, integer_types,
    implements_iterator, implements_to_string, implements_repr,
    input, iterkeys, iteritems, itervalues, long, map,
    native_str, nine, nimport, range, range_list, reraise, str, zip
)

argparse = nimport('argparse')
contextlib = nimport('contextlib')
glob = nimport('glob')
grp = nimport('grp')
os = nimport('os')
pwd = nimport('pwd')
subprocess = nimport('subprocess')
sys = nimport('sys')
time = nimport('time')

DEFAULT_PREFIX = 'home_'
DEFAULT_SUFFIX = '.tar.gz'
DEFAULT_KEEP = 5


def run():
    """Runs this program.

    The program generates a new tar file for a given home directory.
    """
    start_time = time.time()
    options = parse_cmd_line(
        default_prefix=DEFAULT_PREFIX,
        default_suffix=DEFAULT_SUFFIX,
        default_keep=DEFAULT_KEEP)
    remove_old_files(
        pattern='{destdir}{sep}{prefix}*{suffix}'.format(
            destdir=options['destdir'],
            sep=os.sep,
            prefix=options['prefix'],
            suffix=options['suffix']),
        keep=options['keep'] - 1, # subtract 1 because new tar will be created
        remove=options['remove'],
        verbose=options['verbose'])
    dest_tar=generate_tar_name(
        dest_dir=options['destdir'],
        prefix=options['prefix'],
        suffix=options['suffix'])
    tar_src_dir(
        src_dir=options['srcdir'],
        dest_tar=dest_tar,
        verbose=options['verbose'])
    change_file_owner_group(
        filename=dest_tar,
        user=options['user'],
        group=options['user'],
        verbose=options['verbose'])
    if options['verbose']:
        print('Done, in {} seconds!'.format(time.time() - start_time))


def generate_tar_name(**kwargs):
    """Generates a (hopefully) unique fully qualified tar filename.

    Arguments:
        kwargs: a dictionary with the following keys:-
            dest_dir: where to save the tar file
            prefix: the prefix to use for the tar file
            suffix: the suffix to use for the tar file

    Returns:
        A unique tar name of the format:
            <dest_dir>/<prefix><year><month><day><hour><minute>second><suffix>
        where the year, month, day, etc. values represent the local time
        at the point at which this function is called.
    """
    dest_dir = kwargs.pop('dest_dir')
    prefix = kwargs.pop('prefix')
    suffix = kwargs.pop('suffix')
    if kwargs:
        raise TypeError('Unexpected **kwargs: %r' % kwargs)
    localtime = time.localtime()
    return ('{dest_dir}{sep}{prefix}{year:04d}{month:02d}{day:02d}' +
        '{hour:02d}{minute:02d}{second:02d}{suffix}').format(
            dest_dir=dest_dir,
            prefix=prefix,
            suffix=suffix,
            year=localtime[0],
            month=localtime[1],
            day=localtime[2],
            hour=localtime[3],
            minute=localtime[4],
            second=localtime[5],
            sep=os.sep)


def parse_cmd_line(**kwargs):
    """Parses the command-line arguments.

    Arguments:
        kwargs: a dictionary with the following keys:-
            default_prefix: the default prefix for generated tar files
            default_suffix: the default suffix for generated tar files
            default_keep:   the default number of tar files to keep if removing
                            old tar files

    Returns:
        A dictionary with each of the supplied command-line arguments. If
        no value for an item was supplied on the command-line, then the
        default value for that it is returned.
    """
    default_prefix = kwargs.pop('default_prefix')
    default_suffix = kwargs.pop('default_suffix')
    default_keep = kwargs.pop('default_keep')
    if kwargs:
        raise TypeError('Unexpected **kwargs: %r' % kwargs)
    parser = argparse.ArgumentParser(
        description="Archives (tars) the contents of a user's home directory")
    parser.add_argument(
        'user',
        help=' '.join([
            'specify the owner of the archive file;',
            'the archive files user and group will be set to this value']))
    parser.add_argument(
        'srcdir',
        help='specify the owners home directory')
    parser.add_argument(
        'destdir',
        help='specify the directory that the new archive file will be saved to')
    parser.add_argument(
        '--prefix', '-p',
        default=default_prefix,
        help=' '.join([
            'specify the prefix that the newly created archive file will have',
            "(default: '%(default)s')"]))
    parser.add_argument(
        '--suffix', '-x',
        default=default_suffix,
        help=' '.join([
            'specify the suffix that the newly created archive file will have',
            "(default: '%(default)s')"]))
    parser.add_argument(
        '--keep', '-k',
        type=int,
        default=default_keep,
        help='specify how many archives to keep (default: %(default)i)')
    parser.add_argument(
        '--remove', '-r',
        action='store_true',
        default=False,
        help='specify this so that obsolete archives will be removed')
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        default=False,
        help='specify this to display verbose output')
    # vars() turns Namespace into a regular dictionary
    options = vars(parser.parse_args())
    options['srcdir'] = chomp_sep(options['srcdir'])
    options['destdir'] = chomp_sep(options['destdir'])
    options['keep'] = max(options['keep'], 1) # keep must be at least 1
    return options


def chomp_sep(dir_name):
    """Removes any trailing directory separator characters from the given
    directory name.

    Arguments:
        dir_name: the name that has to have any trailing slashes removed

    Returns:
        The directory name with no trailing separator characters
    """
    while dir_name.endswith(os.sep):
        dir_name = dir_name[:-1]
    return dir_name


def remove_old_files(**kwargs):
    """Only keeps the newest files that match the given pattern.

    Arguments:
        kwargs: a dictionary with the following keys:-
            'pattern'   - the glob pattern of existing files
            'keep'      - the number of existing files to keep
            'remove'    - whether or not to remove old files
            'verbose'   - whether or not to output text describing non-fatal
                          events

    Returns:
        None
    """
    pattern = kwargs.pop('pattern')
    keep= kwargs.pop('keep')
    remove = kwargs.pop('remove')
    verbose= kwargs.pop('verbose')
    if kwargs:
        raise TypeError('Unexpected **kwargs: %r' % kwargs)
    if not remove:
        return
    matching = glob.glob(pattern)
    if keep >= len(matching):
        if verbose:
            print("No files matching pattern '{}' need to be removed".format(
                pattern))
        return
    files = [tup[1] for tup in sorted([(os.path.getmtime(f), f)
                    for f in matching])]
    for f in files[keep:]:
        os.remove(f)
        if verbose:
            print("Removed file '{filename}'".format(filename=f))


def tar_src_dir(src_dir, dest_tar, verbose):
    """Archives the source directory into a newly created destination tar.

    Arguments:
        src_dir:    the source directory to be tarred
        dest_tar:   the name of the tar file to create
        verbose:    whether or not to output text describing non-fatal events

    Returns:
        None
   """
    src_dir = os.path.abspath(src_dir)
    index = src_dir.rindex(os.sep)
    base_dir = src_dir[0:index]
    home_dir = src_dir[index+1:]
    if verbose:
        verbose_flag = '--verbose '
    else:
        verbose_flag = ''
    tar_cmd = ''.join(['sudo tar ',
                        '--create --preserve-permissions --atime-preserve ',
                        "--use-compress-program='pigz' ",
                        verbose_flag,
                        "--file={} ".format(dest_tar),
                        home_dir,
                      ])
    with chdir(base_dir):
        if verbose:
            subprocess.call(tar_cmd, shell=True)
        else:
            subprocess.call('{} {}'.format(tar_cmd, '&>/var/null'), shell=True)
    # Ensure that the file is flushed to disk
    subprocess.call('sync; sleep 10; sync', shell=True)


@contextlib.contextmanager
def chdir(new_dir):
    """Context manager which allows code to run in a different directory.

    (1) Stores the current working directory;
    (2) switches to the directory: <new_dir>;
    (3) executes any code run in this context; and then
    (4) returns to the original working directory.

    Returns:
        None
    """
    cwd = os.getcwd()
    os.chdir(new_dir)
    try:
        yield
    finally:
        os.chdir(cwd)


def change_file_owner_group(**kwargs):
    """Changes the owner and group of the named file.

    Arguments:
        kwargs: a dictionary with the following keys:-
            'filename': the file whose ownership is to be changed
            'user':     the new owner to assign to the file
            'group':    the new group to assign to the file;
                        this is the same as 'user' if not given
            'verbose':  whether to output text describing non-fatal events

    Returns:
        None
    """
    filename = kwargs.pop('filename')
    user = kwargs.pop('user')
    group = kwargs.pop('group', user)
    verbose= kwargs.pop('verbose')
    if kwargs:
        raise TypeError('Unexpected **kwargs: %r' % kwargs)
    try:
        os.chown(filename,
                 pwd.getpwnam(user).pw_uid,
                 grp.getgrnam(group).gr_gid)
        if verbose:
            print("Changed ownership of '{filename}' to {user}:{group}".format(
                filename=filename,
                user=user,
                group=group))
    except os.OSError:
        print("Unable to change ownership of "
            + "'{filename}' to {user}:{group}".format(
                filename=filename,
                user=user,
                group=group),
            file=sys.stderr)


if __name__ == '__main__':
    run()