#! /usr/bin/env python
"""Usage: atomic_rm.py [-r] [-f] [files...]

    This works like 'rm' but attempts to be as atomic as possible:
    either the whole operation completely succeeds, or no file is
    removed.  This is only useful in recursive mode or if more
    than one file is specified on the command line.
"""

import os, sys, stat, errno


HIDE_TRACEBACKS = True      # use '--tb' to enable tracebacks

class RmError(Exception):
    pass

def ErrNo(errno, filename, moreinfo=''):
    if moreinfo:
        moreinfo = '\n(%s)' % (moreinfo,)
    return OSError(errno, '%s: %r%s' % (os.strerror(errno),
                                        filename, moreinfo))

def my_rename(file1, file2):
    try:
        os.rename(file1, file2)
    except OSError, e:
        raise ErrNo(e.errno, file1, 'failed to move it to %r' % (file2,))

class WarningPrinter:
    count = 0

    def do(self, function, *args):
        try:
            function(*args)
        except OSError, e:
            print >> sys.stderr, 'warning: %s: %s' % (function.__name__, e)
            self.count += 1

    def set_permissions(self, filename, st):
        self.do(os.chmod, filename, stat.S_IMODE(st.st_mode))
        self.do(os.utime, filename, (st.st_atime, st.st_mtime))
        self.do(os.chown, filename, st.st_uid, st.st_gid)


def create_backup_dir(filename, attempts=100):
    # symlink note: if 'filename' is itself a symlink, we do not dereference
    # it, because we are trying to remove the symlink itself.  If one of
    # the dirs listed in the path of 'filename' is a symlink, that's fine.
    dirname, basename = os.path.split(filename)
    i = 0
    while True:
        tryname = os.path.join(dirname, '%s.atomic_rm%d~' % (basename, i))
        try:
            os.mkdir(tryname)
        except OSError, e:
            if e.errno == errno.EEXIST:
                # the new name we are trying already exists, try again
                i += 1
                if i < attempts:
                    continue
            raise
        return tryname


def atomic_rm(files, **kwds):
    # The idea is to first move all files away without actually killing them
    # so that if something unexpected goes wrong, we can try putting them
    # all back in place.  If something seriously bad occurs, e.g. this
    # process is killed, backup dirs should still contain the files that
    # have been moved away.
    warn = WarningPrinter()
    backup_dirs = []
    try:
        list_all_files(files, backup_dirs, warn, **kwds)
        move_to_backup_dirs(backup_dirs)
        remove_backup_dirs(backup_dirs)
    finally:
        restore_from_backup_dirs(backup_dirs)   # in case of error or crash
    return warn.count


class BackupDir(object):
    def __init__(self, originalfn, originalst, warn, dir=None, force=False,
                 protected_st_list=()):
        self.originalfn = originalfn
        self.originalst = originalst
        self.warn = warn
        self.force = force
        if dir is None:
            self.dir = create_backup_dir(originalfn)
            self.keep = False
        else:
            os.mkdir(dir)
            self.dir = dir
            self.keep = True
        if protected_st_list != ():
            protected_st_list.append(os.lstat(self.dir))
        self.protected_st_list = protected_st_list

    def check(self, fullname):
        if self.force or os.path.islink(fullname):
            return
        if not os.access(fullname, os.W_OK):
            if os.access(fullname, os.F_OK):
                raise OSError(errno.EACCES,
                              '%r is read-only (use -f to remove it)'
                              % (fullname,))
            else:
                self.disappeared(fullname)

    def disappeared(self, fullname):
        raise OSError(errno.ENOENT,
                      '%r went away while the operation was in progress'
                      % (fullname,))

    def join(self, basename):
        return os.path.join(self.dir, basename)

    def rmdir(self):
        if not self.keep:
            self.warn.do(os.rmdir, self.dir)

