view hrm.py @ 22:dafc1cce8a2c default tip

Added another TODO.
author Atul Varma <varmaa@toolness.com>
date Sat, 13 Dec 2008 11:46:44 -0800
parents 7fa60aca382d
children
line wrap: on
line source

#! /usr/bin/env python
# ----------------------------------------------------------------------------
# Copyright (c) 2006, Atul Varma
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#
#   * Redistributions in binary form must reproduce the above
#     copyright notice, this list of conditions and the following
#     disclaimer in the documentation and/or other materials provided
#     with the distribution.
#
#   * Neither the name of hrm nor the names of its
#     contributors may be used to endorse or promote products derived
#     from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
# ----------------------------------------------------------------------------

'''
   This is hrm, a humane replacement for rm, with support for undo.
'''

# TODO:
#
#   * Documentation/tutorial.
#
#   * setup.py script for easy installation.
#
#   * Unit and/or functional tests; doctests would be really nice.
#
#   * If the '-f' parameter isn't supplied, hrm should exit with a
#     nonzero exit code if a file doesn't exist to maintain compatibility
#     with rm.
#
#   * Consider adding support for "max size" based removal, so that
#     if a certain quota is exceeded, the minimum number of files are
#     removed to satisfy the quota.
#
#   * Optional log messages for deletions.
#
#   * Make the transactions atomic, so that if an error occurs partway
#     through removing a set of files, the ones that were deleted
#     before the error happened are undeleted.
#
#   * Also make the transactions atomic so that if hrm is run at the
#     same time under the same account in different terminal sessions,
#     there aren't any race conditions.

import os
import re
import sys
import subprocess
import distutils.dir_util
import time
from optparse import OptionParser

try:
    # Python 2.6 and above
    import json
except ImportError:
    try:
        # Python 2.4 and above
        import simplejson as json
    except ImportError:
        print('Could not import json or simplejson. Please '
              'either upgrade to Python 2.6 or above, or '
              'run "sudo easy_install simplejson".')
        sys.exit(-1)

# Number of hours that must pass before we purge removed files.
HOURS_TO_PURGE = 48

# The above value expressed as seconds.
SECS_TO_PURGE = (60 *                   # Seconds in a minute
                 60 *                   # Minutes in an hour
                 HOURS_TO_PURGE)

# Directory to store state information and removed files (for undo).
HADES_DIR = os.path.expanduser('~/.Hades')

# Where to store state information.
STATE_FILENAME = os.path.join(HADES_DIR, 'state.json')

class Config(object):
    def __init__(self, filename):
        self.__filename = filename
        self.nextid = 0
        if os.path.exists(self.__filename):
            self.__dict__.update(json.load(open(self.__filename, 'r')))
        else:
            self.save()

    def save(self):
        state = {}
        keys = [key for key in self.__dict__
                if not key.startswith('_')]
        for key in keys:
            state[key] = self.__dict__[key]
        json.dump(state, open(self.__filename, 'w'))

def shell(*params):
    popen = subprocess.Popen(params)
    popen.wait()
    if popen.returncode:
        raise Exception('Process failed: %s' % repr(params))

def dir_for_trans(transid):
    return os.path.join(HADES_DIR, '%.9d' % transid)

def purge_old_files(verbose):
    paths = [os.path.abspath(os.path.join(HADES_DIR, name))
             for name in os.listdir(HADES_DIR)]

    min_mtime = time.time() - SECS_TO_PURGE
    to_purge = [path for path in paths
                if os.path.isdir(path)
                and os.stat(path).st_mtime < min_mtime]

    if verbose:
        ids = [str(int(os.path.split(dirname)[-1])) for 
               dirname in to_purge]
        print('Purging transactions: %s' % ', '.join(ids))

    for dirname in to_purge:
        distutils.dir_util.remove_tree(dirname)

if __name__ == '__main__':
    if not (os.path.exists(HADES_DIR) and os.path.isdir(HADES_DIR)):
        shell('mkdir', HADES_DIR)

    parser = OptionParser(description=__import__(__name__).__doc__)
    parser.add_option('-u', '--undo',
                      dest='undo', action='store_true', default=False,
                      help='undo a transaction (default is latest).')
    parser.add_option('-v', '--verbose',
                      dest='verbose', action='store_true', default=False,
                      help='be verbose.')

    args = []
    for arg in sys.argv[1:]:
        # Filter out redundant parameters to the traditional 'rm'
        # command.
        if not re.match('-[dfiRrv]+', arg):
            args.append(arg)
        elif 'v' in arg:
            args.append('-v')
    (options, args) = parser.parse_args(args)

    config = Config(STATE_FILENAME)

    purge_old_files(options.verbose)

    if options.undo:
        transactions = []
        for arg in args:
            try:
                transid = int(arg)
            except ValueError:
                print('Unknown transaction ID: %s' % arg)
                sys.exit(-1)
            transactions.append(transid)
        if not transactions and config.nextid:
            transactions.append(config.nextid - 1)
        for transid in transactions:
            dirname = dir_for_trans(transid)
            if not os.path.exists(dirname):
                print('Transaction ID %d does not exist.' % transid)
                sys.exit(-1)

        transactions.sort(reverse=True)

        for transid in transactions:
            transdirname = dir_for_trans(transid)
            for dirpath, dirnames, filenames in os.walk(transdirname):
                relpath = dirpath[len(transdirname)+1:]
                contents = dirnames + filenames
                for name in contents:
                    srcpath = os.path.join(dirpath, name)
                    destpath = os.path.join('/', relpath, name)
                    if not os.path.exists(destpath):
                        if os.path.isdir(srcpath):
                            dirnames.remove(name)
                        print('Restoring %s.' % destpath)
                        shell('mv', srcpath, destpath)
        print('Done.')
    else:
        if not args:
            parser.print_help()
            sys.exit(-1)

        files = []
        for arg in args:
            filename = os.path.abspath(os.path.expanduser(arg))
            if not os.path.exists(filename):
                print('File does not exist: %s' % arg)
                continue
            realpath = os.path.realpath(filename)
            if (realpath.startswith(HADES_DIR) or
                HADES_DIR.startswith(realpath)):
                print('Cannot move files in or above %s.' % HADES_DIR)
                sys.exit(-1)
            files.append(filename)

        if not files:
            # For compatibility with 'rm -f', just exit w/ no
            # error code.
            sys.exit(0)

        thisid = config.nextid
        basedir = dir_for_trans(thisid)
        shell('mkdir', basedir)
        config.nextid = thisid + 1
        config.save()

        print('The transaction ID for this operation is %s.' % thisid)

        for source in files:
            if options.verbose:
                print('Removing %s.' % source)
            dest = os.path.join(basedir, source[1:])
            distutils.dir_util.mkpath(os.path.dirname(dest))
            shell('mv', source, dest)