install

#!/usr/bin/env python3
import argparse
import collections
import os
import subprocess
import sys
import tempfile


class InstException(Exception):
    pass


def parse_args():
    agp = argparse.ArgumentParser(description="skel installer")
    tgt = agp.add_mutually_exclusive_group()
    tgt.add_argument("--symlink", action="store_true",
                     help="Instead of copying the files, symlink them")
    tgt.add_argument("--hardlink", action="store_true",
                     help="Instead of copying the files, hardlink them")
    tgt.add_argument("--remote", type=str, metavar="(USER@)HOST",
                     help="Install to a remote server via sshfs")
    tgt.add_argument("--target", type=str, metavar="TARGETDIR",
                     help="Install to this directory, instead of ~")
    agp.add_argument("--force-unchanged", action="store_true",
                     help="re-install files even if they didn't change")
    agp.add_argument("--backup", type=str, metavar="FILE.tar",
                     help="backup all overwritten files to this tar archive")
    agp.add_argument("--simulate", action="store_true",
                     help="don't write anything, just print the commands")
    agp.add_argument("--exclude", type=str, action='append', metavar="REGEX",
                     help="don't install files and directories that match " +
                          "this regular expression")
    agp.add_argument("skeldir", nargs="+",
                     help="Source directory for installation")

    return agp.parse_args()