class BackupDirSingle(BackupDir):
    empty = True

    def move_in(self):
        basename = os.path.basename(self.originalfn)
        self.check(self.originalfn)
        try:
            my_rename(self.originalfn, self.join(basename))
        except OSError, e:
            if e.errno == errno.ENOENT and self.force:     # ignore error
                return
            raise
        self.empty = False

    def make_empty(self):
        if not self.empty and not self.keep:
            basename = os.path.basename(self.originalfn)
            self.warn.do(os.unlink, self.join(basename))
            self.empty = True

    def emergency_restore(self):
        if not self.empty:
            basename = os.path.basename(self.originalfn)
            self.warn.do(my_rename, self.join(basename), self.originalfn)
            self.empty = True

class BackupDirRec(BackupDir):
    content = ()
    original_backup_dir_st = None

    def orgjoin(self, basename):
        if basename == '':
            return self.originalfn
        else:
            return os.path.join(self.originalfn, basename)

    def move_in(self):
        # 'content' will contain a list of (relpath, statresult, step)
        # which describe all files and directories that are being removed.
        # The order is depth-first.  A directory is listed both before its
        # content and after its content; the first time with step == "start"
        # and the second time with step == "stop".  The step is ignored
        # for non-directories.
        self.original_backup_dir_st = os.lstat(self.join(''))
        self.content = []
        pending = [('', self.originalst, 'start')]
        while pending:
            relpath, st, step = pending.pop()
            src = self.orgjoin(relpath)

            if not stat.S_ISDIR(st.st_mode):
                # moving a non-directory file object
                self.check(src)
                my_rename(src, self.join(relpath))

            elif step == 'start':
                # moving a directory into the backup dir
                self.check(src)
                newentries = self.find_directory_content(relpath, st)
                pending.extend(newentries)
                if relpath != '':    # the top-level backup dir exists already
                    os.mkdir(self.join(relpath))

            else:
                # finished recursing into a directory: we can rmdir()
                # the original directory.  We also try to preserve
                # the permissions.
                self.warn.set_permissions(self.join(relpath), st)
                os.rmdir(src)

            # done
            self.content.append((relpath, st, step))

    def find_directory_content(self, relpath, st):
        src = self.orgjoin(relpath)
        for st1 in self.protected_st_list:
            if os.path.samestat(st, st1):
                raise RmError("refusing to remove %r" % (src,))
        try:
            names = os.listdir(src)
        except OSError, e:
            if e.errno == errno.ENOENT:
                self.disappeared(src)
            raise
        newentries = []
        for name in names:
            subrelpath = os.path.join(relpath, name)
            fullname = self.orgjoin(subrelpath)
            sub_st = os.lstat(fullname)
            if sub_st.st_dev != st.st_dev:
                # NB. this is probably a good safety measure, but
                # removing this limitation would be a bit of work
                raise ErrNo(errno.EXDEV, fullname,
                            'refusing to remove across multiple devices')
            newentries.append((subrelpath, sub_st, 'start'))
        newentries.append((relpath, st, 'stop'))
        newentries.reverse()
        return newentries

    def make_empty(self):
        if self.keep:
            return
        while self.content:
            relpath, st, step = self.content[-1]

            if not stat.S_ISDIR(st.st_mode):
                # killing a non-directory file object
                self.warn.do(os.unlink, self.join(relpath))

            elif step == 'stop':
                # as we're popping self.content from the end, this marks
                # the beginning of the kill operation for a directory.
                # We might have to add 'write' permissions, otherwise
                # we cannot unlink its files.
                if self.original_backup_dir_st is not None:
                    self.warn.set_permissions(self.join(relpath),
                                              self.original_backup_dir_st)

            else:
                # done emptying a directory: kill the dir in the backup
                # (but not the backup dir itself)
                if relpath != '':
                    self.warn.do(os.rmdir, self.join(relpath))

            del self.content[-1]

    def emergency_restore(self):
        while self.content:
            relpath, st, step = self.content[-1]

            if not stat.S_ISDIR(st.st_mode):
                # restoring a non-directory file object
                self.warn.do(my_rename, self.join(relpath),
                                        self.orgjoin(relpath))

            elif step == 'stop':
                # as we're popping self.content from the end, this marks
                # the beginning of the restore operation for a directory.
                # We might have to add 'write' permissions, otherwise
                # we cannot move files out of it.
                self.warn.do(os.mkdir, self.orgjoin(relpath))
                if self.original_backup_dir_st is not None:
                    self.warn.set_permissions(self.join(relpath),
                                              self.original_backup_dir_st)

            else:
                # done restoring a directory: fix up the permissions
                # on the restored directory
                self.warn.set_permissions(self.orgjoin(relpath), st)
                # kill the empty dir in the backup
                # (but not the backup dir itself)
                if relpath != '':
                    self.warn.do(os.rmdir, self.join(relpath))

            del self.content[-1]


