view profilemanager/manager.py @ 79:145e111903d2 default tip

add MPL license
author Jeff Hammel <jhammel@mozilla.com>
date Mon, 10 May 2010 13:11:38 -0700
parents e091caa41075
children
line wrap: on
line source

# ***** BEGIN LICENSE BLOCK *****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
# 
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
# 
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
# 
# The Original Code is mozilla.org code.
# 
# The Initial Developer of the Original Code is
# Mozilla.org.
# Portions created by the Initial Developer are Copyright (C) 2010
# the Initial Developer. All Rights Reserved.
# 
# Contributor(s):
#     Jeff Hammel <jhammel@mozilla.com>     (Original author)
# 
# Alternatively, the contents of this file may be used under the terms of
# either of the GNU General Public License Version 2 or later (the "GPL"),
# or the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
# 
# ***** END LICENSE BLOCK *****
"""
manage Mozilla/Firefox profiles
"""

import os
import shutil
import string
import sys
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

        # ensure backup directory is not already present
        assert not os.path.exists(dest), "'%s' already exists"

        # 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)
            try:
                backups_dirs.append((backup, int(os.path.getmtime(backup_dir))))
            except OSError:
                print >> sys.stderr, "%s specified in %s, but the directory %s does not exist" % (backup, self.profiles, 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)

            self.write(parser)

        # 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))