<commit_before>import mailpile.plugins
from mailpile.commands import Command, Action
from mailpile.mailutils import Email, ExtractEmails
from mailpile.vcard import SimpleVCard, VCardLine
from mailpile.util import *

class ContactImporter:
    required_parameters = []
    optional_parameters = []
    format_name = ""

    def __init__(self, **kwargs):
        self.args = kwargs

    def load(self):
        pass

    def get_contacts(self):
        return []

    def filter_contacts(self, terms):
        # We put this in each ContactImporter because sometimes fetching
        # everything and then filtering might be really slow, and in those
        # cases there might be a faster way. Example: CardDAV lookup.
        return []

class ContactExporter:
    required_parameters = []
    optional_parameters = []
    format_name = ""

    def __init__(self):
        self.exporting = []

    def add_contact(self, contact):
        self.exporting.append(contact)

    def remove_contact(self, contact):
        self.exporting.remove(contact)

    def save(self):
        pass

class ContactContextProvider:
    provider_name = ""

    def __init__(self, contact):
        self.contact = contact

    def get_recent_context(self, max=10):
        pass

    def get_related_context(self, query, max=10):
        pass

class ContactFieldValidator:
    def __init__(self):
        pass

##[ Pluggable contact management ]########################################

CONTACT_IMPORTERS = {}
CONTACT_EXPORTERS = {}
CONTACT_FIELD_VALIDATORS = {}
CONTACT_CONTEXT_PROVIDERS = {}

def register_contact_importer(importer):
    if not issubclass(importer, ContactImporter):
        raise PluginError("Plugin must be a ContactImporter")
    if importer.format_name in CONTACT_IMPORTERS.keys():
        raise PluginError("Importer for %s already registered"
                          % importer.format_name)
    CONTACT_IMPORTERS[importer.short_name] = importer

def register_contact_exporter(exporter):
    if not issubclass(importer, ContactExporter):
        raise PluginError("Plugin must be a ContactExporter")
    if exporter.format_name in CONTACT_EXPORTERS.keys():
        raise PluginError("Exporter for %s already registered"
                          % exporter.format_name)
    CONTACT_EXPORTERS[exporter.short_name] = exporter

def register_contact_field_validator(field, validator):
    if not issubclass(importer, ContactFieldValidator):
        raise PluginError("Plugin must be a ContactFieldValidator")
    if field in CONTACT_FIELD_VALIDATORS.keys():
        raise PluginError("Field validator for field %s already registered"
                          % field)
    CONTACT_FIELD_VALIDATORS[field] = validator

def register_contact_context_provider(provider):
    if not issubclass(importer, ContactContextProvider):
        raise PluginError("Plugin must be a ContactContextProvider")
    if importer.provider_name in CONTACT_CONTEXT_PROVIDERS.keys():
        raise PluginError("Context provider for %s already registered"
                          % provider.provider_name)
    CONTACT_CONTEXT_PROVIDERS[provider.provider_name] = provider

##[ VCards ]########################################

class VCardCommand(Command):
    VCARD = "vcard"

    def _fparse(self, fromdata):
        email = ExtractEmails(fromdata)[0]
        name = fromdata.replace(email, '').replace('<>', '').strip()
        return email, (name or email)

    def _make_new_vcard(self, handle, name):
        l = [VCardLine(name='fn', value=name),
             VCardLine(name='kind', value=self.KIND)]
        if self.KIND == 'individual':
            return SimpleVCard(VCardLine(name='email', value=handle), *l)
        else:
            return SimpleVCard(VCardLine(name='nickname', value=handle), *l)

    def _valid_vcard_handle(self, vc_handle):
        return (vc_handle and '@' in vc_handle[1:])

    def _add_from_messages(self):
        pairs, idx = [], self._idx()
        for email in [Email(idx, i) for i in self._choose_messages(self.args)]:
            pairs.append(self._fparse(email.get_msg_info(idx.MSG_FROM)))
        return pairs

    def _pre_delete_vcard(self, vcard):
        pass

    def _vcard_list(self, vcards, mode='mpCard', info=None):
        info = info or {}
        if mode == 'lines':
            data = [x.as_lines() for x in vcards]
        else:
            data = [x.as_mpCard() for x in vcards]
        info.update({
            self.VCARD+'s': data,
            "count": len(vcards)
        })
        return info

class VCard(VCardCommand):
    """Add/remove/list/edit vcards"""
    SYNOPSIS = (None, 'vcard', None, '<nickname>')
    ORDER = ('Internals', 6)
    KIND = ''

    def command(self, save=True):
        session, config = self.session, self.session.config
        vcards = []
        for email in self.args:
            vcard = config.vcards.get_vcard(email)
            if vcard:
                vcards.append(vcard)
            else:
                session.ui.warning('No such %s: %s' % (self.VCARD, email))
        return vcards