def main(cleanup):
    args = parse_args()
    if args.target:
        tgtdir = args.target

        # no shiny name required
        shiny_tgtdir = tgtdir
        shiny_tgt = lambda t: t
    elif args.remote:
        # make temp dir (mountpoint)
        tgtdir = tempfile.mkdtemp(prefix="sshfs")
        # schedule temp dir removal
        cleanup.append(lambda: os.rmdir(tgtdir))
        # mount
        print("mounting remote \x1b[1m" + args.remote + "\x1b[m to " +
              "\x1b[3m" + tgtdir + "\x1b[m ... ", end="")
        sys.stdout.flush()
        retval = subprocess.call(['sshfs', args.remote + ":", tgtdir])
        if retval != 0:
            raise InstException("could not mount remote " + args.remote)
        print("\x1b[32;1mOK\x1b[m\n")
        # schedule umount
        cleanup.append(lambda: subprocess.call(['fusermount', '-u', tgtdir]))

        # shiny name: replace tgtdir by args.remote:
        shiny_tgtdir = args.remote + ":"
        shiny_tgt = lambda t: shiny_tgtdir + os.path.relpath(t, tgtdir)
    else:
        tgtdir = os.environ["HOME"]

        # shiny name: replace tgtdir by ~/
        shiny_tgtdir = "~"
        shiny_tgt = lambda t: shiny_tgtdir + "/" + os.path.relpath(t, tgtdir)

    # lists of file actions to perform
    mkdirlist = []
    installlist = []
    dupeinstalls = {}
    warnings = []
    excluded_srcfiles = []
    overwritten_tgtfiles = []

    # scan skeldir, check for actions to perform (fill mkdirlist, installlist)
    for skeldir in args.skeldir:
        if not os.path.isdir(skeldir):
            raise InstException("not a directory: " + skeldir)

        searchqueue = collections.deque([(skeldir, tgtdir)])

        while len(searchqueue) > 0:
            currentsrcdirname, currenttgtdirname = searchqueue.pop()

            for f in sorted(os.listdir(currentsrcdirname)):
                srcfile = currentsrcdirname + "/" + f
                tgtfile = currenttgtdirname + "/" + f

                if os.path.islink(srcfile):
                    raise InstException("symlink: " + srcfile)

                if args.exclude:
                    import re

                    excluded = False
                    for expr in args.exclude:
                        if re.match(expr, srcfile):
                            excluded = True
                            break

                    if excluded:
                        relfile = os.path.relpath(srcfile, skeldir)
                        isdir = os.path.isdir(srcfile)
                        excluded_srcfiles.append((relfile, skeldir, isdir))
                        continue

                if os.path.isdir(srcfile):
                    # it's a dir, so we need to create that if not exists
                    if os.path.exists(tgtfile):
                        if os.path.isdir(tgtfile):
                            # dir already exists, everything's fine
                            pass
                        else:
                            # there is a file with that name
                            raise InstException("Existing file " + tgtfile +
                                                " blocks directory creation")
                    else:
                        # we need to create the dir
                        if tgtfile not in mkdirlist:
                            mkdirlist.append(tgtfile)

                    # recursively search that dir, too
                    searchqueue.appendleft((srcfile, tgtfile))

                else:
                    # it's a file, so we need to copy it
                    if tgtfile in dupeinstalls:
                        # the file is already being installed from a different
                        # folder; undo that.
                        warnings.append(shiny_tgt(tgtfile) + " from \x1b[1m" +
                                        dupeinstalls[tgtfile] + "\x1b[m " +
                                        "obsolted by \x1b[1m" + skeldir +
                                        "\x1b[m")

                        installlist = [entry for entry in installlist if
                                       entry[1] != tgtfile]

                    unchanged = False
                    if os.path.exists(tgtfile):
                        if os.path.isdir(tgtfile):
                            raise InstException("Existing dir " + tgtfile +
                                                " prevents file installation")

                        # check whether the file differs
                        tgtdata = open(tgtfile, 'rb').read()
                        srcdata = open(srcfile, 'rb').read()
                        if tgtdata == srcdata:
                            # file unchanged
                            if args.force_unchanged:
                                # overwrite it anyways
                                unchanged = True
                            else:
                                # skip it
                                continue

                        exists = True
                    else:
                        exists = False

                    dupeinstalls[tgtfile] = skeldir

                    installlist.append((srcfile, tgtfile,
                                        exists, skeldir, unchanged))

    if not installlist and not mkdirlist:
        raise InstException("Nothing to be installed")

    # ask for user confirmation
    if mkdirlist:
        print("The following directories will be created:\n")
        for tgtfile in mkdirlist:
            print("\x1b[3m" + shiny_tgt(tgtfile) + "\x1b[m")

        print("")

    def maxlen(l, key=lambda x: x, default=0):
        """
        finds the len of the longest list element.
        use 'key' to pre-process elements.
        returns default for empty lists.

        returns the tuple of len(key(element)), element
        """
        if not l:
            return default
        else:
            longestelement = max(l, key=lambda e: len(key(e)))
            return len(key(longestelement)), longestelement

    if installlist:
        print("The following files will be installed:\n")
        maxtgtfilelen, _ = maxlen(installlist, key=lambda e: shiny_tgt(e[1]))
        maxskeldirlen, _ = maxlen(installlist, key=lambda e: e[3])

        for srcfile, tgtfile, exists, skeldir, unchanged in installlist:
            if exists:
                if args.backup:
                    warning = "\x1b[33m[File will be backed up]\x1b[m"
                else:
                    warning = "\x1b[31;1m[File will be overwritten]\x1b[m"

                if unchanged:
                    warning += " (unchanged)"

                shortname = os.path.relpath(tgtfile, tgtdir)
                overwritten_tgtfiles.append((tgtfile, srcfile, shortname))
            else:
                warning = ""

            print(shiny_tgt(tgtfile).ljust(maxtgtfilelen) + " from \x1b[1m" +
                  skeldir.ljust(maxskeldirlen) + "\x1b[m " + warning)

        print("")

    if args.backup and not overwritten_tgtfiles:
        warnings.append("nothing to backup")
        args.backup = False
    elif not args.backup and overwritten_tgtfiles:
        warnings.append("files will be overwritten, " +
                        "and no backup will be performed!")

    if args.exclude and not excluded_srcfiles:
        warnings.append("no files were matched by the exclusion expression!")

    if excluded_srcfiles:
        print("The following files will \x1b[1mnot\x1b[m be installed:\n")
        maxrelfilelen, _ = maxlen(excluded_srcfiles, key=lambda e: e[0])
        for relfile, skeldir, isdir in excluded_srcfiles:
            if isdir:
                relfile = "\x1b[3m" + relfile.ljust(maxrelfilelen) + "\x1b[m"
            print(relfile + " from \x1b[1m" + skeldir + "\x1b[m")

        print("")

    if args.simulate:
        warnings.append("simulation run: none of these actions will actually "
                        "be performed.")
        if args.backup:
            warnings.append("simulation run: will not actually create backup")

    if warnings:
        print("\x1b[1;33m[Warnings]\x1b[m\n")
        for warning in warnings:
            print("- " + warning)
        prompt_default = "n"
    else:
        prompt_default = "y"

    def prompt(default):
        options = []

        def yes():
            return True
        yes.helptext = "Continue"
        yes.expectedargs = []
        yes.opt = "y"
        options.append(yes)

        def no():
            raise InstException("Aborted by user")
        no.helptext = "Abort"
        no.expectedargs = []
        no.opt = "n"
        options.append(no)

        if overwritten_tgtfiles:
            def diff(filename):
                for tgtfile, srcfile, shortname in overwritten_tgtfiles:
                    if shortname == filename:
                        print("\x1b[1mdiff\x1b[m " + shiny_tgt(tgtfile) +
                              " " + srcfile)

                        # try to use colordiff, which has fancier output
                        try:
                            subprocess.call(["colordiff", "-u", "--",
                                             tgtfile, srcfile])

                            return False
                        except:
                            pass

                        subprocess.call(["diff", "-u", "--",
                                         tgtfile, srcfile])
                        return False

                print("Not an overwritten file: " + filename)
                return False

            diff.helptext = "Show diff for overwritten file FILE"
            diff.expectedargs = ["FILE"]
            diff.opt = "d"

            def completer(args):
                if len(args) == 0:
                    prefix = ""
                elif len(args) == 1:
                    prefix = args[0]
                else:
                    return None

                return ['d ' + shortname
                        for _, _, shortname in overwritten_tgtfiles
                        if shortname.startswith(prefix)]

            diff.completions = completer
            options.append(diff)

        if len(options) > 2:
            def shell():
                shell = os.environ.get('SHELL', '/bin/sh')
                print("Spawning an interactive shell... press ^D to return.")
                subprocess.call([shell])
                return False

            shell.helptext = "Spawn an interactive shell"
            shell.expectedargs = []
            shell.opt = "s"
            options.append(shell)

            def help():
                for f in options:
                    print(f.opt, end="")
                    if f.expectedargs:
                        print(" " + ", ".join(f.expectedargs), end="")
                    print(": " + f.helptext)
                return False
            help.helptext = "Show this help text"
            help.expectedargs = []
            help.opt = "?"
            options.append(help)

        optstring = "".join((f.opt for f in options))
        optstring = optstring.replace(default, default.upper())

        def find_completions(text):
            text = [x for x in text.split(' ') if x]
            if not text:
                return [f.opt for f in options]

            opt = text[0].lower()
            args = text[1:]

            for f in options:
                if f.opt == opt and hasattr(f, 'completions'):
                    return f.completions(args)

            return None

        class Completer:
            def complete(self, text, state):
                if state == 0:
                    self.matches = find_completions(text)
                try:
                    return self.matches[state]
                except IndexError:
                    return None

        import readline
        readline.set_completer(Completer().complete)
        readline.parse_and_bind("tab: complete")
        readline.set_completer_delims('')

        print("")
        try:
            userinput = input("Continue [" + optstring + "]? ")
        except EOFError:
            userinput = 'n'
            print("^D")
        except KeyboardInterrupt:
            userinput = 'n'
            print("")

        readline.set_completer(None)

        userinput = [x for x in userinput.split(' ') if x]
        if not userinput:
            userinput = [default]

        opt = userinput[0].lower()
        args = userinput[1:]

        for f in options:
            if opt == f.opt:
                if len(args) != len(f.expectedargs):
                    if not f.expectedargs:
                        print("'" + opt + "' expects no arguments")
                        return False
                    elif len(f.expectedargs) == 1:
                        print("'" + opt + "' expects the argument " +
                              f.expectedargs[0])
                        return False
                    else:
                        print("'" + opt + "' expects the arguments " +
                              ", ".join(f.expectedargs))
                        return False

                return f(*args)

        print("Unknown input: " + opt)
        return False

    while not prompt(prompt_default):
        prompt_default = 'n'

    print("")

    # create backup
    if args.backup:
        print("\x1b[1mtar cf\x1b[m " + args.backup + " \x1b[1m-C\x1b[m " +
              "\x1b[3m" + shiny_tgtdir + "\x1b[m", end="")
        sys.stdout.flush()

        if not args.simulate:
            import tarfile
            tar = tarfile.open(args.backup, 'w')

        for tgtfile, _, shortname in overwritten_tgtfiles:
            print(" " + shortname, end="")
            sys.stdout.flush()
            if not args.simulate:
                tar.add(tgtfile, shortname)

        if not args.simulate:
            tar.close()

        print("")

    # do installation
    # create directories
    for tgtfile in mkdirlist:
        print("\x1b[1mmkdir\x1b[m \x1b[3m" + shiny_tgt(tgtfile) + "\x1b[m")
        if not args.simulate:
            os.mkdir(tgtfile)

    # install files
    for srcfile, tgtfile, exists, _, _ in installlist:
        if exists:
            print("\x1b[1mrm\x1b[m " + shiny_tgt(tgtfile))
            if not args.simulate:
                os.unlink(tgtfile)

        if args.symlink:
            srcfile = os.path.relpath(srcfile, os.path.dirname(tgtfile))
            print("\x1b[1mln -s\x1b[m " + srcfile + " " + shiny_tgt(tgtfile))
            if not args.simulate:
                os.symlink(srcfile, tgtfile)
        elif args.hardlink:
            print("\x1b[1mln\x1b[m " + srcfile + " " + shiny_tgt(tgtfile))
            if not args.simulate:
                os.link(srcfile, tgtfile)
        else:
            print("\x1b[1mcp\x1b[m " + srcfile + " " + shiny_tgt(tgtfile))
            if not args.simulate:
                with open(srcfile, "rb") as f:
                    data = f.read()
                with open(tgtfile, "wb") as f:
                    f.write(data)

            if os.stat(srcfile).st_mode & (64 + 8 + 1):
                # set x flag
                print("\x1b[1mchmod +x\x1b[m " + shiny_tgt(tgtfile))
                if not args.simulate:
                    os.chmod(tgtfile, os.stat(tgtfile).st_mode | (64 + 8 + 1))

if __name__ == "__main__":
    cleanup = []

    try:
        main(cleanup)
    except InstException as e:
        print(e)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
    except Exception:
        import traceback
        traceback.print_exc()

    for c in reversed(cleanup):
        c()