def list_all_files(files, backup_dirs, warn,
                   backup=None, recursive=False, force=False):
    if backup is not None and len(files) > 1:
        raise RmError("cannot specify multiple files with --backup")
    protected_st_list = [os.lstat('/')]
    for fn in files:
        try:
            st = os.lstat(fn)
        except OSError, e:
            if e.errno == errno.ENOENT and force:
                continue     # completely ignore the file
            raise
        for st1 in protected_st_list:
            if os.path.samestat(st, st1):
                raise RmError("refusing to remove %r" % (fn,))
        if stat.S_ISDIR(st.st_mode):
            fn = fn.rstrip('/')   # remove trailing slashes
            assert fn    # if 'fn' contains only slashes, it's the root
            if os.path.basename(fn) in ['.', '..']:
                raise RmError("refusing to remove '.' or '..'")
            if not recursive:
                raise ErrNo(errno.EISDIR, fn,
                            'use -r to remove all its content')
            Cls = BackupDirRec
        else:
            Cls = BackupDirSingle
        dirfn = os.path.dirname(fn) or os.curdir
        if not os.access(dirfn, os.W_OK):
            raise ErrNo(errno.EACCES, dirfn,
                        'the parent directory is not writable')
        backup_dirs.append(Cls(fn, st, dir=backup, warn=warn, force=force,
                               protected_st_list=protected_st_list))

def move_to_backup_dirs(backup_dirs):
    for backup_dir in backup_dirs:
        backup_dir.move_in()

def remove_backup_dirs(backup_dirs):
    while backup_dirs:
        backup_dirs[-1].make_empty()
        backup_dirs[-1].rmdir()
        del backup_dirs[-1]

def restore_from_backup_dirs(backup_dirs):
    while backup_dirs:
        backup_dirs[-1].emergency_restore()
        backup_dirs[-1].rmdir()
        del backup_dirs[-1]

def remove_trailing_slash(fn):
    if len(fn) > 1 and fn.endswith('/') and not fn.endswith('//'):
        fn = fn[:-1]
    return fn


def main(args):
    from getopt import gnu_getopt as getopt
    options, files = getopt(args, 'rRf', ['help', 'version', 'backup=',
                                          'recursive', 'force', 'tb'])
    kwds = {}
    for option, value in options:
        if option == '--help':
            print >> sys.stderr, __doc__
            return 2
        elif option == '--version':
            if files:
                raise RmError("cannot give both files and --version")
            print >> sys.stderr, '$LastChangedRevision$'
            return 0
        elif option in ('-r', '-R', '--recursive'):
            kwds['recursive'] = True
        elif option in ('-f', '--force'):
            kwds['force'] = True
        elif option == '--backup':
            kwds['backup'] = value
        elif option == '--tb':
            global HIDE_TRACEBACKS
            HIDE_TRACEBACKS = False
        else:
            raise ValueError(option)
    if not files and not kwds.get('force'):
        raise RmError("missing operand, try --help")
    files = [remove_trailing_slash(fn) for fn in files]
    warncount = atomic_rm(files, **kwds)
    if warncount == 0:
        return 0
    else:
        return 3

if __name__ == '__main__':
    try:
        exitstatus = main(sys.argv[1:])
    except (OSError, RmError), e:
        if HIDE_TRACEBACKS:
            print >> sys.stderr, '%s: %s' % (os.path.basename(sys.argv[0]), e)
            exitstatus = 1
        else:
            raise
    sys.exit(exitstatus)