class AddVCard(VCardCommand):
    """Add one or more vcards"""
    SYNOPSIS = (None, 'vcard/add', None, '<msgs>', '<email> = <name>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'PUT')

    def command(self):
        session, config, idx = self.session, self.session.config, self._idx()

        if (len(self.args) > 2
        and self.args[1] == '='
        and self._valid_vcard_handle(self.args[0])):
            pairs = [(self.args[0], ' '.join(self.args[2:]))]
        elif self.data:
            if "@contactname" in self.data and "@contactemail" in self.data:
                pairs = [(self.data["@contactemail"],
                          self.data["@contactname"])]
            elif "contactnames" in self.data and "contactemails" in self.data:
                pairs = zip(self.data["contactemails"],
                            self.data["contactnames"])
        else:
            pairs = self._add_from_messages()

        if pairs:
            vcards = []
            for handle, name in pairs:
                if handle.lower() not in config.vcards:
                    vcard = self._make_new_vcard(handle, name)
                    config.vcards.add_vcard(vcard)
                    vcards.append(vcard)
                else:
                    session.ui.warning('Already exists: %s' % handle)
        else:
            return self._error('Nothing to do!')
        return {"contacts": [x.as_mpCard() for x in vcards]}

class VCardAddLines(VCardCommand):
    """Add a lines to a VCard"""
    SYNOPSIS = (None, 'vcard/addline', None, '<email> <lines>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'UPDATE')

    def command(self):
        session, config = self.session, self.session.config
        handle, var, lines = self.args[0], self.args[1], self.args[2:]
        vcard = config.vcards.get_vcard(handle)
        if not vcard:
            return self._error('%s not found: %s' % (self.VCARD, handle))
        config.vcards.deindex_vcard(vcard)
        try:
            vcard.add(*[VCardLine(l) for l in lines])
            vcard.save()
            return self._vcard_list([vcard], info={
                'updated': handle,
                'added': len(lines)
            })
        except:
            config.vcards.index_vcard(vcard)
            self._ignore_exception()
            return self._error('Error setting %s = %s' % (var, val))
        finally:
            config.vcards.index_vcard(vcard)

class RemoveVCard(VCardCommand):
    """Delete vcards"""
    SYNOPSIS = (None, 'vcard/remove', None, '<email>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'DELETE')

    def command(self):
        session, config = self.session, self.session.config
        for handle in self.args:
            vcard = config.vcards.get_vcard(handle)
            if vcard:
                self._pre_delete_vcard(vcard)
                config.vcards.del_vcard(handle)
            else:
                session.ui.error('No such contact: %s' % handle)
        return True

class ListVCards(VCardCommand):
    """Find vcards"""
    SYNOPSIS = (None, 'vcard/list', None, '[--lines] [<terms>]')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_QUERY_VARS = {
        'q': 'search terms',
        'format': 'lines or mpCard (default)'
    }

    def command(self):
        session, config = self.session, self.session.config
        kinds = self.KIND and [self.KIND] or []

        if 'format' in self.data:
            fmt = self.data['format'][0]
        elif self.args and self.args[0] == '--lines':
            self.args.pop(0)
            fmt = 'lines'
        else:
            fmt = 'mpCard'

        if 'q' in self.data:
            terms = self.data['q']
        else:
            terms = self.args

        vcards = config.vcards.find_vcards(terms, kinds=kinds)
        return self._vcard_list(vcards, mode=fmt, info={
            'terms': self.args
        })

def ContactVCard(parent):
    """A factory for generating contact commands"""
    synopsis = [(t and t.replace('vcard', 'contact') or t)
                for t in parent.SYNOPSIS]
    synopsis[2] = synopsis[1]

    class ContactVCardCommand(parent):
        SYNOPSIS = tuple(synopsis)
        KIND = 'individual'
        ORDER = ('Tagging', 3)
        VCARD = "contact"

    return ContactVCardCommand

class Contact(ContactVCard(VCard)):
    """View contacts"""

class AddContact(ContactVCard(AddVCard)):
    """Add contacts"""

class ContactAddLines(ContactVCard(VCardAddLines)):
    """Set contact variables"""

class RemoveContact(ContactVCard(RemoveVCard)):
    """Remove a contact"""

class ListContacts(ContactVCard(ListVCards)):
    SYNOPSIS = (None, 'contact/list', 'contact/list', '[--lines] [<terms>]')
    """Find contacts"""

class ContactImport(Command):
    """Import contacts"""
    SYNOPSIS = (None, 'contact/import', 'contact/import', '[<parameters>]')
    ORDER = ('Internals', 6)
    HTTP_CALLABLE = ('GET', )

    def command(self, format, terms=None, **kwargs):
        session, config = self.session, self.session.config

        if not format in mailpile.plugins.CONTACT_IMPORTERS.keys():
            session.ui.error("No such import format")
            return False

        importer = mailpile.plugins.CONTACT_IMPORTERS[format]

        if not all([x in kwargs.keys() for x in importer.required_parameters]):
            session.ui.error("Required paramter missing. Required parameters "
                + "are: %s" % ", ".join(importer.required_parameters))
            return False

        allparams = importer.required_parameters + importer.optional_parameters

        if not all([x in allparams for x in kwargs.keys()]):
            session.ui.error("Unknown parameter passed to importer."
                + " Provided %s; but known parameters are: %s"
                % (", ".join(kwargs), ", ".join(allparams)))
            return False

        imp = importer(kwargs)
        if terms:
            contacts = imp.filter_contacts(terms)
        else:
            contacts = imp.get_contacts()

        for importedcontact in contacts:
            # Check if contact exists. If yes, then update. Else create.
            pass

class ContactImporters(Command):
    """Return a list of contact importers"""
    SYNOPSIS = (None, 'contact/importers', 'contact/importers', '')
    ORDER = ('Internals', 6)
    HTTP_CALLABLE = ('GET', )

    def command(self):
        res = []
        for iname, importer in CONTACT_IMPORTERS.iteritems():
            r = {}
            r["short_name"] = iname
            r["format_name"] = importer.format_name
            r["format_description"] = importer.format_description
            r["optional_parameters"] = importer.optional_parameters
            r["required_parameters"] = importer.required_parameters
            res.append(r)

        return res

class AddressSearch(VCardCommand):
    """Find addresses (in contacts or mail index)"""
    SYNOPSIS = (None, 'search/address', 'search/address', '[<terms>]')
    ORDER = ('Searching', 6)
    HTTP_QUERY_VARS = {
        'q': 'search terms',
    }

    def _vcard_addresses(self, cfg, terms):
        addresses = []
        for vcard in cfg.vcards.find_vcards(terms, kinds='individual'):
             fn = vcard.get('fn')
             keys = [{'fingerprint': k.value, 'type': 'openpgp'}
                     for k in vcard.get_all('openpgp-key')]
             photos = vcard.get_all('photo')
             for email_vcl in vcard.get_all('email'):
                 rank = 3.0 + 5*len(keys) + len(photos)
                 for term in terms:
                     if term in fn.value.lower():
                         rank += 10*(float(len(term)) / len(fn.value))
                     if term in email_vcl.value:
                         rank += 10*(float(len(term)) / len(email_vcl.value))
                 info = {
                     'rank': int(rank),
                     'fn': fn.value,
                     'proto': 'smtp',
                     'address': email_vcl.value,
                     'keys': [k for k in keys[:1]]
                 }
                 addresses.append(info)
        return addresses

    def _index_addresses(self, cfg, terms, vcard_addresses):
        existing = dict([(k['address'], k) for k in vcard_addresses])
        index = self._idx()
        # FIXME: 1st, go through the last 100 or so messages and search for
        #        matching senders or recipients, give medium priority.
        # FIXME: 2nd, search the social graph for matches, give low priority.
        addresses = []
        return addresses

    def command(self):
        session, config = self.session, self.session.config
        if 'q' in self.data:
            terms = self.data['q']
        else:
            terms = self.args

        vcard_addrs = self._vcard_addresses(config, terms)
        index_addrs = self._index_addresses(config, terms, vcard_addrs)
        addresses = vcard_addrs + index_addrs
        addresses.sort(key=lambda k: -k['rank'])
        return {
            'addresses': addresses[:10]
        }

mailpile.plugins.register_commands(VCard, AddVCard, VCardAddLines,
                                   RemoveVCard, ListVCards)
mailpile.plugins.register_commands(Contact, AddContact, ContactAddLines,
                                   RemoveContact, ListContacts,
                                   AddressSearch)
mailpile.plugins.register_commands(ContactImport, ContactImporters)
<commit_msg>1st draft of mixing in results from index into address results<commit_after>import mailpile.plugins
from mailpile.commands import Command, Action
from mailpile.mailutils import Email, ExtractEmails
from mailpile.vcard import SimpleVCard, VCardLine
from mailpile.util import *

class ContactImporter:
    required_parameters = []
    optional_parameters = []
    format_name = ""

    def __init__(self, **kwargs):
        self.args = kwargs

    def load(self):
        pass

    def get_contacts(self):
        return []

    def filter_contacts(self, terms):
        # We put this in each ContactImporter because sometimes fetching
        # everything and then filtering might be really slow, and in those
        # cases there might be a faster way. Example: CardDAV lookup.
        return []

class ContactExporter:
    required_parameters = []
    optional_parameters = []
    format_name = ""

    def __init__(self):
        self.exporting = []

    def add_contact(self, contact):
        self.exporting.append(contact)

    def remove_contact(self, contact):
        self.exporting.remove(contact)

    def save(self):
        pass

class ContactContextProvider:
    provider_name = ""

    def __init__(self, contact):
        self.contact = contact

    def get_recent_context(self, max=10):
        pass

    def get_related_context(self, query, max=10):
        pass

class ContactFieldValidator:
    def __init__(self):
        pass

##[ Pluggable contact management ]########################################

CONTACT_IMPORTERS = {}
CONTACT_EXPORTERS = {}
CONTACT_FIELD_VALIDATORS = {}
CONTACT_CONTEXT_PROVIDERS = {}

def register_contact_importer(importer):
    if not issubclass(importer, ContactImporter):
        raise PluginError("Plugin must be a ContactImporter")
    if importer.format_name in CONTACT_IMPORTERS.keys():
        raise PluginError("Importer for %s already registered"
                          % importer.format_name)
    CONTACT_IMPORTERS[importer.short_name] = importer

def register_contact_exporter(exporter):
    if not issubclass(importer, ContactExporter):
        raise PluginError("Plugin must be a ContactExporter")
    if exporter.format_name in CONTACT_EXPORTERS.keys():
        raise PluginError("Exporter for %s already registered"
                          % exporter.format_name)
    CONTACT_EXPORTERS[exporter.short_name] = exporter

def register_contact_field_validator(field, validator):
    if not issubclass(importer, ContactFieldValidator):
        raise PluginError("Plugin must be a ContactFieldValidator")
    if field in CONTACT_FIELD_VALIDATORS.keys():
        raise PluginError("Field validator for field %s already registered"
                          % field)
    CONTACT_FIELD_VALIDATORS[field] = validator

def register_contact_context_provider(provider):
    if not issubclass(importer, ContactContextProvider):
        raise PluginError("Plugin must be a ContactContextProvider")
    if importer.provider_name in CONTACT_CONTEXT_PROVIDERS.keys():
        raise PluginError("Context provider for %s already registered"
                          % provider.provider_name)
    CONTACT_CONTEXT_PROVIDERS[provider.provider_name] = provider

##[ VCards ]########################################

class VCardCommand(Command):
    VCARD = "vcard"

    def _fparse(self, fromdata):
        email = ExtractEmails(fromdata)[0]
        name = fromdata.replace(email, '').replace('<>', '').strip()
        return email, (name or email)

    def _make_new_vcard(self, handle, name):
        l = [VCardLine(name='fn', value=name),
             VCardLine(name='kind', value=self.KIND)]
        if self.KIND == 'individual':
            return SimpleVCard(VCardLine(name='email', value=handle), *l)
        else:
            return SimpleVCard(VCardLine(name='nickname', value=handle), *l)

    def _valid_vcard_handle(self, vc_handle):
        return (vc_handle and '@' in vc_handle[1:])

    def _add_from_messages(self):
        pairs, idx = [], self._idx()
        for email in [Email(idx, i) for i in self._choose_messages(self.args)]:
            pairs.append(self._fparse(email.get_msg_info(idx.MSG_FROM)))
        return pairs

    def _pre_delete_vcard(self, vcard):
        pass

    def _vcard_list(self, vcards, mode='mpCard', info=None):
        info = info or {}
        if mode == 'lines':
            data = [x.as_lines() for x in vcards]
        else:
            data = [x.as_mpCard() for x in vcards]
        info.update({
            self.VCARD+'s': data,
            "count": len(vcards)
        })
        return info

class VCard(VCardCommand):
    """Add/remove/list/edit vcards"""
    SYNOPSIS = (None, 'vcard', None, '<nickname>')
    ORDER = ('Internals', 6)
    KIND = ''

    def command(self, save=True):
        session, config = self.session, self.session.config
        vcards = []
        for email in self.args:
            vcard = config.vcards.get_vcard(email)
            if vcard:
                vcards.append(vcard)
            else:
                session.ui.warning('No such %s: %s' % (self.VCARD, email))
        return vcards

class AddVCard(VCardCommand):
    """Add one or more vcards"""
    SYNOPSIS = (None, 'vcard/add', None, '<msgs>', '<email> = <name>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'PUT')

    def command(self):
        session, config, idx = self.session, self.session.config, self._idx()

        if (len(self.args) > 2
        and self.args[1] == '='
        and self._valid_vcard_handle(self.args[0])):
            pairs = [(self.args[0], ' '.join(self.args[2:]))]
        elif self.data:
            if "@contactname" in self.data and "@contactemail" in self.data:
                pairs = [(self.data["@contactemail"],
                          self.data["@contactname"])]
            elif "contactnames" in self.data and "contactemails" in self.data:
                pairs = zip(self.data["contactemails"],
                            self.data["contactnames"])
        else:
            pairs = self._add_from_messages()

        if pairs:
            vcards = []
            for handle, name in pairs:
                if handle.lower() not in config.vcards:
                    vcard = self._make_new_vcard(handle, name)
                    config.vcards.add_vcard(vcard)
                    vcards.append(vcard)
                else:
                    session.ui.warning('Already exists: %s' % handle)
        else:
            return self._error('Nothing to do!')
        return {"contacts": [x.as_mpCard() for x in vcards]}

class VCardAddLines(VCardCommand):
    """Add a lines to a VCard"""
    SYNOPSIS = (None, 'vcard/addline', None, '<email> <lines>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'UPDATE')

    def command(self):
        session, config = self.session, self.session.config
        handle, var, lines = self.args[0], self.args[1], self.args[2:]
        vcard = config.vcards.get_vcard(handle)
        if not vcard:
            return self._error('%s not found: %s' % (self.VCARD, handle))
        config.vcards.deindex_vcard(vcard)
        try:
            vcard.add(*[VCardLine(l) for l in lines])
            vcard.save()
            return self._vcard_list([vcard], info={
                'updated': handle,
                'added': len(lines)
            })
        except:
            config.vcards.index_vcard(vcard)
            self._ignore_exception()
            return self._error('Error setting %s = %s' % (var, val))
        finally:
            config.vcards.index_vcard(vcard)

class RemoveVCard(VCardCommand):
    """Delete vcards"""
    SYNOPSIS = (None, 'vcard/remove', None, '<email>')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_CALLABLE = ('POST', 'DELETE')

    def command(self):
        session, config = self.session, self.session.config
        for handle in self.args:
            vcard = config.vcards.get_vcard(handle)
            if vcard:
                self._pre_delete_vcard(vcard)
                config.vcards.del_vcard(handle)
            else:
                session.ui.error('No such contact: %s' % handle)
        return True

class ListVCards(VCardCommand):
    """Find vcards"""
    SYNOPSIS = (None, 'vcard/list', None, '[--lines] [<terms>]')
    ORDER = ('Internals', 6)
    KIND = ''
    HTTP_QUERY_VARS = {
        'q': 'search terms',
        'format': 'lines or mpCard (default)'
    }

    def command(self):
        session, config = self.session, self.session.config
        kinds = self.KIND and [self.KIND] or []

        if 'format' in self.data:
            fmt = self.data['format'][0]
        elif self.args and self.args[0] == '--lines':
            self.args.pop(0)
            fmt = 'lines'
        else:
            fmt = 'mpCard'

        if 'q' in self.data:
            terms = self.data['q']
        else:
            terms = self.args

        vcards = config.vcards.find_vcards(terms, kinds=kinds)
        return self._vcard_list(vcards, mode=fmt, info={
            'terms': self.args
        })

def ContactVCard(parent):
    """A factory for generating contact commands"""
    synopsis = [(t and t.replace('vcard', 'contact') or t)
                for t in parent.SYNOPSIS]
    synopsis[2] = synopsis[1]

    class ContactVCardCommand(parent):
        SYNOPSIS = tuple(synopsis)
        KIND = 'individual'
        ORDER = ('Tagging', 3)
        VCARD = "contact"

    return ContactVCardCommand

class Contact(ContactVCard(VCard)):
    """View contacts"""

class AddContact(ContactVCard(AddVCard)):
    """Add contacts"""

class ContactAddLines(ContactVCard(VCardAddLines)):
    """Set contact variables"""

class RemoveContact(ContactVCard(RemoveVCard)):
    """Remove a contact"""

class ListContacts(ContactVCard(ListVCards)):
    SYNOPSIS = (None, 'contact/list', 'contact/list', '[--lines] [<terms>]')
    """Find contacts"""

class ContactImport(Command):
    """Import contacts"""
    SYNOPSIS = (None, 'contact/import', 'contact/import', '[<parameters>]')
    ORDER = ('Internals', 6)
    HTTP_CALLABLE = ('GET', )

    def command(self, format, terms=None, **kwargs):
        session, config = self.session, self.session.config

        if not format in mailpile.plugins.CONTACT_IMPORTERS.keys():
            session.ui.error("No such import format")
            return False

        importer = mailpile.plugins.CONTACT_IMPORTERS[format]

        if not all([x in kwargs.keys() for x in importer.required_parameters]):
            session.ui.error("Required paramter missing. Required parameters "
                + "are: %s" % ", ".join(importer.required_parameters))
            return False

        allparams = importer.required_parameters + importer.optional_parameters

        if not all([x in allparams for x in kwargs.keys()]):
            session.ui.error("Unknown parameter passed to importer."
                + " Provided %s; but known parameters are: %s"
                % (", ".join(kwargs), ", ".join(allparams)))
            return False

        imp = importer(kwargs)
        if terms:
            contacts = imp.filter_contacts(terms)
        else:
            contacts = imp.get_contacts()

        for importedcontact in contacts:
            # Check if contact exists. If yes, then update. Else create.
            pass

class ContactImporters(Command):
    """Return a list of contact importers"""
    SYNOPSIS = (None, 'contact/importers', 'contact/importers', '')
    ORDER = ('Internals', 6)
    HTTP_CALLABLE = ('GET', )

    def command(self):
        res = []
        for iname, importer in CONTACT_IMPORTERS.iteritems():
            r = {}
            r["short_name"] = iname
            r["format_name"] = importer.format_name
            r["format_description"] = importer.format_description
            r["optional_parameters"] = importer.optional_parameters
            r["required_parameters"] = importer.required_parameters
            res.append(r)

        return res

class AddressSearch(VCardCommand):
    """Find addresses (in contacts or mail index)"""
    SYNOPSIS = (None, 'search/address', 'search/address', '[<terms>]')
    ORDER = ('Searching', 6)
    HTTP_QUERY_VARS = {
        'q': 'search terms',
    }

    def _boost_rank(self, term, *matches):
        boost = 0.0
        for match in matches:
            match = match.lower()
            if term in match:
                if match.startswith(term):
                    boost += 25*(float(len(term)) / len(match))
                else:
                    boost += 5*(float(len(term)) / len(match))
        return int(boost)

    def _vcard_addresses(self, cfg, terms):
        addresses = []
        for vcard in cfg.vcards.find_vcards(terms, kinds='individual'):
             fn = vcard.get('fn')
             keys = [{'fingerprint': k.value, 'type': 'openpgp'}
                     for k in vcard.get_all('openpgp-key')]
             photos = vcard.get_all('photo')
             for email_vcl in vcard.get_all('email'):
                 rank = 10.0 + 25*len(keys) + 5*len(photos)
                 for term in terms:
                     rank += self._boost_rank(term, fn.value, email_vcl.value)
                 info = {
                     'rank': int(rank),
                     'fn': fn.value,
                     'proto': 'smtp',
                     'address': email_vcl.value,
                 }
                 if keys:
                     info['keys'] = [k for k in keys[:1]]
                 addresses.append(info)
        return addresses

    def _index_addresses(self, cfg, terms, vcard_addresses):
        existing = dict([(k['address'].lower(), k) for k in vcard_addresses])
        index = self._idx()

        # 1st, go through the last 1000 or so messages in the index and search
        # for matching senders or recipients, give medium priority.
        matches = {}
        addresses = []
        for msg_idx in range(0, len(index.INDEX))[-1000:]:
            frm = index.get_msg_at_idx_pos(msg_idx)[index.MSG_FROM]
            match = True
            for term in terms:
                if term not in frm.lower():
                    match = False
            if match:
                matches[frm] = matches.get(frm, 0) + 1
        for frm in matches:
            email, fn = self._fparse(frm)
            boost = min(10, matches[frm])
            for term in terms:
                boost += self._boost_rank(term, fn, email)

            if email.lower() in existing:
                 existing[email.lower()]['rank'] += min(20, boost)
            else:
                 info = {
                     'rank': boost,
                     'fn': fn,
                     'proto': 'smtp',
                     'address': email,
                 }
                 existing[email.lower()] = info
                 addresses.append(info)

        # FIXME: 2nd, search the social graph for matches, give low priority.

        return addresses

    def command(self):
        session, config = self.session, self.session.config
        if 'q' in self.data:
            terms = [t.lower() for t in self.data['q']]
        else:
            terms = [t.lower() for t in self.args]

        vcard_addrs = self._vcard_addresses(config, terms)
        index_addrs = self._index_addresses(config, terms, vcard_addrs)
        addresses = vcard_addrs + index_addrs
        addresses.sort(key=lambda k: -k['rank'])
        return {
            'addresses': addresses[:10]
        }

mailpile.plugins.register_commands(VCard, AddVCard, VCardAddLines,
                                   RemoveVCard, ListVCards)
mailpile.plugins.register_commands(Contact, AddContact, ContactAddLines,
                                   RemoveContact, ListContacts,
                                   AddressSearch)
mailpile.plugins.register_commands(ContactImport, ContactImporters)
<|endoftext|>


<commit_before>#!/usr/bin/python2.5
# Copyright (c) 2010 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Snapshot Build Bisect Tool

This script bisects a snapshot archive using binary search. It starts at
a bad revision (it will try to guess HEAD) and asks for a last known-good
revision. It will then binary search across this revision range by downloading,
unzipping, and opening Chromium for you. After testing the specific revision,
it will ask you whether it is good or bad before continuing the search.
"""

# Base URL to download snapshots from.
BUILD_BASE_URL = 'http://build.chromium.org/buildbot/snapshots/'

# The type (platform) of the build archive. This is what's passed in to the
# '-a/--archive' option.
BUILD_ARCHIVE_TYPE = ''

# The selected archive to bisect.
BUILD_ARCHIVE_DIR = ''

# The location of the builds.
BUILD_ARCHIVE_URL = '/%d/'

# Name of the build archive.
BUILD_ZIP_NAME = ''

# Directory name inside the archive.
BUILD_DIR_NAME = ''

# Name of the executable.
BUILD_EXE_NAME = ''

# URL to the ViewVC commit page.
BUILD_VIEWVC_URL = 'http://src.chromium.org/viewvc/chrome?view=rev&revision=%d'

# Changelogs URL
CHANGELOG_URL = 'http://build.chromium.org/buildbot/' \
                'perf/dashboard/ui/changelog.html?url=/trunk/src&range=%d:%d'

###############################################################################

import math
import optparse
import os
import pipes
import re
import shutil
import sys
import tempfile
import urllib
import zipfile

def UnzipFilenameToDir(filename, dir):
  """Unzip |filename| to directory |dir|."""
  zf = zipfile.ZipFile(filename)
  # Make base.
  pushd = os.getcwd()
  try:
    if not os.path.isdir(dir):
      os.mkdir(dir)
    os.chdir(dir)
    # Extract files.
    for info in zf.infolist():
      name = info.filename
      if name.endswith('/'):  # dir
        if not os.path.isdir(name):
          os.makedirs(name)
      else:  # file
        dir = os.path.dirname(name)
        if not os.path.isdir(dir):
          os.makedirs(dir)
        out = open(name, 'wb')
        out.write(zf.read(name))
        out.close()
      # Set permissions. Permission info in external_attr is shifted 16 bits.
      os.chmod(name, info.external_attr >> 16L)
    os.chdir(pushd)
  except Exception, e:
    print >>sys.stderr, e
    sys.exit(1)

def SetArchiveVars(archive):
  """Set a bunch of global variables appropriate for the specified archive."""
  global BUILD_ARCHIVE_TYPE
  global BUILD_ARCHIVE_DIR
  global BUILD_ZIP_NAME
  global BUILD_DIR_NAME
  global BUILD_EXE_NAME
  global BUILD_BASE_URL

  BUILD_ARCHIVE_TYPE = archive
  BUILD_ARCHIVE_DIR = 'chromium-rel-' + BUILD_ARCHIVE_TYPE

  if BUILD_ARCHIVE_TYPE in ('linux', 'linux-64', 'linux-chromiumos'):
    BUILD_ZIP_NAME = 'chrome-linux.zip'
    BUILD_DIR_NAME = 'chrome-linux'
    BUILD_EXE_NAME = 'chrome'
  elif BUILD_ARCHIVE_TYPE in ('mac'):
    BUILD_ZIP_NAME = 'chrome-mac.zip'
    BUILD_DIR_NAME = 'chrome-mac'
    BUILD_EXE_NAME = 'Chromium.app/Contents/MacOS/Chromium'
  elif BUILD_ARCHIVE_TYPE in ('xp'):
    BUILD_ZIP_NAME = 'chrome-win32.zip'
    BUILD_DIR_NAME = 'chrome-win32'
    BUILD_EXE_NAME = 'chrome.exe'

  BUILD_BASE_URL += BUILD_ARCHIVE_DIR

def ParseDirectoryIndex(url):
  """Parses the HTML directory listing into a list of revision numbers."""
  handle = urllib.urlopen(url)
  dirindex = handle.read()
  handle.close()
  return re.findall(r'\1/', dirindex)

def GetRevList(good, bad):
  """Gets the list of revision numbers between |good| and |bad|."""
  # Download the main revlist.
  revlist = ParseDirectoryIndex(BUILD_BASE_URL)
  revlist = map(int, revlist)
  revlist = filter(lambda r: range(good, bad).__contains__(int(r)), revlist)
  revlist.sort()
  return revlist

def TryRevision(rev, profile, args):
  """Downloads revision |rev|, unzips it, and opens it for the user to test.
  |profile| is the profile to use."""
  # Do this in a temp dir so we don't collide with user files.
  cwd = os.getcwd()
  tempdir = tempfile.mkdtemp(prefix='bisect_tmp')
  os.chdir(tempdir)

  # Download the file.
  download_url = BUILD_BASE_URL + (BUILD_ARCHIVE_URL % rev) + BUILD_ZIP_NAME
  try:
    print 'Fetching ' + download_url
    urllib.urlretrieve(download_url, BUILD_ZIP_NAME)
  except Exception, e:
    print('Could not retrieve the download. Sorry.')
    sys.exit(-1)

  # Unzip the file.
  print 'Unziping ...'
  UnzipFilenameToDir(BUILD_ZIP_NAME, os.curdir)

  # Tell the system to open the app.
  args = ['--user-data-dir=%s' % profile] + args
  flags = ' '.join(map(pipes.quote, args))
  exe = os.path.join(os.getcwd(), BUILD_DIR_NAME, BUILD_EXE_NAME)
  cmd = '%s %s' % (exe, flags)
  print 'Running %s' % cmd
  os.system(cmd)

  os.chdir(cwd)
  print 'Cleaning temp dir ...'
  try:
    shutil.rmtree(tempdir, True)
  except Exception, e:
    pass

def AskIsGoodBuild(rev):
  """Ask the user whether build |rev| is good or bad."""
  # Loop until we get a response that we can parse.
  while True:
    response = raw_input('\nBuild %d is [(g)ood/(b)ad]: ' % int(rev))
    if response and response in ('g', 'b'):
      return response == 'g'

def main():
  usage = ('%prog [options] [-- chromium-options]\n'
           'Perform binary search on the snapshot builds.')
  parser = optparse.OptionParser(usage=usage)
  # Strangely, the default help output doesn't include the choice list.
  choices = ['mac', 'xp', 'linux', 'linux-64', 'linux-chromiumos']
  parser.add_option('-a', '--archive',
                    choices = choices,
                    help = 'The buildbot archive to bisect [%s].' %
                           '|'.join(choices))
  parser.add_option('-b', '--bad', type = 'int',
                    help = 'The bad revision to bisect to.')
  parser.add_option('-g', '--good', type = 'int',
                    help = 'The last known good revision to bisect from.')
  parser.add_option('-p', '--profile', '--user-data-dir', type = 'str',
                    help = 'Profile to use; this will not reset every run. ' +
                    'Defaults to a clean profile.')
  (opts, args) = parser.parse_args()

  if opts.archive is None:
    parser.print_help()
    return 1

  if opts.bad and opts.good and (opts.good > opts.bad):
    print ('The good revision (%d) must precede the bad revision (%d).\n' %
           (opts.good, opts.bad))
    parser.print_help()
    return 1

  SetArchiveVars(opts.archive)

  # Pick a starting point, try to get HEAD for this.
  if opts.bad:
    bad_rev = opts.bad
  else:
    bad_rev = 0
    try:
      # Location of the latest build revision number
      BUILD_LATEST_URL = '%s/LATEST' % (BUILD_BASE_URL)
      nh = urllib.urlopen(BUILD_LATEST_URL)
      latest = int(nh.read())
      nh.close()
      bad_rev = raw_input('Bad revision [HEAD:%d]: ' % latest)
      if (bad_rev == ''):
        bad_rev = latest
      bad_rev = int(bad_rev)
    except Exception, e:
      print('Could not determine latest revision. This could be bad...')
      bad_rev = int(raw_input('Bad revision: '))

  # Find out when we were good.
  if opts.good:
    good_rev = opts.good
  else:
    good_rev = 0
    try:
      good_rev = int(raw_input('Last known good [0]: '))
    except Exception, e:
      pass

  # Get a list of revisions to bisect across.
  revlist = GetRevList(good_rev, bad_rev)
  if len(revlist) < 2:  # Don't have enough builds to bisect
    print 'We don\'t have enough builds to bisect. revlist: %s' % revlist
    sys.exit(1)

  # If we don't have a |good_rev|, set it to be the first revision possible.
  if good_rev == 0:
    good_rev = revlist[0]

  # These are indexes of |revlist|.
  good = 0
  bad = len(revlist) - 1
  last_known_good_rev = revlist[good]

  # Binary search time!
  while good < bad:
    candidates = revlist[good:bad]
    num_poss = len(candidates)
    if num_poss > 10:
      print('%d candidates. %d tries left.' %
          (num_poss, round(math.log(num_poss, 2))))
    else:
      print('Candidates: %s' % revlist[good:bad])

    # Cut the problem in half...
    test = int((bad - good) / 2) + good
    test_rev = revlist[test]

    # Let the user give this rev a spin (in her own profile, if she wants).
    profile = opts.profile
    if not profile:
      profile = 'profile'  # In a temp dir.
    TryRevision(test_rev, profile, args)
    if AskIsGoodBuild(test_rev):
      last_known_good_rev = revlist[good]
      good = test + 1
    else:
      bad = test

  # We're done. Let the user know the results in an official manner.
  print('You are probably looking for build %d.' % revlist[bad])
  print('CHANGELOG URL:')
  print(CHANGELOG_URL % (last_known_good_rev, revlist[bad]))
  print('Built at revision:')
  print(BUILD_VIEWVC_URL % revlist[bad])

if __name__ == '__main__':
  sys.exit(main())
<commit_msg>build-bisect: clarify usage<commit_after>#!/usr/bin/python
# Copyright (c) 2010 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Snapshot Build Bisect Tool

This script bisects a snapshot archive using binary search. It starts at
a bad revision (it will try to guess HEAD) and asks for a last known-good
revision. It will then binary search across this revision range by downloading,
unzipping, and opening Chromium for you. After testing the specific revision,
it will ask you whether it is good or bad before continuing the search.
"""

# Base URL to download snapshots from.
BUILD_BASE_URL = 'http://build.chromium.org/buildbot/snapshots/'

# The type (platform) of the build archive. This is what's passed in to the
# '-a/--archive' option.
BUILD_ARCHIVE_TYPE = ''

# The selected archive to bisect.
BUILD_ARCHIVE_DIR = ''

# The location of the builds.
BUILD_ARCHIVE_URL = '/%d/'

# Name of the build archive.
BUILD_ZIP_NAME = ''

# Directory name inside the archive.
BUILD_DIR_NAME = ''

# Name of the executable.
BUILD_EXE_NAME = ''

# URL to the ViewVC commit page.
BUILD_VIEWVC_URL = 'http://src.chromium.org/viewvc/chrome?view=rev&revision=%d'

# Changelogs URL
CHANGELOG_URL = 'http://build.chromium.org/buildbot/' \
                'perf/dashboard/ui/changelog.html?url=/trunk/src&range=%d:%d'

###############################################################################

import math
import optparse
import os
import pipes
import re
import shutil
import sys
import tempfile
import urllib
import zipfile

def UnzipFilenameToDir(filename, dir):
  """Unzip |filename| to directory |dir|."""
  zf = zipfile.ZipFile(filename)
  # Make base.
  pushd = os.getcwd()
  try:
    if not os.path.isdir(dir):
      os.mkdir(dir)
    os.chdir(dir)
    # Extract files.
    for info in zf.infolist():
      name = info.filename
      if name.endswith('/'):  # dir
        if not os.path.isdir(name):
          os.makedirs(name)
      else:  # file
        dir = os.path.dirname(name)
        if not os.path.isdir(dir):
          os.makedirs(dir)
        out = open(name, 'wb')
        out.write(zf.read(name))
        out.close()
      # Set permissions. Permission info in external_attr is shifted 16 bits.
      os.chmod(name, info.external_attr >> 16L)
    os.chdir(pushd)
  except Exception, e:
    print >>sys.stderr, e
    sys.exit(1)

def SetArchiveVars(archive):
  """Set a bunch of global variables appropriate for the specified archive."""
  global BUILD_ARCHIVE_TYPE
  global BUILD_ARCHIVE_DIR
  global BUILD_ZIP_NAME
  global BUILD_DIR_NAME
  global BUILD_EXE_NAME
  global BUILD_BASE_URL

  BUILD_ARCHIVE_TYPE = archive
  BUILD_ARCHIVE_DIR = 'chromium-rel-' + BUILD_ARCHIVE_TYPE

  if BUILD_ARCHIVE_TYPE in ('linux', 'linux-64', 'linux-chromiumos'):
    BUILD_ZIP_NAME = 'chrome-linux.zip'
    BUILD_DIR_NAME = 'chrome-linux'
    BUILD_EXE_NAME = 'chrome'
  elif BUILD_ARCHIVE_TYPE in ('mac'):
    BUILD_ZIP_NAME = 'chrome-mac.zip'
    BUILD_DIR_NAME = 'chrome-mac'
    BUILD_EXE_NAME = 'Chromium.app/Contents/MacOS/Chromium'
  elif BUILD_ARCHIVE_TYPE in ('xp'):
    BUILD_ZIP_NAME = 'chrome-win32.zip'
    BUILD_DIR_NAME = 'chrome-win32'
    BUILD_EXE_NAME = 'chrome.exe'

  BUILD_BASE_URL += BUILD_ARCHIVE_DIR

def ParseDirectoryIndex(url):
  """Parses the HTML directory listing into a list of revision numbers."""
  handle = urllib.urlopen(url)
  dirindex = handle.read()
  handle.close()
  return re.findall(r'\1/', dirindex)

def GetRevList(good, bad):
  """Gets the list of revision numbers between |good| and |bad|."""
  # Download the main revlist.
  revlist = ParseDirectoryIndex(BUILD_BASE_URL)
  revlist = map(int, revlist)
  revlist = filter(lambda r: range(good, bad).__contains__(int(r)), revlist)
  revlist.sort()
  return revlist

def TryRevision(rev, profile, args):
  """Downloads revision |rev|, unzips it, and opens it for the user to test.
  |profile| is the profile to use."""
  # Do this in a temp dir so we don't collide with user files.
  cwd = os.getcwd()
  tempdir = tempfile.mkdtemp(prefix='bisect_tmp')
  os.chdir(tempdir)

  # Download the file.
  download_url = BUILD_BASE_URL + (BUILD_ARCHIVE_URL % rev) + BUILD_ZIP_NAME
  try:
    print 'Fetching ' + download_url
    urllib.urlretrieve(download_url, BUILD_ZIP_NAME)
  except Exception, e:
    print('Could not retrieve the download. Sorry.')
    sys.exit(-1)

  # Unzip the file.
  print 'Unziping ...'
  UnzipFilenameToDir(BUILD_ZIP_NAME, os.curdir)

  # Tell the system to open the app.
  args = ['--user-data-dir=%s' % profile] + args
  flags = ' '.join(map(pipes.quote, args))
  exe = os.path.join(os.getcwd(), BUILD_DIR_NAME, BUILD_EXE_NAME)
  cmd = '%s %s' % (exe, flags)
  print 'Running %s' % cmd
  os.system(cmd)

  os.chdir(cwd)
  print 'Cleaning temp dir ...'
  try:
    shutil.rmtree(tempdir, True)
  except Exception, e:
    pass

def AskIsGoodBuild(rev):
  """Ask the user whether build |rev| is good or bad."""
  # Loop until we get a response that we can parse.
  while True:
    response = raw_input('\nBuild %d is [(g)ood/(b)ad]: ' % int(rev))
    if response and response in ('g', 'b'):
      return response == 'g'

def main():
  usage = ('%prog [options] [-- chromium-options]\n'
           'Perform binary search on the snapshot builds.\n'
           '\n'
           'Tip: add "-- --no-first-run" to bypass the first run prompts.')
  parser = optparse.OptionParser(usage=usage)
  # Strangely, the default help output doesn't include the choice list.
  choices = ['mac', 'xp', 'linux', 'linux-64', 'linux-chromiumos']
  parser.add_option('-a', '--archive',
                    choices = choices,
                    help = 'The buildbot archive to bisect [%s].' %
                           '|'.join(choices))
  parser.add_option('-b', '--bad', type = 'int',
                    help = 'The bad revision to bisect to.')
  parser.add_option('-g', '--good', type = 'int',
                    help = 'The last known good revision to bisect from.')
  parser.add_option('-p', '--profile', '--user-data-dir', type = 'str',
                    help = 'Profile to use; this will not reset every run. ' +
                    'Defaults to a clean profile.')
  (opts, args) = parser.parse_args()

  if opts.archive is None:
    print 'Error: missing required parameter: --archive'
    print
    parser.print_help()
    return 1

  if opts.bad and opts.good and (opts.good > opts.bad):
    print ('The good revision (%d) must precede the bad revision (%d).\n' %
           (opts.good, opts.bad))
    parser.print_help()
    return 1

  SetArchiveVars(opts.archive)

  # Pick a starting point, try to get HEAD for this.
  if opts.bad:
    bad_rev = opts.bad
  else:
    bad_rev = 0
    try:
      # Location of the latest build revision number
      BUILD_LATEST_URL = '%s/LATEST' % (BUILD_BASE_URL)
      nh = urllib.urlopen(BUILD_LATEST_URL)
      latest = int(nh.read())
      nh.close()
      bad_rev = raw_input('Bad revision [HEAD:%d]: ' % latest)
      if (bad_rev == ''):
        bad_rev = latest
      bad_rev = int(bad_rev)
    except Exception, e:
      print('Could not determine latest revision. This could be bad...')
      bad_rev = int(raw_input('Bad revision: '))

  # Find out when we were good.
  if opts.good:
    good_rev = opts.good
  else:
    good_rev = 0
    try:
      good_rev = int(raw_input('Last known good [0]: '))
    except Exception, e:
      pass

  # Get a list of revisions to bisect across.
  revlist = GetRevList(good_rev, bad_rev)
  if len(revlist) < 2:  # Don't have enough builds to bisect
    print 'We don\'t have enough builds to bisect. revlist: %s' % revlist
    sys.exit(1)

  # If we don't have a |good_rev|, set it to be the first revision possible.
  if good_rev == 0:
    good_rev = revlist[0]

  # These are indexes of |revlist|.
  good = 0
  bad = len(revlist) - 1
  last_known_good_rev = revlist[good]

  # Binary search time!
  while good < bad:
    candidates = revlist[good:bad]
    num_poss = len(candidates)
    if num_poss > 10:
      print('%d candidates. %d tries left.' %
          (num_poss, round(math.log(num_poss, 2))))
    else:
      print('Candidates: %s' % revlist[good:bad])

    # Cut the problem in half...
    test = int((bad - good) / 2) + good
    test_rev = revlist[test]

    # Let the user give this rev a spin (in her own profile, if she wants).
    profile = opts.profile
    if not profile:
      profile = 'profile'  # In a temp dir.
    TryRevision(test_rev, profile, args)
    if AskIsGoodBuild(test_rev):
      last_known_good_rev = revlist[good]
      good = test + 1
    else:
      bad = test

  # We're done. Let the user know the results in an official manner.
  print('You are probably looking for build %d.' % revlist[bad])
  print('CHANGELOG URL:')
  print(CHANGELOG_URL % (last_known_good_rev, revlist[bad]))
  print('Built at revision:')
  print(BUILD_VIEWVC_URL % revlist[bad])

if __name__ == '__main__':
  sys.exit(main())
<|endoftext|>


<commit_before># vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (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.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
Base classes for our unit tests.
Allows overriding of flags for use of fakes,
and some black magic for inline callbacks.
"""

import datetime
import unittest

import mox
import stubout

from nova import context
from nova import db
from nova import fakerabbit
from nova import flags
from nova import rpc
from nova.network import manager as network_manager
from nova.tests import fake_flags

FLAGS = flags.FLAGS
flags.DEFINE_bool('flush_db', True,
                  'Flush the database before running fake tests')
flags.DEFINE_bool('fake_tests', True,
                  'should we use everything for testing')

def skip_if_fake(func):
    """Decorator that skips a test if running in fake mode"""
    def _skipper(*args, **kw):
        """Wrapped skipper function"""
        if FLAGS.fake_tests:
            raise unittest.SkipTest('Test cannot be run in fake mode')
        else:
            return func(*args, **kw)
    return _skipper

class TestCase(unittest.TestCase):
    """Test case base class for all unit tests"""
    def setUp(self):
        """Run before each test method to initialize test environment"""
        super(TestCase, self).setUp()
        # NOTE(vish): We need a better method for creating fixtures for tests
        #             now that we have some required db setup for the system
        #             to work properly.
        self.start = datetime.datetime.utcnow()
        ctxt = context.get_admin_context()
        if db.network_count(ctxt) != 5:
            network_manager.VlanManager().create_networks(ctxt,
                                                          FLAGS.fixed_range,
                                                          5, 16,
                                                          FLAGS.fixed_range_v6,
                                                          FLAGS.vlan_start,
                                                          FLAGS.vpn_start,
                                                          )

        # emulate some of the mox stuff, we can't use the metaclass
        # because it screws with our generators
        self.mox = mox.Mox()
        self.stubs = stubout.StubOutForTesting()
        self.flag_overrides = {}
        self.injected = []
        self._monkey_patch_attach()
        self._original_flags = FLAGS.FlagValuesDict()

    def tearDown(self):
        """Runs after each test method to finalize/tear down test
        environment."""
        try:
            self.mox.UnsetStubs()
            self.stubs.UnsetAll()
            self.stubs.SmartUnsetAll()
            self.mox.VerifyAll()
            # NOTE(vish): Clean up any ips associated during the test.
            ctxt = context.get_admin_context()
            db.fixed_ip_disassociate_all_by_timeout(ctxt, FLAGS.host,
                                                    self.start)
            db.network_disassociate_all(ctxt)
            rpc.Consumer.attach_to_eventlet = self.originalAttach
            for x in self.injected:
                try:
                    x.stop()
                except AssertionError:
                    pass

            if FLAGS.fake_rabbit:
                fakerabbit.reset_all()

            db.security_group_destroy_all(ctxt)
            super(TestCase, self).tearDown()
        finally:
            self.reset_flags()

    def flags(self, **kw):
        """Override flag variables for a test"""
        for k, v in kw.iteritems():
            if k in self.flag_overrides:
                self.reset_flags()
                raise Exception(
                        'trying to override already overriden flag: %s' % k)
            self.flag_overrides[k] = getattr(FLAGS, k)
            setattr(FLAGS, k, v)

    def reset_flags(self):
        """Resets all flag variables for the test.  Runs after each test"""
        FLAGS.Reset()
        for k, v in self._original_flags.iteritems():
            setattr(FLAGS, k, v)

    def _monkey_patch_attach(self):
        self.originalAttach = rpc.Consumer.attach_to_eventlet

        def _wrapped(innerSelf):
            rv = self.originalAttach(innerSelf)
            self.injected.append(rv)
            return rv

        _wrapped.func_name = self.originalAttach.func_name
        rpc.Consumer.attach_to_eventlet = _wrapped
<commit_msg>add a start_service method to our test baseclass<commit_after># vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2010 United States Government as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (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.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
Base classes for our unit tests.
Allows overriding of flags for use of fakes,
and some black magic for inline callbacks.
"""

import datetime
import uuid
import unittest

import mox
import stubout

from nova import context
from nova import db
from nova import fakerabbit
from nova import flags
from nova import rpc
from nova import service
from nova.network import manager as network_manager
from nova.tests import fake_flags

FLAGS = flags.FLAGS
flags.DEFINE_bool('flush_db', True,
                  'Flush the database before running fake tests')
flags.DEFINE_bool('fake_tests', True,
                  'should we use everything for testing')

def skip_if_fake(func):
    """Decorator that skips a test if running in fake mode"""
    def _skipper(*args, **kw):
        """Wrapped skipper function"""
        if FLAGS.fake_tests:
            raise unittest.SkipTest('Test cannot be run in fake mode')
        else:
            return func(*args, **kw)
    return _skipper

class TestCase(unittest.TestCase):
    """Test case base class for all unit tests"""
    def setUp(self):
        """Run before each test method to initialize test environment"""
        super(TestCase, self).setUp()
        # NOTE(vish): We need a better method for creating fixtures for tests
        #             now that we have some required db setup for the system
        #             to work properly.
        self.start = datetime.datetime.utcnow()
        ctxt = context.get_admin_context()
        if db.network_count(ctxt) != 5:
            network_manager.VlanManager().create_networks(ctxt,
                                                          FLAGS.fixed_range,
                                                          5, 16,
                                                          FLAGS.fixed_range_v6,
                                                          FLAGS.vlan_start,
                                                          FLAGS.vpn_start,
                                                          )

        # emulate some of the mox stuff, we can't use the metaclass
        # because it screws with our generators
        self.mox = mox.Mox()
        self.stubs = stubout.StubOutForTesting()
        self.flag_overrides = {}
        self.injected = []
        self._services = []
        self._monkey_patch_attach()
        self._original_flags = FLAGS.FlagValuesDict()

    def tearDown(self):
        """Runs after each test method to finalize/tear down test
        environment."""
        try:
            self.mox.UnsetStubs()
            self.stubs.UnsetAll()
            self.stubs.SmartUnsetAll()
            self.mox.VerifyAll()
            super(TestCase, self).tearDown()
        finally:
            try:
                # Clean up any ips associated during the test.
                ctxt = context.get_admin_context()
                db.fixed_ip_disassociate_all_by_timeout(ctxt, FLAGS.host,
                                                        self.start)
                db.network_disassociate_all(ctxt)

                db.security_group_destroy_all(ctxt)
            except Exception:
                pass

            # Clean out fake_rabbit's queue if we used it
            if FLAGS.fake_rabbit:
                fakerabbit.reset_all()

            # Reset any overriden flags
            self.reset_flags()

            # Reset our monkey-patches
            rpc.Consumer.attach_to_eventlet = self.originalAttach

            # Stop any timers
            for x in self.injected:
                try:
                    x.stop()
                except AssertionError:
                    pass

            # Kill any services
            for x in self._services:
                try:
                    x.kill()
                except Exception:
                    pass

    def flags(self, **kw):
        """Override flag variables for a test"""
        for k, v in kw.iteritems():
            if k in self.flag_overrides:
                self.reset_flags()
                raise Exception(
                        'trying to override already overriden flag: %s' % k)
            self.flag_overrides[k] = getattr(FLAGS, k)
            setattr(FLAGS, k, v)

    def reset_flags(self):
        """Resets all flag variables for the test.  Runs after each test"""
        FLAGS.Reset()
        for k, v in self._original_flags.iteritems():
            setattr(FLAGS, k, v)

    def start_service(self, name=None, host=None, **kwargs):
        host = host and host or uuid.uuid4().hex
        kwargs.setdefault('host', host)
        kwargs.setdefault('binary', 'nova-%s' % name)
        svc = service.Service.create(**kwargs)
        svc.start()
        self._services.append(svc)
        return svc

    def _monkey_patch_attach(self):
        self.originalAttach = rpc.Consumer.attach_to_eventlet

        def _wrapped(innerSelf):
            rv = self.originalAttach(innerSelf)
            self.injected.append(rv)
            return rv

        _wrapped.func_name = self.originalAttach.func_name
        rpc.Consumer.attach_to_eventlet = _wrapped
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Contains the crawling models. We use namedtuple for most models (easier to
pickle, lower footprint, indicates that it is immutable) and we use classes for
objects with mutable states and helper methods.

Classes with crawling logic are declared in the crawler module.
"""
from __future__ import unicode_literals, absolute_import

from collections import namedtuple
from optparse import OptionParser, OptionGroup

from pylinkchecker.compat import get_safe_str
from pylinkchecker.urlutil import get_clean_url_split

DEFAULT_TYPES = ['a', 'img', 'script', 'link']

TYPE_ATTRIBUTES = {
    'a': 'href',
    'img': 'src',
    'script': 'src',
    'link': 'href',
}

DEFAULT_TIMEOUT = 10

MODE_THREAD = "thread"
MODE_PROCESS = "process"
MODE_GREEN = "green"

DEFAULT_WORKERS = {
    MODE_THREAD: 1,
    MODE_PROCESS: 1,
    MODE_GREEN: 1000,
}

PARSER_STDLIB = "html.parser"
PARSER_LXML = "lxml"
PARSER_HTML5 = "html5lib"

# TODO Add support for gumbo. Will require some refactoring of the parsing
# logic.
# PARSER_GUMBO = "gumbo"

FORMAT_PLAIN = "plain"
FORMAT_HTML = "html"
FORMAT_JSON = "json"

WHEN_ALWAYS = "always"
WHEN_ON_ERROR = "error"

REPORT_TYPE_ERRORS = "errors"
REPORT_TYPE_SUMMARY = "summary"
REPORT_TYPE_ALL = "all"

VERBOSE_QUIET = "0"
VERBOSE_NORMAL = "1"
VERBOSE_INFO = "2"

HTML_MIME_TYPE = "text/html"

PAGE_QUEUED = '__PAGE_QUEUED__'
PAGE_CRAWLED = '__PAGE_CRAWLED__'

# Note: we use namedtuple to exchange data with workers because they are
# immutable and easy to pickle (as opposed to a class).

WorkerInit = namedtuple("WorkerInit", ["worker_config", "input_queue",
        "output_queue", "logger"])

WorkerConfig = namedtuple("WorkerConfig", ["username", "password", "types",
        "timeout", "parser"])

WorkerInput = namedtuple("WorkerInput", ["url_split", "should_crawl"])

Response = namedtuple("Response", ["content", "status", "exception",
        "original_url", "final_url", "is_redirect", "is_timeout"])

ExceptionStr = namedtuple("ExceptionStr", ["type_name", "message"])

Link = namedtuple("Link", ["type", "url_split", "original_url_split",
        "source_str"])

PageCrawl = namedtuple("PageCrawl", ["original_url_split", "final_url_split",
        "status", "is_timeout", "is_redirect", "links", "exception", "is_html"])

PageStatus = namedtuple("PageStatus", ["status", "sources"])

PageSource = namedtuple("PageSource", ["origin", "origin_str"])

class UTF8Class(object):
    """Handles unicode string from __unicode__() in: __str__() and __repr__()
    """
    def __str__(self):
        return get_safe_str(self.__unicode__())

    def __repr__(self):
        return get_safe_str(self.__unicode__())

class LazyLogParam(object):
    """Lazy Log Parameter that is only evaluated if the logging statement
       is printed"""

    def __init__(self, func):
        self.func=func

    def __str__(self):
        return str(self.func())

class Config(UTF8Class):
    """Contains all the configuration options."""

    def __init__(self):
        # Design note: we only use attributes when options need to be
        # transformed. Otherwise, we use options.
        self.parser = self._build_parser()
        self.options = None
        self.start_urls = []
        self.worker_config = None
        self.accepted_hosts = []
        self.ignored_prefixes = []
        self.worker_size = 0

    def should_crawl(self, url_split):
        """Returns True if url split is local AND run_once is False"""
        return not self.options.run_once and self.is_local(url_split)

    def is_local(self, url_split):
        """Returns true if url split is in the accepted hosts"""
        return url_split.netloc in self.accepted_hosts

    def should_download(self, url_split):
        """Returns True if the url does not start with an ignored prefix and if
        it is local or outside links are allowed."""
        local = self.is_local(url_split)

        if not self.options.test_outside and not local:
            return False

        url = url_split.geturl()

        for ignored_prefix in self.ignored_prefixes:
            if url.startswith(ignored_prefix):
                return False

        return True

    def parse_cli_config(self):
        """Builds the options and args based on the command line options."""
        (self.options, self.start_urls) = self.parser.parse_args()
        self._parse_config()

    def parse_api_config(self, start_urls, options_dict=None):
        """Builds the options and args based on passed parameters."""
        # TODO Add options
        options = self._get_options(options_dict)
        (self.options, self.start_urls) = self.parser.parse_args(options + start_urls)
        self._parse_config()

    def _get_options(self, options_dict):
        if not options_dict:
            options_dict = {}
        options = []
        for key, value in options_dict.items():
            if isinstance(value, bool) and value:
                options.append("--{0}".format(key))
            else:
                options.append("--{0}={1}".format(key, value))
        return options

    def _parse_config(self):
        self.worker_config = self._build_worker_config(self.options)
        self.accepted_hosts = self._build_accepted_hosts(self.options,
                self.start_urls)

        if self.options.ignored_prefixes:
            self.ignored_prefixes = self.options.ignored_prefixes.split(',')

        if self.options.workers:
            self.worker_size = self.options.workers
        else:
            self.worker_size = DEFAULT_WORKERS[self.options.mode]

    def _build_worker_config(self, options):
        types = options.types.split(',')
        for element_type in types:
            if element_type not in DEFAULT_TYPES:
                raise ValueError("This type is not supported: {0}"
                        .format(element_type))

        return WorkerConfig(options.username, options.password, types,
                options.timeout, options.parser)

    def _build_accepted_hosts(self, options, start_urls):
        hosts = set()
        urls = []

        if self.options.accepted_hosts:
            urls = self.options.accepted_hosts.split(',')
        urls = urls + start_urls

        for url in urls:
            split_result = get_clean_url_split(url)
            hosts.add(split_result.netloc)

        return hosts

    def _build_parser(self):
        # avoid circular references
        import pylinkchecker
        version = pylinkchecker.__version__

        parser = OptionParser(usage="%prog [options] URL ...",
                version="%prog {0}".format(version))

        parser.add_option("-V", "--verbose", dest="verbose", action="store",
                default=VERBOSE_QUIET, choices=[VERBOSE_QUIET, VERBOSE_NORMAL,
                VERBOSE_INFO])

        crawler_group = OptionGroup(parser, "Crawler Options",
                "These options modify the way the crawler traverses the site.")
        crawler_group.add_option("-O", "--test-outside", dest="test_outside",
                action="store_true", default=False,
                help="fetch resources from other domains without crawling them")
        crawler_group.add_option("-H", "--accepted-hosts",
                dest="accepted_hosts",  action="store", default=None,
                help="comma-separated list of additional hosts to crawl (e.g., "
                "example.com,subdomain.another.com)")
        crawler_group.add_option("-i", "--ignore", dest="ignored_prefixes",
                action="store", default=None,
                help="comma-separated list of host/path prefixes to ignore "
                "(e.g., www.example.com/ignore_this_and_after/)")
        crawler_group.add_option("-u", "--username", dest="username",
                action="store", default=None,
                help="username to use with basic HTTP authentication")
        crawler_group.add_option("-p", "--password", dest="password",
                action="store", default=None,
                help="password to use with basic HTTP authentication")
        # crawler_group.add_option("-U", "--unique", dest="unique",
        #         action="store_true", default=False)
        crawler_group.add_option("-t", "--types", dest="types", action="store",
                default=",".join(DEFAULT_TYPES),
                help="Comma-separated values of tags to look for when crawling"
                "a site. Default (and supported types): a,img,link,script")
        crawler_group.add_option("-T", "--timeout", dest="timeout",
                type="int", action="store", default=DEFAULT_TIMEOUT,
                help="Seconds to wait before considering that a page timed out")
        crawler_group.add_option("-P", "--progress", dest="progress",
                action="store_true", default=False,
                help="Prints crawler progress in the console")
        crawler_group.add_option("-N", "--run-once", dest="run_once",
                action="store_true", default=False,
                help="Only crawl the first page.")
        # TODO Add follow redirect option.

        parser.add_option_group(crawler_group)

        perf_group = OptionGroup(parser, "Performance Options",
                "These options can impact the performance of the crawler.")

        perf_group.add_option("-w", "--workers", dest="workers", action="store",
                default=None, type="int",
                help="Number of workers to spawn")
        perf_group.add_option("-m", "--mode", dest="mode", action="store",
                default=MODE_THREAD, choices=[MODE_THREAD, MODE_PROCESS,
                MODE_GREEN],
                help="Types of workers: thread (default), process, or green")
        perf_group.add_option("-R", "--parser", dest="parser", action="store",
                default=PARSER_STDLIB, choices=[PARSER_STDLIB, PARSER_LXML,
                PARSER_HTML5],
                help="Types of HTML parse: html.parser (default), lxml, html5lib")

        parser.add_option_group(perf_group)

        output_group = OptionGroup(parser, "Output Options",
                "These options change the output of the crawler.")

        output_group.add_option("-f", "--format", dest="format", action="store",
                default=FORMAT_PLAIN, choices=[FORMAT_PLAIN],
                help="Format of the report: plain")
        output_group.add_option("-o", "--output", dest="output", action="store",
                default=None,
                help="Path of the file where the report will be printed.")
        output_group.add_option("-W", "--when", dest="when", action="store",
                default=WHEN_ALWAYS, choices=[WHEN_ALWAYS, WHEN_ON_ERROR],
                help="When to print the report. error (only if a "
                "crawling error occurs) or always (default)")
        output_group.add_option("-E", "--report-type", dest="report_type",
                action="store", default=REPORT_TYPE_ERRORS, choices=[
                REPORT_TYPE_ERRORS, REPORT_TYPE_SUMMARY, REPORT_TYPE_ALL],
                help="Type of report to print: errors (default, summary and "
                "erroneous links), summary, all (summary and all links)")
        output_group.add_option("-c", "--console", dest="console",
                action="store_true", default=False,
                help="Prints report to the console in addition to other output"
                " options such as file or email.")
        crawler_group.add_option("-S", "--show-source", dest="show_source",
                action="store_true", default=False,
                help="Show source of links (html) in the report.")

        parser.add_option_group(output_group)

        email_group = OptionGroup(parser, "Email Options",
                "These options allows the crawler to send a report by email.")

        email_group.add_option("-a", "--address", dest="address", action="store",
                default=None,
                help="Comma-separated list of email addresses used to send a "
                "report")
        email_group.add_option("--from", dest="from_address", action="store",
                default=None,
                help="Email address to use in the from field of the email "
                "(optional)")
        email_group.add_option("-s", "--smtp", dest="smtp", action="store",
                default=None,
                help="Host of the smtp server")
        email_group.add_option("--port", dest="port", action="store",
                default=25, type="int",
                help="Port of the smtp server (optional)")
        email_group.add_option("--tls", dest="tls", action="store_true",
                default=False,
                help="Use TLS with the email server.")
        email_group.add_option("--subject", dest="subject", action="store",
                default=None,
                help="Subject of the email (optional)")
        email_group.add_option("--smtp-username", dest="smtp_username",
                action="store", default=None,
                help="Username to use with the smtp server (optional)")
        email_group.add_option("--smtp-password", dest="smtp_password",
                action="store", default=None,
                help="Password to use with the smtp server (optional)")

        parser.add_option_group(email_group)

        return parser

    def __unicode__(self):
        return "Configuration - Start URLs: {0} - Options: {1}".format(
                self.start_urls, self.options)

class SitePage(UTF8Class):
    """Contains the crawling result for a page.

    This is a class because we need to keep track of the various sources
    linking to this page and it must be modified as the crawl progresses.
    """

    def __init__(self, url_split, status=200, is_timeout=False, exception=None,
            is_html=True, is_local=True):
        self.url_split = url_split

        self.original_source = None
        self.sources = []

        self.type = type
        self.status = status
        self.is_timeout = is_timeout
        self.exception = exception
        self.is_html = is_html
        self.is_local = is_local
        self.is_ok = status and status < 400

    def add_sources(self, page_sources):
        self.sources.extend(page_sources)

    def get_status_message(self):
        if self.status:
            if self.status < 400:
                return "ok ({0})".format(self.status)
            elif self.status == 404:
                return "not found (404)"
        elif self.is_timeout:
            return "error (timeout)"
        elif self.exception:
            return "error ({0}): {1}".format(self.exception.type_name,
                    self.exception.message)
        else:
            return "error"

    def __unicode__(self):
        return "Resource {0} - {1}".format(self.url_split.geturl(), self.status)

<commit_msg>display status code when available. refs #13<commit_after># -*- coding: utf-8 -*-
"""
Contains the crawling models. We use namedtuple for most models (easier to
pickle, lower footprint, indicates that it is immutable) and we use classes for
objects with mutable states and helper methods.

Classes with crawling logic are declared in the crawler module.
"""
from __future__ import unicode_literals, absolute_import

from collections import namedtuple
from optparse import OptionParser, OptionGroup

from pylinkchecker.compat import get_safe_str
from pylinkchecker.urlutil import get_clean_url_split

DEFAULT_TYPES = ['a', 'img', 'script', 'link']

TYPE_ATTRIBUTES = {
    'a': 'href',
    'img': 'src',
    'script': 'src',
    'link': 'href',
}

DEFAULT_TIMEOUT = 10

MODE_THREAD = "thread"
MODE_PROCESS = "process"
MODE_GREEN = "green"

DEFAULT_WORKERS = {
    MODE_THREAD: 1,
    MODE_PROCESS: 1,
    MODE_GREEN: 1000,
}

PARSER_STDLIB = "html.parser"
PARSER_LXML = "lxml"
PARSER_HTML5 = "html5lib"

# TODO Add support for gumbo. Will require some refactoring of the parsing
# logic.
# PARSER_GUMBO = "gumbo"

FORMAT_PLAIN = "plain"
FORMAT_HTML = "html"
FORMAT_JSON = "json"

WHEN_ALWAYS = "always"
WHEN_ON_ERROR = "error"

REPORT_TYPE_ERRORS = "errors"
REPORT_TYPE_SUMMARY = "summary"
REPORT_TYPE_ALL = "all"

VERBOSE_QUIET = "0"
VERBOSE_NORMAL = "1"
VERBOSE_INFO = "2"

HTML_MIME_TYPE = "text/html"

PAGE_QUEUED = '__PAGE_QUEUED__'
PAGE_CRAWLED = '__PAGE_CRAWLED__'

# Note: we use namedtuple to exchange data with workers because they are
# immutable and easy to pickle (as opposed to a class).

WorkerInit = namedtuple("WorkerInit", ["worker_config", "input_queue",
        "output_queue", "logger"])

WorkerConfig = namedtuple("WorkerConfig", ["username", "password", "types",
        "timeout", "parser"])

WorkerInput = namedtuple("WorkerInput", ["url_split", "should_crawl"])

Response = namedtuple("Response", ["content", "status", "exception",
        "original_url", "final_url", "is_redirect", "is_timeout"])

ExceptionStr = namedtuple("ExceptionStr", ["type_name", "message"])

Link = namedtuple("Link", ["type", "url_split", "original_url_split",
        "source_str"])

PageCrawl = namedtuple("PageCrawl", ["original_url_split", "final_url_split",
        "status", "is_timeout", "is_redirect", "links", "exception", "is_html"])

PageStatus = namedtuple("PageStatus", ["status", "sources"])

PageSource = namedtuple("PageSource", ["origin", "origin_str"])

class UTF8Class(object):
    """Handles unicode string from __unicode__() in: __str__() and __repr__()
    """
    def __str__(self):
        return get_safe_str(self.__unicode__())

    def __repr__(self):
        return get_safe_str(self.__unicode__())

class LazyLogParam(object):
    """Lazy Log Parameter that is only evaluated if the logging statement
       is printed"""

    def __init__(self, func):
        self.func=func

    def __str__(self):
        return str(self.func())

class Config(UTF8Class):
    """Contains all the configuration options."""

    def __init__(self):
        # Design note: we only use attributes when options need to be
        # transformed. Otherwise, we use options.
        self.parser = self._build_parser()
        self.options = None
        self.start_urls = []
        self.worker_config = None
        self.accepted_hosts = []
        self.ignored_prefixes = []
        self.worker_size = 0

    def should_crawl(self, url_split):
        """Returns True if url split is local AND run_once is False"""
        return not self.options.run_once and self.is_local(url_split)

    def is_local(self, url_split):
        """Returns true if url split is in the accepted hosts"""
        return url_split.netloc in self.accepted_hosts

    def should_download(self, url_split):
        """Returns True if the url does not start with an ignored prefix and if
        it is local or outside links are allowed."""
        local = self.is_local(url_split)

        if not self.options.test_outside and not local:
            return False

        url = url_split.geturl()

        for ignored_prefix in self.ignored_prefixes:
            if url.startswith(ignored_prefix):
                return False

        return True

    def parse_cli_config(self):
        """Builds the options and args based on the command line options."""
        (self.options, self.start_urls) = self.parser.parse_args()
        self._parse_config()

    def parse_api_config(self, start_urls, options_dict=None):
        """Builds the options and args based on passed parameters."""
        # TODO Add options
        options = self._get_options(options_dict)
        (self.options, self.start_urls) = self.parser.parse_args(options + start_urls)
        self._parse_config()

    def _get_options(self, options_dict):
        if not options_dict:
            options_dict = {}
        options = []
        for key, value in options_dict.items():
            if isinstance(value, bool) and value:
                options.append("--{0}".format(key))
            else:
                options.append("--{0}={1}".format(key, value))
        return options

    def _parse_config(self):
        self.worker_config = self._build_worker_config(self.options)
        self.accepted_hosts = self._build_accepted_hosts(self.options,
                self.start_urls)

        if self.options.ignored_prefixes:
            self.ignored_prefixes = self.options.ignored_prefixes.split(',')

        if self.options.workers:
            self.worker_size = self.options.workers
        else:
            self.worker_size = DEFAULT_WORKERS[self.options.mode]

    def _build_worker_config(self, options):
        types = options.types.split(',')
        for element_type in types:
            if element_type not in DEFAULT_TYPES:
                raise ValueError("This type is not supported: {0}"
                        .format(element_type))

        return WorkerConfig(options.username, options.password, types,
                options.timeout, options.parser)

    def _build_accepted_hosts(self, options, start_urls):
        hosts = set()
        urls = []

        if self.options.accepted_hosts:
            urls = self.options.accepted_hosts.split(',')
        urls = urls + start_urls

        for url in urls:
            split_result = get_clean_url_split(url)
            hosts.add(split_result.netloc)

        return hosts

    def _build_parser(self):
        # avoid circular references
        import pylinkchecker
        version = pylinkchecker.__version__

        parser = OptionParser(usage="%prog [options] URL ...",
                version="%prog {0}".format(version))

        parser.add_option("-V", "--verbose", dest="verbose", action="store",
                default=VERBOSE_QUIET, choices=[VERBOSE_QUIET, VERBOSE_NORMAL,
                VERBOSE_INFO])

        crawler_group = OptionGroup(parser, "Crawler Options",
                "These options modify the way the crawler traverses the site.")
        crawler_group.add_option("-O", "--test-outside", dest="test_outside",
                action="store_true", default=False,
                help="fetch resources from other domains without crawling them")
        crawler_group.add_option("-H", "--accepted-hosts",
                dest="accepted_hosts",  action="store", default=None,
                help="comma-separated list of additional hosts to crawl (e.g., "
                "example.com,subdomain.another.com)")
        crawler_group.add_option("-i", "--ignore", dest="ignored_prefixes",
                action="store", default=None,
                help="comma-separated list of host/path prefixes to ignore "
                "(e.g., www.example.com/ignore_this_and_after/)")
        crawler_group.add_option("-u", "--username", dest="username",
                action="store", default=None,
                help="username to use with basic HTTP authentication")
        crawler_group.add_option("-p", "--password", dest="password",
                action="store", default=None,
                help="password to use with basic HTTP authentication")
        # crawler_group.add_option("-U", "--unique", dest="unique",
        #         action="store_true", default=False)
        crawler_group.add_option("-t", "--types", dest="types", action="store",
                default=",".join(DEFAULT_TYPES),
                help="Comma-separated values of tags to look for when crawling"
                "a site. Default (and supported types): a,img,link,script")
        crawler_group.add_option("-T", "--timeout", dest="timeout",
                type="int", action="store", default=DEFAULT_TIMEOUT,
                help="Seconds to wait before considering that a page timed out")
        crawler_group.add_option("-P", "--progress", dest="progress",
                action="store_true", default=False,
                help="Prints crawler progress in the console")
        crawler_group.add_option("-N", "--run-once", dest="run_once",
                action="store_true", default=False,
                help="Only crawl the first page.")
        # TODO Add follow redirect option.

        parser.add_option_group(crawler_group)

        perf_group = OptionGroup(parser, "Performance Options",
                "These options can impact the performance of the crawler.")

        perf_group.add_option("-w", "--workers", dest="workers", action="store",
                default=None, type="int",
                help="Number of workers to spawn")
        perf_group.add_option("-m", "--mode", dest="mode", action="store",
                default=MODE_THREAD, choices=[MODE_THREAD, MODE_PROCESS,
                MODE_GREEN],
                help="Types of workers: thread (default), process, or green")
        perf_group.add_option("-R", "--parser", dest="parser", action="store",
                default=PARSER_STDLIB, choices=[PARSER_STDLIB, PARSER_LXML,
                PARSER_HTML5],
                help="Types of HTML parse: html.parser (default), lxml, html5lib")

        parser.add_option_group(perf_group)

        output_group = OptionGroup(parser, "Output Options",
                "These options change the output of the crawler.")

        output_group.add_option("-f", "--format", dest="format", action="store",
                default=FORMAT_PLAIN, choices=[FORMAT_PLAIN],
                help="Format of the report: plain")
        output_group.add_option("-o", "--output", dest="output", action="store",
                default=None,
                help="Path of the file where the report will be printed.")
        output_group.add_option("-W", "--when", dest="when", action="store",
                default=WHEN_ALWAYS, choices=[WHEN_ALWAYS, WHEN_ON_ERROR],
                help="When to print the report. error (only if a "
                "crawling error occurs) or always (default)")
        output_group.add_option("-E", "--report-type", dest="report_type",
                action="store", default=REPORT_TYPE_ERRORS, choices=[
                REPORT_TYPE_ERRORS, REPORT_TYPE_SUMMARY, REPORT_TYPE_ALL],
                help="Type of report to print: errors (default, summary and "
                "erroneous links), summary, all (summary and all links)")
        output_group.add_option("-c", "--console", dest="console",
                action="store_true", default=False,
                help="Prints report to the console in addition to other output"
                " options such as file or email.")
        crawler_group.add_option("-S", "--show-source", dest="show_source",
                action="store_true", default=False,
                help="Show source of links (html) in the report.")

        parser.add_option_group(output_group)

        email_group = OptionGroup(parser, "Email Options",
                "These options allows the crawler to send a report by email.")

        email_group.add_option("-a", "--address", dest="address", action="store",
                default=None,
                help="Comma-separated list of email addresses used to send a "
                "report")
        email_group.add_option("--from", dest="from_address", action="store",
                default=None,
                help="Email address to use in the from field of the email "
                "(optional)")
        email_group.add_option("-s", "--smtp", dest="smtp", action="store",
                default=None,
                help="Host of the smtp server")
        email_group.add_option("--port", dest="port", action="store",
                default=25, type="int",
                help="Port of the smtp server (optional)")
        email_group.add_option("--tls", dest="tls", action="store_true",
                default=False,
                help="Use TLS with the email server.")
        email_group.add_option("--subject", dest="subject", action="store",
                default=None,
                help="Subject of the email (optional)")
        email_group.add_option("--smtp-username", dest="smtp_username",
                action="store", default=None,
                help="Username to use with the smtp server (optional)")
        email_group.add_option("--smtp-password", dest="smtp_password",
                action="store", default=None,
                help="Password to use with the smtp server (optional)")

        parser.add_option_group(email_group)

        return parser

    def __unicode__(self):
        return "Configuration - Start URLs: {0} - Options: {1}".format(
                self.start_urls, self.options)

class SitePage(UTF8Class):
    """Contains the crawling result for a page.

    This is a class because we need to keep track of the various sources
    linking to this page and it must be modified as the crawl progresses.
    """

    def __init__(self, url_split, status=200, is_timeout=False, exception=None,
            is_html=True, is_local=True):
        self.url_split = url_split

        self.original_source = None
        self.sources = []

        self.type = type
        self.status = status
        self.is_timeout = is_timeout
        self.exception = exception
        self.is_html = is_html
        self.is_local = is_local
        self.is_ok = status and status < 400

    def add_sources(self, page_sources):
        self.sources.extend(page_sources)

    def get_status_message(self):
        if self.status:
            if self.status < 400:
                return "ok ({0})".format(self.status)
            elif self.status == 404:
                return "not found (404)"
            else:
                return "error (status={0})".format(self.status)
        elif self.is_timeout:
            return "error (timeout)"
        elif self.exception:
            return "error ({0}): {1}".format(self.exception.type_name,
                    self.exception.message)
        else:
            return "error"

    def __unicode__(self):
        return "Resource {0} - {1}".format(self.url_split.geturl(), self.status)

<|endoftext|>