view profilemanager/manager.py @ 74:b6bb59b79525

* return the backed-up profile * further test restoring profiles
author Jeff Hammel <jhammel@mozilla.com>
date Fri, 07 May 2010 17:09:32 -0700
parents 1cfd259f74cf
children a63dd19807c8
line wrap: on
line source

"""
manage Mozilla/Firefox profiles
"""

import os
import shutil
import string
import tempfile
import time
import ConfigParser
from datetime import datetime
from dateutil.parser import parse
from random import Random
from utils import format_tabular
from ConfigParser import SafeConfigParser as ConfigParser

# Changes to profiles.ini:
# - add a ``backups`` field for each profile:
# Backups = /path/to/backup:1273104310 /path/to/other/backup:1273104355

class ProfileNotFound(Exception):
    """
    exception when a profile is specified but is not present in a given
    .ini file
    """
    def __init__(self, profile, config):
        self.profile = profile
        self.config = config
        Exception.__init__(self,
                           'Profile %s not found in %s' % (profile, config))

class NoBackupError(Exception):
    """
    exception when a sought backup is not found
    """

class ProfileManager(object):

    backups_dir = 'backups' # directory for backups relative to profile_dir

    def __init__(self, profiles):
        """
        - profiles: filesystem path to profiles.ini file
        """
        self.profiles = profiles
        self.profile_dir = os.path.abspath(os.path.dirname(profiles))

    ### (public) API

    def new(self, name, directory=None, hash=True):
        """
        generate a new clean profile
        - name: name of the profile to generate
        - directory: where to create the profile [DEFAULT: relative to profiles.ini]
        - hash: whether to generate a a random hash tag
        """

        # path to the profile directory
        dirname = name
        relative = False
        if hash:
            dirname = '%s.%s' % (self.hash(), dirname)
        if not directory:
            directory = self.profile_dir
            relative = True
        path = os.path.join(directory, dirname)

        # create directory
        # TODO: (optionally) pre-populate the directory a la Firefox
        os.mkdir(path)

        # update profiles.ini
        self.add(name, relative and dirname or path, relative)
        
        # return the directory name
        return path

    def remove(self, name, delete=True):
        """
        remove a profile from profiles.ini
        - delete: delete the profile directory as well
        """
        parser = self.parser()
        section = self.section(name, parser)
        if section is None:
            raise ProfileNotFound(profile, self.profiles)
        if delete: # remove the profile from disk
            shutil.rmtree(self.path(name))
        parser.remove_section(section)
        self.write(parser)

    def list(self, directories=False):
        """
        lists the profiles available in the config file
        - directories : display the directories
        """
        profiles = self.profiles_dict()
        if not directories:
            return sorted(profiles.keys())
        return dict([(name, self.path(name)) for name in profiles.keys()])

    def clone(self, source, dest, add=True):
        """
        clones the profile `source` and output to `dest`
        - add: add the profile to the profiles.ini file
        """

        # filesystem path of the `from` profile
        source_path = self.path(source)    

        # dest: fs path to back up to
        relative = False
        if os.path.sep in dest:
            if not os.path.isabs(dest):
                dest = os.path.abspath(dest)
            dirname = dest
            name = os.path.basename(dest)
        else:
            name = dest
            relative = True
            if add:
                dirname = '%s.%s' % (self.hash(), name)
            dest = os.path.join(self.profile_dir, dirname)

        # update profiles.ini
        if add:
            self.add(name, dirname, relative)

        # copy the files
        shutil.copytree(source_path, dest, symlinks=False)

        return dest

    def backup(self, profile, dest=None):
        """
        backup a profile
        - profile: name of the profile to be backed up
        - dest: name of the destination; if not given, a default is used
        """

        # get the profile section
        parser = self.parser()
        section = self.section(profile)
        if section is None:
            raise ProfileNotFound(profile, self.profiles)

        # determine destination directory
        if dest is None:
            dest = '%s.%d.bak' % (profile, int(time.time()))
            name = dest
            backups_dir = os.path.join(self.profile_dir, self.backups_dir)
            if not os.path.exists(backups_dir):
                os.mkdir(backups_dir)
            dest = os.path.join(backups_dir, dest)
        else:
            if not os.path.isabs(dest):
                dest = os.path.abspath(dest)
            name = dest

        # copy the files
        self.clone(profile, dest, add=False)
        
        # add backup entry to profiles.ini (colon separated):
        # `Backup=$(profile)s.$(datestamp)s.bak`
        if parser.has_option(section, 'Backups'):
            backups = '%s:%s' % (parser.get(section, 'Backups'), name)
        else:
            backups = name
        parser.set(section, 'Backups', backups)
        self.write(parser)

    def backups(self, profile=None, datestamp=None):
        """
        list backups for a given profile, or all profiles if the
        profile is not given;  returns a list of backups if
        profile is given or a dictionary of lists otherwise
        - datestamp: format of date; otherwise, seconds since epoch
        """
        if profile is None: # all profiles
            retval = {}
            profiles_dict = self.profiles_dict()
            for profile in profiles_dict:
                if 'Backups' in profiles_dict[profile]:
                    retval[profile] = self.backups(profile, datestamp)
            return retval

        # single profile
        profile_dict = self.profile_dict(profile)
        if 'Backups' not in profile_dict:
            return []
        backups = profile_dict['Backups'].split(':')
        backups_dirs = []
        for backup in backups:
            if os.path.isabs(backup):
                backup_dir = backup
            else:
                backup_dir = os.path.join(self.profile_dir, self.backups_dir, backup)
            backups_dirs.append((backup, int(os.path.getmtime(backup_dir))))
        # TODO: check if the getmtime == the datestamp for relative backups

        # sort by reverse datestamp
        backups_dirs.sort(key=lambda x: x[1], reverse=True)

        # format to datestamps, if applicable
        if datestamp:
            backups_dirs = [ (i[0], time.strftime(datestamp,
                                                 time.localtime(int(i[1]))))
                             for i in backups_dirs ]

        return backups_dirs

    def restore(self, profile, date=None, delete=False):
        """
        restore the profile from a backup
        the most recent backup is used unless `date` is given
        - date : date to restore from (will not delete)
        - delete : delete the backup after restoration
        """

        assert not (date and delete), 'date and delete cannot be used in conjunction'
        orig_date = date

        # get the path to the profile
        path = self.path(profile)

        # get the possible backups
        backups = self.backups(profile)
        # TODO: check to see if these all exist (print warnings if not)

        # determine the backup to use
        if date:

            if isinstance(date, basestring):

                # test for s since epoch
                seconds_snice_epoch = None
                try:
                    seconds_since_epoch = int(date)
                    
                except ValueError:
                    pass
                if seconds_since_epoch and seconds_since_epoch > time.localtime().tm_year:
                    date = seconds_since_epoch
                else:
                    date = parse(date) # parse date from string

            if isinstance(date, datetime):
                # convert to s since epoch
                date = time.mktime(date.timetuple())
            
            for backup in backups:
                if backup[1] < date:
                    break
            else:
                raise NoBackupError("No backups for profile %s in %s earlier than %s" % (profile, self.profiles, orig_date))
        else:
            if not backups:
                raise NoBackupError("No backups for profile %s in %s" % (profile, self.profiles))
            backup = backups[0]

        # determine path of the backup
        if os.path.isabs(backup[0]):
            backup_path = backup
        else:
            backup_path = os.path.join(self.profile_dir, self.backups_dir, backup[0])

        # restore the backup over ``profile``
        # make a backup copy first to avoid data loss
        tempdir = tempfile.mktemp()
        shutil.move(path, tempdir)
        try:
            shutil.copytree(backup_path, path)
        except Exception, e:
            shutil.rmtree(path)
            shutil.move(tempdir, path)
            raise IOError("Couldn't restore backup: %s" % str(e))
        shutil.rmtree(tempdir)
        # TODO: (optionally) make a backup of the current profile before restore

        if delete: # delete the backup

            # delete the directory
            shutil.rmtree(backup_path)
            
            # delete the entry from ``profiles.ini``
            parser = self.parser()
            section = self.section(profile)
            backups = parser.get(section, 'Backups').split(':')
            backups = [ i for i in backups if i != backup[0] ]
            parser.set(section, 'Backups', ':'.join(backups))

            # if there are no backups, delete the ``backups`` line
            # and the backups directory
            if not backups:
                parser.remove_option(section, 'Backups')
                backups_dir = os.path.join(self.profile_dir, self.backups_dir)
                if not os.listdir(backups_dir):
                    shutil.rmtree(backups_dir)

        # return which backup is restored
        return backup

    def temporary(self):
        """make a temporary profile"""
        raise NotImplementedError

    def merge(self, output, *profiles):
        """merge a set of profiles (not trivial!)"""
        raise NotImplementedError

    ### internal functions

    def add(self, profile, path, relative=True):
        """
        add a profile entry to profiles.ini
        """

        # ensure name is not already present
        assert profile not in self.profiles_dict(), 'Profile "%s" already in %s' % (name, self.profiles)
        parser = self.parser()

        # find and add the section 
        ctr = 0
        section = 'Profile%d' % ctr # unsure of this naming convention
        while section in parser.sections():
            ctr += 1
            section = 'Profile%d' % ctr
        parser.add_section(section)

        # add metadata
        parser.set(section, 'Name', profile)
        parser.set(section, 'IsRelative', '%d' % int(relative))
        parser.set(section, 'Path', path)
        if not ctr:
            parser.set(section, 'Default', '1')

        # write the file
        self.write(parser)

    def path(self, profile):
        """returns the path to the profile"""
        profile = self.profile_dict(profile)
        if profile.get('IsRelative', None) == '1':
            return os.path.join(self.profile_dir, profile['Path'])
        return profile['Path']

    def parser(self):
        """
        return a ConfigParser instance appropriate to profiles.ini
        """
        parser = ConfigParser()
        parser.optionxform = str
        parser.read(self.profiles)
        return parser

    def write(self, parser):
        f = file(self.profiles, 'w')
        parser.write(f)
        f.close()

    def section(self, profile, parser=None):
        """
        returns the name of the section that a profile is in or None
        if not found
        """
        if parser is None:
            parser = self.parser()
        for section in parser.sections():
            if not parser.has_option(section, 'Name'):
                continue # not a profile
            if parser.get(section, 'Name') == profile:
                return section

    def profile_dict(self, profile):
        """
        return option dictionary for a single profile
        """
        parser = self.parser()
        section = self.section(profile, parser)
        if section is None:
            raise ProfileNotFound(profile, self.profiles)
        return dict(parser.items(section))

    def profiles_dict(self):
        """
        return nested dict of all profiles
        """
        # assumes profiles have unique names
        parser = self.parser()
        retval = {}
        for section in parser.sections():
            if section == 'General':
                continue
            if not parser.has_option(section, 'Name'):
                continue
            name = parser.get(section, 'Name')
            retval[name] = self.profile_dict(name)
        return retval

    def hash(self):
        """
        generate a random hash for a new profile
        """
        population = string.lowercase + string.digits
        return ''.join(Random().sample(population, 8))