<commit_before>#!/usr/bin/env python

'''Identifies keyword selected domains from InterPro database.'''

import argparse
import re

__description__ = ""
__author__ = "Malvika Sharan <malvika.sharan@uni-wuerzburg.de>"
__email__ = "malvika.sharan@uni-wuerzburg.de"
__version__ = ""

def main():
    parser = argparse.ArgumentParser(description=__description__)
    parser.add_argument("keywords_file")
    parser.add_argument("ipr_whole_data_file")
    parser.add_argument("interpro_mapped_cdd")
    parser.add_argument("domain_data_path")
    args = parser.parse_args()

    keyword_selected_ipr_selection = RnaRelatedIPRSelection(
        args.keywords_file, args.ipr_whole_data_file,
        args.interpro_mapped_cdd, args.domain_data_path)
    keyword_selected_ipr_selection.read_keyword_file()
    keyword_selected_ipr_selection.read_interpro_mapped_cdd_file()
    keyword_selected_ipr_selection.read_ipr_whole_data_file()
    keyword_selected_ipr_selection.create_keyword_selected_domain_file()

class RnaRelatedIPRSelection(object):
    def __init__(self, keywords_file,
                 ipr_whole_data_file,
                 interpro_mapped_cdd,
                 domain_data_path):
        self._keywords_file = keywords_file
        self._ipr_whole_data_file = ipr_whole_data_file
        self._interpro_mapped_cdd = interpro_mapped_cdd
        self._domain_data_path = domain_data_path

        self.ipr_whole_data_list = []
        self._ipr_dict = {}
        self._mapped_interpro_length = {}

    def select_ipr_domains(self):
        '''To call from apricot'''
        self.read_keyword_file()
        self.read_interpro_mapped_cdd_file()
        self.read_ipr_whole_data_file()
        self.create_keyword_selected_domain_file()

    def read_keyword_file(self):
        '''reads user provided keywords for domain selection'''
        self._keyword_list = [
            rna_keyword.strip() for rna_keyword in open(
                self._keywords_file)]
        return self._keyword_list

    def read_interpro_mapped_cdd_file(self):
        '''Parses cdd interpro mapped file for
        the extraction of common information'''
        with open(self._interpro_mapped_cdd, 'r') as in_fh:
            for entry in in_fh:
                ipr_id = entry.split('\t')[0]
                ipr_members = entry.split('\t')[1]
                cdd_id = entry.split('\t')[2]
                cdd_member = entry.split('\t')[3]
                domain_length = entry.split('\t')[4].strip()
                self._mapped_interpro_length[cdd_member] = domain_length
        self._mapped_interpro_length

    def read_ipr_whole_data_file(self):
        '''Parses InterPro annotation data from table'''
        with open(self._ipr_whole_data_file, 'r') as in_fh:
            for ipr_id_entry in in_fh:
                self.ipr_whole_data_list.append(ipr_id_entry.strip())
                self._ipr_dict[ipr_id_entry.strip()] = ipr_id_entry.strip(
                    ).split('\t')[2]
        return self.ipr_whole_data_list, self._ipr_dict

    def create_keyword_selected_domain_file(self):
        '''Creates a file with InterPro derived domain of interest'''
        self._keyword_annotation_dict = {}
        self._keyword_selected_domain = {}
        for ipr_entry in self.ipr_whole_data_list:
            domain_full_name = list(self._ipr_dict[
                ipr_entry].lower().replace('-', ' ').replace(
                    '_', ' ').split(' '))
            for keyword in self._keyword_list:
                check_keyword = keyword.lower()
                if '-' in check_keyword:
                    ref_kw_list = check_keyword.split('-')
                    new_kw_list = []
                    for domain_name in domain_full_name:
                        for string in ref_kw_list:
                            if string == 'rna':
                                if re.findall(
                                        r"(?:\s|^)%s(?=\s|$)" % '*rna',
                                        domain_name):
                                    if not string in set(new_kw_list):
                                        new_kw_list.append(string)
                            else:
                                if re.findall(
                                        r"(?:\s|^)%s(?=\s|$)" % string,
                                        domain_name):
                                    if not string in set(new_kw_list):
                                        new_kw_list.append(string)
                    if str(ref_kw_list) == str(new_kw_list):
                        self._keyword_annotation_dict.setdefault(
                            keyword, []).append(ipr_entry)
                else:
                    for domain_name in domain_full_name:
                        if re.findall(
                                r"(?:\s|^)%s(?=\s|$)" % check_keyword,
                                domain_name):
                            self._keyword_annotation_dict.setdefault(
                                keyword, []).append(ipr_entry)
        for fkeyword in self._keyword_list:
            fkeyword = fkeyword.replace(' ', '_')
            with open(self._domain_data_path+'/'+fkeyword+'_related_ipr_ids.tab',
                      'w') as key_fh:
                if self._keyword_annotation_dict.get(fkeyword):
                    for each_entry in self._keyword_annotation_dict[fkeyword]:
                        each_entry = each_entry.replace(each_entry.split(
                            '\t')[3], " ".join(each_entry.split('\t')[3].split(
                            ))).replace(';', ',')
                        cdd_id = each_entry.split('\t')[1]
                        if cdd_id in set(self._mapped_interpro_length.keys()):
                            length = self._mapped_interpro_length[cdd_id]
                        else:
                            length = 'NA'
                        self._keyword_selected_domain.setdefault("%s\t%s" % (
                            each_entry, length), set()).add(fkeyword)
                        key_fh.write("%s\t%s\t%s\n" % (
                            each_entry, length, fkeyword))
        with open(self._domain_data_path+'/all_keyword_selected_ipr_data.tab',
                  'w') as keyword_selected_domain_file:
            for domain_entry in self._keyword_selected_domain.keys():
                keywords = ', '.join(list(
                    self._keyword_selected_domain[domain_entry]))
                keyword_selected_domain_file.write('%s\t%s\n' % (
                    domain_entry, keywords))

if __name__ == "__main__":

    main()
<commit_msg>Update keyword_selected_ipr_domain.py<commit_after>#!/usr/bin/env python

import argparse
import re

__description__ = "Identifies keyword selected domains from InterPro database."
__author__ = "Malvika Sharan <malvika.sharan@uni-wuerzburg.de>"
__email__ = "malvika.sharan@uni-wuerzburg.de"

def main():
    parser = argparse.ArgumentParser(description=__description__)
    parser.add_argument("keywords_file")
    parser.add_argument("ipr_whole_data_file")
    parser.add_argument("interpro_mapped_cdd")
    parser.add_argument("domain_data_path")
    args = parser.parse_args()

    keyword_selected_ipr_selection = RnaRelatedIPRSelection(
        args.keywords_file, args.ipr_whole_data_file,
        args.interpro_mapped_cdd, args.domain_data_path)
    keyword_selected_ipr_selection.read_keyword_file()
    keyword_selected_ipr_selection.read_interpro_mapped_cdd_file()
    keyword_selected_ipr_selection.read_ipr_whole_data_file()
    keyword_selected_ipr_selection.create_keyword_selected_domain_file()

class RnaRelatedIPRSelection(object):
    def __init__(self, keywords_file,
                 ipr_whole_data_file,
                 interpro_mapped_cdd,
                 domain_data_path):
        self._keywords_file = keywords_file
        self._ipr_whole_data_file = ipr_whole_data_file
        self._interpro_mapped_cdd = interpro_mapped_cdd
        self._domain_data_path = domain_data_path

        self.ipr_whole_data_list = []
        self._ipr_dict = {}
        self._mapped_interpro_length = {}

    def select_ipr_domains(self):
        '''To call from apricot'''
        self.read_keyword_file()
        self.read_interpro_mapped_cdd_file()
        self.read_ipr_whole_data_file()
        self.create_keyword_selected_domain_file()

    def read_keyword_file(self):
        '''reads user provided keywords for domain selection'''
        self._keyword_list = [
            rna_keyword.strip() for rna_keyword in open(
                self._keywords_file)]
        return self._keyword_list

    def read_interpro_mapped_cdd_file(self):
        '''Parses cdd interpro mapped file for
        the extraction of common information'''
        with open(self._interpro_mapped_cdd, 'r') as in_fh:
            for entry in in_fh:
                ipr_id = entry.split('\t')[0]
                ipr_members = entry.split('\t')[1]
                cdd_id = entry.split('\t')[2]
                cdd_member = entry.split('\t')[3]
                domain_length = entry.split('\t')[4].strip()
                self._mapped_interpro_length[cdd_member] = domain_length
        self._mapped_interpro_length

    def read_ipr_whole_data_file(self):
        '''Parses InterPro annotation data from table'''
        with open(self._ipr_whole_data_file, 'r') as in_fh:
            for ipr_id_entry in in_fh:
                self.ipr_whole_data_list.append(ipr_id_entry.strip())
                self._ipr_dict[ipr_id_entry.strip()] = ipr_id_entry.strip(
                    ).split('\t')[2]
        return self.ipr_whole_data_list, self._ipr_dict

    def create_keyword_selected_domain_file(self):
        '''Creates a file with InterPro derived domain of interest'''
        self._keyword_annotation_dict = {}
        self._keyword_selected_domain = {}
        for ipr_entry in self.ipr_whole_data_list:
            domain_full_name = list(self._ipr_dict[
                ipr_entry].lower().replace('-', ' ').replace(
                    '_', ' ').split(' '))
            for keyword in self._keyword_list:
                check_keyword = keyword.lower()
                if '-' in check_keyword:
                    ref_kw_list = check_keyword.split('-')
                    new_kw_list = []
                    for domain_name in domain_full_name:
                        for string in ref_kw_list:
                            if string == 'rna':
                                if re.findall(
                                        r"(?:\s|^)%s(?=\s|$)" % '*rna',
                                        domain_name):
                                    if not string in set(new_kw_list):
                                        new_kw_list.append(string)
                            else:
                                if re.findall(
                                        r"(?:\s|^)%s(?=\s|$)" % string,
                                        domain_name):
                                    if not string in set(new_kw_list):
                                        new_kw_list.append(string)
                    if str(ref_kw_list) == str(new_kw_list):
                        self._keyword_annotation_dict.setdefault(
                            keyword, []).append(ipr_entry)
                else:
                    for domain_name in domain_full_name:
                        if re.findall(
                                r"(?:\s|^)%s(?=\s|$)" % check_keyword,
                                domain_name):
                            self._keyword_annotation_dict.setdefault(
                                keyword, []).append(ipr_entry)
        for fkeyword in self._keyword_list:
            fkeyword = fkeyword.replace(' ', '_')
            with open(self._domain_data_path+'/'+fkeyword+'_related_ipr_ids.tab',
                      'w') as key_fh:
                if self._keyword_annotation_dict.get(fkeyword):
                    for each_entry in self._keyword_annotation_dict[fkeyword]:
                        each_entry = each_entry.replace(each_entry.split(
                            '\t')[3], " ".join(each_entry.split('\t')[3].split(
                            ))).replace(';', ',')
                        cdd_id = each_entry.split('\t')[1]
                        if cdd_id in set(self._mapped_interpro_length.keys()):
                            length = self._mapped_interpro_length[cdd_id]
                        else:
                            length = 'NA'
                        self._keyword_selected_domain.setdefault("%s\t%s" % (
                            each_entry, length), set()).add(fkeyword)
                        key_fh.write("%s\t%s\t%s\n" % (
                            each_entry, length, fkeyword))
        with open(self._domain_data_path+'/all_keyword_selected_ipr_data.tab',
                  'w') as keyword_selected_domain_file:
            for domain_entry in self._keyword_selected_domain.keys():
                keywords = ', '.join(list(
                    self._keyword_selected_domain[domain_entry]))
                keyword_selected_domain_file.write('%s\t%s\n' % (
                    domain_entry, keywords))

if __name__ == "__main__":

    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python

from PyQt4 import QtGui, QtCore
from item import Item

class PointItem(Item):

    def __init__(self, scene, rpm, x, y, parent, fixed=False):

        self.outline_diameter = 0.025

        super(PointItem, self).__init__(scene, rpm, x, y, parent)
        self.fixed = fixed

    def boundingRect(self):

        es = self.m2pix(self.outline_diameter * 1.8)

        return QtCore.QRectF(-es / 2, -es / 2, es, es)

    def shape(self):

        path = QtGui.QPainterPath()
        es = self.m2pix(self.outline_diameter)
        path.addEllipse(QtCore.QPoint(0, 0), es / 2, es / 2)
        return path

    def item_moved(self):

        self.parentItem().point_changed()  # TODO change it to callback

    def cursor_release(self):

        self.parentItem().point_changed(True)

    def paint(self, painter, option, widget):

        painter.setClipRect(option.exposedRect)
        painter.setRenderHint(QtGui.QPainter.Antialiasing)

        es = self.m2pix(self.outline_diameter)

        if self.hover:
            # TODO coordinates
            painter.setBrush(QtCore.Qt.gray)
            painter.setPen(QtCore.Qt.gray)
            painter.drawEllipse(QtCore.QPoint(0, 0), es / 2 * 1.8, es / 2 * 1.8)

        painter.setBrush(QtCore.Qt.cyan)
        painter.setPen(QtCore.Qt.cyan)

        painter.drawEllipse(QtCore.QPoint(0, 0), es / 2, es / 2)
<commit_msg>format<commit_after>#!/usr/bin/env python

from PyQt4 import QtGui, QtCore
from item import Item

class PointItem(Item):

    def __init__(self, scene, rpm, x, y, parent, fixed=False):

        self.outline_diameter = 0.025

        super(PointItem, self).__init__(scene, rpm, x, y, parent)
        self.fixed = fixed

    def boundingRect(self):

        es = self.m2pix(self.outline_diameter * 1.8)

        return QtCore.QRectF(-es / 2, -es / 2, es, es)

    def shape(self):

        path = QtGui.QPainterPath()
        es = self.m2pix(self.outline_diameter)
        path.addEllipse(QtCore.QPoint(0, 0), es / 2, es / 2)
        return path

    def item_moved(self):

        self.parentItem().point_changed()  # TODO change it to callback

    def cursor_release(self):

        self.parentItem().point_changed(True)

    def paint(self, painter, option, widget):

        painter.setClipRect(option.exposedRect)
        painter.setRenderHint(QtGui.QPainter.Antialiasing)

        es = self.m2pix(self.outline_diameter)

        if self.hover:
            # TODO coordinates
            painter.setBrush(QtCore.Qt.gray)
            painter.setPen(QtCore.Qt.gray)
            painter.drawEllipse(QtCore.QPoint(0, 0), es / 2 * 1.8, es / 2 * 1.8)

        painter.setBrush(QtCore.Qt.cyan)
        painter.setPen(QtCore.Qt.cyan)

        painter.drawEllipse(QtCore.QPoint(0, 0), es / 2, es / 2)
<|endoftext|>


<commit_before>import re
import json
import logging

from django.conf import settings
from functools import wraps
from static_replace import replace_urls
from mitxmako.shortcuts import render_to_string
from xmodule.seq_module import SequenceModule
from xmodule.vertical_module import VerticalModule

log = logging.getLogger("mitx.xmodule_modifiers")

def wrap_xmodule(get_html, module, template):
    """
    Wraps the results of get_html in a standard <section> with identifying
    data so that the appropriate javascript module can be loaded onto it.

    get_html: An XModule.get_html method or an XModuleDescriptor.get_html method
    module: An XModule
    template: A template that takes the variables:
        content: the results of get_html,
        class_: the module class name
        module_name: the js_module_name of the module
    """

    @wraps(get_html)
    def _get_html():
        return render_to_string(template, {
            'content': get_html(),
            'class_': module.__class__.__name__,
            'module_name': module.js_module_name
        })
    return _get_html

def replace_static_urls(get_html, prefix, module):
    """
    Updates the supplied module with a new get_html function that wraps
    the old get_html function and substitutes urls of the form /static/...
    with urls that are /static/<prefix>/...
    """

    @wraps(get_html)
    def _get_html():
        if type(module) in [SequenceModule, VerticalModule]:	# TODO: make this more general, eg use an XModule attribute instead
            return get_html()
        return replace_urls(get_html(), staticfiles_prefix=prefix)
    return _get_html

def grade_histogram(module_id):
    ''' Print out a histogram of grades on a given problem.
        Part of staff member debug info.
    '''
    from django.db import connection
    cursor = connection.cursor()

    q = """SELECT courseware_studentmodule.grade,
                  COUNT(courseware_studentmodule.student_id)
    FROM courseware_studentmodule
    WHERE courseware_studentmodule.module_id=%s
    GROUP BY courseware_studentmodule.grade"""
    # Passing module_id this way prevents sql-injection.
    cursor.execute(q, [module_id])

    grades = list(cursor.fetchall())
    grades.sort(key=lambda x: x[0])          # Add ORDER BY to sql query?
    if len(grades) == 1 and grades[0][0] is None:
        return []
    return grades

def add_histogram(get_html, module):
    """
    Updates the supplied module with a new get_html function that wraps
    the output of the old get_html function with additional information
    for admin users only, including a histogram of student answers and the
    definition of the xmodule

    Does nothing if module is a SequenceModule
    """
    @wraps(get_html)
    def _get_html():

        if type(module) in [SequenceModule, VerticalModule]:	# TODO: make this more general, eg use an XModule attribute instead
            return get_html()

        module_id = module.id
        histogram = grade_histogram(module_id)
        render_histogram = len(histogram) > 0

        # TODO (ichuang): Remove after fall 2012 LMS migration done
        if settings.MITX_FEATURES.get('ENABLE_LMS_MIGRATION'):
            [filepath, filename] = module.definition.get('filename', ['', None])
            osfs = module.system.filestore
            if filename is not None and osfs.exists(filename):
                # if original, unmangled filename exists then use it (github
                # doesn't like symlinks)
                filepath = filename
            data_dir = osfs.root_path.rsplit('/')[-1]
            edit_link = "https://github.com/MITx/%s/tree/master/%s" % (data_dir,filepath)
        else:
            edit_link = False

        staff_context = {'definition': module.definition.get('data'),
                         'metadata': json.dumps(module.metadata, indent=4),
                         'element_id': module.location.html_id(),
                         'edit_link': edit_link,
                         'histogram': json.dumps(histogram),
                         'render_histogram': render_histogram,
                         'module_content': get_html()}
        return render_to_string("staff_problem_info.html", staff_context)

    return _get_html

<commit_msg>minor change to allow github url to be overridden by metadata<commit_after>import re
import json
import logging

from django.conf import settings
from functools import wraps
from static_replace import replace_urls
from mitxmako.shortcuts import render_to_string
from xmodule.seq_module import SequenceModule
from xmodule.vertical_module import VerticalModule

log = logging.getLogger("mitx.xmodule_modifiers")

def wrap_xmodule(get_html, module, template):
    """
    Wraps the results of get_html in a standard <section> with identifying
    data so that the appropriate javascript module can be loaded onto it.

    get_html: An XModule.get_html method or an XModuleDescriptor.get_html method
    module: An XModule
    template: A template that takes the variables:
        content: the results of get_html,
        class_: the module class name
        module_name: the js_module_name of the module
    """

    @wraps(get_html)
    def _get_html():
        return render_to_string(template, {
            'content': get_html(),
            'class_': module.__class__.__name__,
            'module_name': module.js_module_name
        })
    return _get_html

def replace_static_urls(get_html, prefix, module):
    """
    Updates the supplied module with a new get_html function that wraps
    the old get_html function and substitutes urls of the form /static/...
    with urls that are /static/<prefix>/...
    """

    @wraps(get_html)
    def _get_html():
        if type(module) in [SequenceModule, VerticalModule]:	# TODO: make this more general, eg use an XModule attribute instead
            return get_html()
        return replace_urls(get_html(), staticfiles_prefix=prefix)
    return _get_html

def grade_histogram(module_id):
    ''' Print out a histogram of grades on a given problem.
        Part of staff member debug info.
    '''
    from django.db import connection
    cursor = connection.cursor()

    q = """SELECT courseware_studentmodule.grade,
                  COUNT(courseware_studentmodule.student_id)
    FROM courseware_studentmodule
    WHERE courseware_studentmodule.module_id=%s
    GROUP BY courseware_studentmodule.grade"""
    # Passing module_id this way prevents sql-injection.
    cursor.execute(q, [module_id])

    grades = list(cursor.fetchall())
    grades.sort(key=lambda x: x[0])          # Add ORDER BY to sql query?
    if len(grades) == 1 and grades[0][0] is None:
        return []
    return grades

def add_histogram(get_html, module):
    """
    Updates the supplied module with a new get_html function that wraps
    the output of the old get_html function with additional information
    for admin users only, including a histogram of student answers and the
    definition of the xmodule

    Does nothing if module is a SequenceModule
    """
    @wraps(get_html)
    def _get_html():

        if type(module) in [SequenceModule, VerticalModule]:	# TODO: make this more general, eg use an XModule attribute instead
            return get_html()

        module_id = module.id
        histogram = grade_histogram(module_id)
        render_histogram = len(histogram) > 0

        # TODO (ichuang): Remove after fall 2012 LMS migration done
        if settings.MITX_FEATURES.get('ENABLE_LMS_MIGRATION'):
            [filepath, filename] = module.definition.get('filename', ['', None])
            osfs = module.system.filestore
            if filename is not None and osfs.exists(filename):
                # if original, unmangled filename exists then use it (github
                # doesn't like symlinks)
                filepath = filename
            data_dir = osfs.root_path.rsplit('/')[-1]
            giturl = module.metadata.get('giturl','https://github.com/MITx')
            edit_link = "%s/%s/tree/master/%s" % (giturl,data_dir,filepath)
        else:
            edit_link = False

        staff_context = {'definition': module.definition.get('data'),
                         'metadata': json.dumps(module.metadata, indent=4),
                         'element_id': module.location.html_id(),
                         'edit_link': edit_link,
                         'histogram': json.dumps(histogram),
                         'render_histogram': render_histogram,
                         'module_content': get_html()}
        return render_to_string("staff_problem_info.html", staff_context)

    return _get_html

<|endoftext|>


<commit_before>'''
calibre_plugins.comicvine - A calibre metadata source for comicvine
'''
#pylint: disable-msg=R0913,R0904
from functools import partial
import logging
from Queue import Queue

from calibre import setup_cli_handlers
from calibre.ebooks.metadata.sources.base import Source
from calibre.utils.config import OptionParser
import calibre.utils.logging as calibre_logging
from calibre_plugins.comicvine.config import PREFS
from calibre_plugins.comicvine import utils

class Comicvine(Source):
  ''' Metadata source implementation '''
  name = 'Comicvine'
  description = 'Downloads metadata and covers from Comicvine'
  author = 'Russell Heilling'
  version = (0, 7, 4)
  #TODO(xchewtoyx): Implement cover capability
  capabilities = frozenset(['identify', 'cover'])
  touched_fields = frozenset([
      'title', 'authors', 'identifier:comicvine', 'comments', 'publisher',
      'pubdate', 'series'
      ])

  has_html_comments = True
  can_get_multiple_covers = True

  def __init__(self, *args, **kwargs):
    self.logger = logging.getLogger('urls')
    self.logger.setLevel(logging.DEBUG)
    self.logger.addHandler(utils.CalibreHandler(logging.DEBUG))
    Source.__init__(self, *args, **kwargs)

  def config_widget(self):
    from calibre_plugins.comicvine.config import ConfigWidget
    return ConfigWidget()

  def save_settings(self, config_widget):
    config_widget.save_settings()

  def is_configured(self):
    return bool(PREFS.get('api_key'))

  def cli_main(self, args):
    'Perform comicvine lookups from the calibre-debug cli'
    def option_parser():
      'Parse command line options'
      parser = OptionParser(
        usage='Comicvine [t:title] [a:authors] [i:id]')
      parser.add_option('--verbose', '-v', default=False,
                        action='store_true', dest='verbose')
      parser.add_option('--debug_api', default=False,
                        action='store_true', dest='debug_api')
      return parser

    opts, args = option_parser().parse_args(args)
    if opts.debug_api:
      calibre_logging.default_log = calibre_logging.Log(
        level=calibre_logging.DEBUG)
    if opts.verbose:
      level = 'DEBUG'
    else:
      level = 'INFO'
    setup_cli_handlers(logging.getLogger('comicvine'),
                       getattr(logging, level))
    log = calibre_logging.ThreadSafeLog(level=getattr(calibre_logging, level))

    (title, authors, ids) = (None, [], {})
    for arg in args:
      if arg.startswith('t:'):
        title = arg.split(':', 1)[1]
      if arg.startswith('a:'):
        authors.append(arg.split(':', 1)[1])
      if arg.startswith('i:'):
        (idtype, identifier) = arg.split(':', 2)[1:]
        ids[idtype] = int(identifier)
    result_queue = Queue()
    self.identify(
      log, result_queue, False, title=title, authors=authors, identifiers=ids)
    ranking = self.identify_results_keygen(title, authors, ids)
    for result in sorted(result_queue.queue, key=ranking):
      if result.pubdate:
        pubdate = str(result.pubdate.date())
      else:
        pubdate = 'Unknown'
      log.info('(%04d) - %s: %s [%s]' % (
          ranking(result), result.identifiers['comicvine'],
          result.title, pubdate))

  def enqueue(self, log, result_queue, issue_id):
    'Add a result entry to the result queue'
    log.debug('Adding Issue(%d) to queue' % issue_id)
    metadata = utils.build_meta(log, issue_id)
    if metadata:
      self.clean_downloaded_metadata(metadata)
      result_queue.put(metadata)
      log.debug('Added Issue(%s) to queue' % metadata.title)

  def identify_results_keygen(self, title=None, authors=None,
                              identifiers=None):
    'Provide a keying function for result comparison'
    (issue_number, title_tokens) = utils.normalised_title(self, title)
    return partial(
      utils.keygen, title=title, authors=authors, identifiers=identifiers,
      issue_number=issue_number, title_tokens=title_tokens)

  def identify(self, log, result_queue, abort,
               title=None, authors=None, identifiers=None, timeout=30):
    '''Attempt to identify comicvine Issue matching given parameters'''

    # Do a simple lookup if comicvine identifier present
    if identifiers:
      comicvine_id = int(identifiers.get('comicvine'))
      if comicvine_id is not None:
        log.debug('Looking up Issue(%d)' % comicvine_id)
        self.enqueue(log, result_queue, comicvine_id)
        return None

    if title:
      # Look up candidate volumes based on title
      (issue_number, candidate_volumes) = utils.find_title(self, title, log)

      # Look up candidate authors
      candidate_authors = utils.find_authors(self, authors, log)

      # Look up candidate issues
      candidate_issues = utils.find_issues(
        candidate_volumes, issue_number, log)

      # Refine issue selection based on authors
      for issue in candidate_issues:
        if candidate_authors:
          for author in candidate_authors:
            if issue in author.issues:
              self.enqueue(log, result_queue, issue.id)
              break
        else:
          self.enqueue(log, result_queue, issue.id)

    return None

  def download_cover(self, log, result_queue, abort,
                     title=None, authors=None, identifiers=None,
                     timeout=30, get_best_cover=False):
    if identifiers and 'comicvine' in identifiers:
      for url in utils.cover_urls(identifiers['comicvine'], get_best_cover):
        browser = self.browser
        log('Downloading cover from:', url)
        try:
          cdata = browser.open_novisit(url, timeout=timeout).read()
          result_queue.put((self, cdata))
        except:
          log.exception('Failed to download cover from:', url)

<commit_msg>Bump version to 0.8.0<commit_after>'''
calibre_plugins.comicvine - A calibre metadata source for comicvine
'''
#pylint: disable-msg=R0913,R0904
from functools import partial
import logging
from Queue import Queue

from calibre import setup_cli_handlers
from calibre.ebooks.metadata.sources.base import Source
from calibre.utils.config import OptionParser
import calibre.utils.logging as calibre_logging
from calibre_plugins.comicvine.config import PREFS
from calibre_plugins.comicvine import utils

class Comicvine(Source):
  ''' Metadata source implementation '''
  name = 'Comicvine'
  description = 'Downloads metadata and covers from Comicvine'
  author = 'Russell Heilling'
  version = (0, 8, 0)
  #TODO(xchewtoyx): Implement cover capability
  capabilities = frozenset(['identify', 'cover'])
  touched_fields = frozenset([
      'title', 'authors', 'identifier:comicvine', 'comments', 'publisher',
      'pubdate', 'series'
      ])

  has_html_comments = True
  can_get_multiple_covers = True

  def __init__(self, *args, **kwargs):
    self.logger = logging.getLogger('urls')
    self.logger.setLevel(logging.DEBUG)
    self.logger.addHandler(utils.CalibreHandler(logging.DEBUG))
    Source.__init__(self, *args, **kwargs)

  def config_widget(self):
    from calibre_plugins.comicvine.config import ConfigWidget
    return ConfigWidget()

  def save_settings(self, config_widget):
    config_widget.save_settings()

  def is_configured(self):
    return bool(PREFS.get('api_key'))

  def cli_main(self, args):
    'Perform comicvine lookups from the calibre-debug cli'
    def option_parser():
      'Parse command line options'
      parser = OptionParser(
        usage='Comicvine [t:title] [a:authors] [i:id]')
      parser.add_option('--verbose', '-v', default=False,
                        action='store_true', dest='verbose')
      parser.add_option('--debug_api', default=False,
                        action='store_true', dest='debug_api')
      return parser

    opts, args = option_parser().parse_args(args)
    if opts.debug_api:
      calibre_logging.default_log = calibre_logging.Log(
        level=calibre_logging.DEBUG)
    if opts.verbose:
      level = 'DEBUG'
    else:
      level = 'INFO'
    setup_cli_handlers(logging.getLogger('comicvine'),
                       getattr(logging, level))
    log = calibre_logging.ThreadSafeLog(level=getattr(calibre_logging, level))

    (title, authors, ids) = (None, [], {})
    for arg in args:
      if arg.startswith('t:'):
        title = arg.split(':', 1)[1]
      if arg.startswith('a:'):
        authors.append(arg.split(':', 1)[1])
      if arg.startswith('i:'):
        (idtype, identifier) = arg.split(':', 2)[1:]
        ids[idtype] = int(identifier)
    result_queue = Queue()
    self.identify(
      log, result_queue, False, title=title, authors=authors, identifiers=ids)
    ranking = self.identify_results_keygen(title, authors, ids)
    for result in sorted(result_queue.queue, key=ranking):
      if result.pubdate:
        pubdate = str(result.pubdate.date())
      else:
        pubdate = 'Unknown'
      log.info('(%04d) - %s: %s [%s]' % (
          ranking(result), result.identifiers['comicvine'],
          result.title, pubdate))

  def enqueue(self, log, result_queue, issue_id):
    'Add a result entry to the result queue'
    log.debug('Adding Issue(%d) to queue' % issue_id)
    metadata = utils.build_meta(log, issue_id)
    if metadata:
      self.clean_downloaded_metadata(metadata)
      result_queue.put(metadata)
      log.debug('Added Issue(%s) to queue' % metadata.title)

  def identify_results_keygen(self, title=None, authors=None,
                              identifiers=None):
    'Provide a keying function for result comparison'
    (issue_number, title_tokens) = utils.normalised_title(self, title)
    return partial(
      utils.keygen, title=title, authors=authors, identifiers=identifiers,
      issue_number=issue_number, title_tokens=title_tokens)

  def identify(self, log, result_queue, abort,
               title=None, authors=None, identifiers=None, timeout=30):
    '''Attempt to identify comicvine Issue matching given parameters'''

    # Do a simple lookup if comicvine identifier present
    if identifiers:
      comicvine_id = int(identifiers.get('comicvine'))
      if comicvine_id is not None:
        log.debug('Looking up Issue(%d)' % comicvine_id)
        self.enqueue(log, result_queue, comicvine_id)
        return None

    if title:
      # Look up candidate volumes based on title
      (issue_number, candidate_volumes) = utils.find_title(self, title, log)

      # Look up candidate authors
      candidate_authors = utils.find_authors(self, authors, log)

      # Look up candidate issues
      candidate_issues = utils.find_issues(
        candidate_volumes, issue_number, log)

      # Refine issue selection based on authors
      for issue in candidate_issues:
        if candidate_authors:
          for author in candidate_authors:
            if issue in author.issues:
              self.enqueue(log, result_queue, issue.id)
              break
        else:
          self.enqueue(log, result_queue, issue.id)

    return None

  def download_cover(self, log, result_queue, abort,
                     title=None, authors=None, identifiers=None,
                     timeout=30, get_best_cover=False):
    if identifiers and 'comicvine' in identifiers:
      for url in utils.cover_urls(identifiers['comicvine'], get_best_cover):
        browser = self.browser
        log('Downloading cover from:', url)
        try:
          cdata = browser.open_novisit(url, timeout=timeout).read()
          result_queue.put((self, cdata))
        except:
          log.exception('Failed to download cover from:', url)

<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

from sys import exc_info

class SpamBLError(Exception):
    ''' Base exception class for spambl module '''

class UnknownCodeError(SpamBLError):
    ''' Raise when trying to use an unexpected value of dnsbl return code '''

class DNSBL(object):
    ''' Represents a DNSBL service provider '''

    def __init__(self, identifier, query_suffix, code_item_class, lists_ips, lists_uris):
        ''' Create new DNSBL object

        :param identifier: a value designating DNSBL service provider: its name or url address.
        :param query_suffix: a suffix added to DNSBL query address
        :param code_item_class: item classes associated with DNSBL query return codes
        :param lists_ips: information if this object represents an ip blocklist
        :param lists_uris: information if this object represents a domain name blocklist
        '''

        self.identifier = identifier
        self._query_suffix = query_suffix
        self._code_item_class = code_item_class
        self.lists_ips = lists_ips
        self.lists_uris = lists_uris

    def get_classification(self, code):
        ''' Return classification for given code

        :param code: a valid return code extracted from response to DNSBL query
        :raises UnknownCodeError: raised when given code is not specified in self._code_item_class
        :returns: a value associated with a valid return code
        '''

        try:
            return self._code_item_class[code]

        except KeyError:

            msg_template = 'Using a code value "{}" unsupported by DNSBL instance representing {}'
            raise UnknownCodeError(msg_template.format(code, self.identifier)), None, exc_info()[2]

if __name__ == '__main__':
    pass<commit_msg>Add _get_host_and_query_prefix method to DNSBL class<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

from sys import exc_info

class SpamBLError(Exception):
    ''' Base exception class for spambl module '''

class UnknownCodeError(SpamBLError):
    ''' Raise when trying to use an unexpected value of dnsbl return code '''

class DNSBL(object):
    ''' Represents a DNSBL service provider '''

    def __init__(self, identifier, query_suffix, code_item_class, lists_ips, lists_uris):
        ''' Create new DNSBL object

        :param identifier: a value designating DNSBL service provider: its name or url address.
        :param query_suffix: a suffix added to DNSBL query address
        :param code_item_class: item classes associated with DNSBL query return codes
        :param lists_ips: information if this object represents an ip blocklist
        :param lists_uris: information if this object represents a domain name blocklist
        '''

        self.identifier = identifier
        self._query_suffix = query_suffix
        self._code_item_class = code_item_class
        self.lists_ips = lists_ips
        self.lists_uris = lists_uris

    def get_classification(self, code):
        ''' Return classification for given code

        :param code: a valid return code extracted from response to DNSBL query
        :raises UnknownCodeError: raised when given code is not specified in self._code_item_class
        :returns: a value associated with a valid return code
        '''

        try:
            return self._code_item_class[code]

        except KeyError:

            msg_template = 'Using a code value "{}" unsupported by DNSBL instance representing {}'
            raise UnknownCodeError(msg_template.format(code, self.identifier)), None, exc_info()[2]

    def _get_host_and_query_prefix(self, host_collection):
        ''' Get valid hosts and  query prefixes from hosts given in host_collection

        :param host_collection: a container with valid host values
        :returns: a tuple with original host and a value prepared for lookup
        '''
        if self.lists_ips:
            for ip in host_collection.ips:
                ip = str(ip)
                suffix = '.in-addr.arpa' if ip.version == 4 else '.ip6.arpa'
                reverse = ip.replace(suffix, '')

                yield ip, reverse

        if self.lists_uris:
            for hostname in host_collection.hostnames:
                hostname = str(hostname)
                yield hostname, hostname.rstrip('.')

if __name__ == '__main__':
    pass<|endoftext|>


<commit_before>#!/usr/bin/env
# -*- coding: utf-8 -*-
# __author__ = 'QY-Y'
"""Spider.

Usage:
  spider.py [-u <URL>]
            [<Keyword> <Keyword>...]
            [-d <Depth>]
            [--testself]
            [-f <LogFile>]
            [-l <Level>]
            [--thread <ThreadNumber>]
            [--dbfile <DatabaseFileName>]
  spider.py (-h | --help)
  spider.py --version

Options:
  -h --help     Show this screen.
  --version     Show version.
  -d <Depth>    Depth of the spider.                                [default: 1].
  --testself    Test module                                         [default: False]
  -f <LogFile>  Name/Path of the log file.                          [default: spider.log]
  -l <Level>    How detailed the logging should be(1/less-5/more).  [default: 1]
  --thread <ThreadNumber>                                           [default: 10]
  --dbfile <DatabaseFileName>                                       [default: spider.db]

Description:
  -u <URL> Root URL.
  <Keyword> Keyword list, split by ' '.
"""

import os
import time
import sys
import Queue
import re
from docopt import docopt
import logging
import threading
import signal
import sqlite3
import urllib2
import lxml.html
import hashlib
import gzip
from StringIO import StringIO
from progressbar import *

IS_EXIT = False

# Description : chang global variable IS_EXIT.

def handler(signum, frame):
    global IS_EXIT
    IS_EXIT = True

# Description : send request to url,return contentof response.

# Output :
#     req.get_type()
#     req.get_host()
#     content

def request_url(url):
    req = urllib2.Request(url)
    req.add_header(
        'User-Agent',
        'Mozilla/5.0 (Windows NT 6.2; rv:16.0) Gecko/20100101 Firefox/48.0')
    req.add_header('Accept-Encoding', 'gzip')
    req.add_header('Accept-Language', 'en-US,en;q=0.5')
    try:
        res = urllib2.urlopen(url, timeout=10)
        content = res.read()
        if res.info().get('Content-Encoding') == 'gzip':
            buf = StringIO(content)
            content = gzip.GzipFile(fileobj=buf).read()
        return req.get_type(), req.get_host(), content
    except Exception as e:
        logging.warning(' open[' + url + '] failed ' + str(e))
        return req.get_type(), req.get_host(), False
# Description : init database file create table(if necessary)
# Input:
#     dbFile: Name of the Database file
# Output:
#     db: A sqlite3 connection obejct
#     c: A cursor from db

def init_data_base(database):

    exist = False
    files = os.listdir('.')
    if database in files:
        exist = True
    db = sqlite3.connect(database, check_same_thread=False)
    c = db.cursor()
    if not exist:
        try:
            c.execute(
                'create table spider(id integer primary key,url text,key text,content text)')
            db.commit()
        except sqlite3.OperationalError:
            logging.cratical(database + 'Error in creating table')
    return db, c

# Description : insert data to database

def insert_data(url, key_word, content, saved_count):
    content = urllib2.quote(str(content))
    try:
        c.execute('insert into spider(url,key,content) values("'
                  + url + '","' + key_word + '","' + content + '")')
        db.commit()
        saved_count.put(url)
    except sqlite3.OperationalError:
        logging.critical(' insert [' + url + '] error')

# Description : WORKER: search and save pages

class WORKER(threading.Thread):

    def __init__(self, links, keys, rlock, url_set, md5_set, saved_count):
        threading.Thread.__init__(self)
        self.task_queue = links
        self.key_list = keys
        self.rlock = rlock
        self.link = None
        self.depth = None
        self.setDaemon(True)
        self.start()
        self.url_set = url_set
        self.md5_set = md5_set
        self.count = 0
        self.saved_count = saved_count

#    Description : 1.get task from task queue.
#                   2.get sub links
#                   3.save pages
    def run(self):
        while True:
            try:
                self.link, self.depth = self.task_queue.get(timeout=2)
            except Exception as e:
                logging.info("get task error")
                break
            try:
                res_host, res_type, data = request_url(self.link)
                if not data:
                    self.task_queue.task_done()
                    continue
                self.count += 1
                md5 = hashlib.md5(data).hexdigest()
                if self.depth > 0:
                    self.depth -= 1
                    sub_link_list = self.getLinks(data, res_host, res_type)
                else:
                    sub_link_list = []
                for sub_link in sub_link_list:
                    if not (sub_link[0] in self.url_set):
                        self.url_set.add((sub_link[0]))
                        self.task_queue.put(sub_link)
                if not (md5 in self.md5_set):
                    self.save(data)
                    self.md5_set.add(md5)
            except Exception as e:
                logging.critical(str(e))
            self.task_queue.task_done()

#    Description : save pages if key word found
    def save(self, data):
        if not data:
            return
        for key in self.key_list:
            if data.find(key) > 0:
                logging.info(" [" + key + "] found in [" + self.link + "]")
                self.rlock.acquire()
                insert_data(self.link, key, data, self.saved_count)
                self.rlock.release()
                break
#    Description : get sub links,returns a list

    def getLinks(self, data, res_type, res_host):
        if not data:
            return []
        host = res_type + '://' + res_host
        try:
            data = data.decode('utf8', 'ignore')
            doc = lxml.html.document_fromstring(data)
        except Exception as e:
            logging.critical(str(e))
        tags = ['a', 'iframe', 'frame']
        doc.make_links_absolute(host)
        links = doc.iterlinks()
        new_link_list = []
        absolute_Link_list = []
        for l in links:
            if l[0].tag in tags:
                new_link_list.append(l)
        del links
        for link in new_link_list:
            if re.match('http', link[2]):
                absolute_Link_list.append((link[2], self.depth))
        return absolute_Link_list

# Description : manage the task queue and the thread pool,meanwhile show
# progress

class THREAD_POOL:

    def __init__(self, num, event, url, depth, key_word_list):
        self.num = num
        self.event = event
        self.threads = []
        self.task_queue = Queue.Queue()
        self.task_queue.put((url, depth))
        self.key_list = key_word_list
        self.url_set = set()
        self.url_set.add(url)
        self.md5_set = set()
        self.saved_count = Queue.Queue()
        for i in range(self.num):
            new_thread = WORKER(
                self.task_queue,
                self.key_list,
                rlock,
                self.url_set,
                self.md5_set,
                self.saved_count)
            self.threads.append(new_thread)
        logging.info(" pool init done, " + str(self.num) + " woreker created")
# Description : show progress every 0.5 second and check the global variable IS_EXIT
#           return if IS_EXIT is true or all tasks is done.

    def show_percent(self):
        global IS_EXIT
        while self.task_queue.unfinished_tasks:
            if IS_EXIT:
                try:
                    self.task_queue.empty()
                    self.event.set()
                except Exception as e:
                    logging.critical(str(e))
                return
            width = 50
            now = self.task_queue.qsize()
            total = len(self.url_set)
            percent = (float(now) / float(total))
            percent = int((1 - percent) * 100)
            info = ('[%%-%ds' % width) % (width * percent / 100 * '=')
            info += '] ' + str(percent) + "%    "
            info += str(len(self.url_set) - self.task_queue.qsize())
            info += '/' + str(len(self.url_set)) + '\r'
            sys.stdout.write((len(info) + 5) * ' ' + '\r')
            sys.stdout.flush()
            sys.stdout.write(info)
            sys.stdout.flush()

            time.sleep(0.5)

        sys.stdout.write('\n')
        print self.saved_count.qsize(), "pages saved. All tasks done at", time.ctime()
        self.event.set()
        return

class testSameDB(threading.Thread):

    def __init__(self, cursor, md5, progress):
        threading.Thread.__init__(self)
        self.c = cursor
        self.count = 0
        self.md5 = md5
        self.progress = progress
        self.start()

    def run(self):
        while True:
                # get 10,000 contents each time
            self.c.execute('select content from spider limit %s,%s'
                           % (self.count, self.count + 10000))
            self.count += 10000
            contents = self.c.fetchall()
            if len(contents) == 0:
                break
            for c in contents:
                res = hashlib.md5(c[0].encode('utf8'))
                self.md5.add(res.hexdigest())
                self.progress[0] = self.count

# checke if there are duplicate pages in database(by MD5)
def test(dbFile):
    progress = [0]
    md5 = set()
    # get count of all pages in database.
    db = sqlite3.connect(dbFile, check_same_thread=False)
    c = db.cursor()
    c.execute('select count(*) from spider')
    total_num = c.fetchall()[0][0]
    if total_num is 0:
        logging.warning("0 content in database")
        return
    t = testSameDB(c, md5, progress)
    t.join()
    pBar = ProgressBar(widgets=[Percentage(), Bar()],
                       maxval=total_num).start()
    while progress[0] < total_num:
        pBar.update(progress[0] + 1)
    pBar.finish()
    if len(md5) == total_num:
        print('no duplicate pages found in database')
    else:
        print('duplicate pages found in database')

def main_handler(thread_num, url, depth, key_word_list, is_test):
    event = threading.Event()
    event.clear()
    pool = THREAD_POOL(
        num=thread_num,
        event=event,
        url=url,
        depth=depth,
        key_word_list=key_word_list)
    pool.show_percent()
    if is_test:
        test(dbFile)

if __name__ == '__main__':
    arguments = docopt(__doc__, version='0.1')
    try:
        arguments['-l'] = int(arguments['-l'])
        arguments['-d'] = int(arguments['-d'])
        arguments['--thread'] = int(arguments['--thread'])
    except:
        print "-l,-d,--thread must be numbers"
        sys.exit()
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)
    rlock = threading.RLock()
    if (not 0 < arguments['-l'] < 6) or (arguments['-d'] <
                                         1) or (arguments['--thread'] < -100) or (not arguments['-u']):
        print 'input error'
        _usage()
        exit()
    db, c = init_data_base(arguments['--dbfile'])
    logLevel = {
        1: logging.CRITICAL,
        2: logging.ERROR,
        3: logging.WARNING,
        4: logging.INFO,
        5: logging.DEBUG,
    }
    logging.basicConfig(
        filename=arguments['-f'],
        level=arguments['-l'],
        format='%(asctime)s-%(levelname)s-%(filename)s-'
        '%(threadName)s-%(message)s',
        datefmt='[%d/%b/%Y %H:%M:%S]',
    )
    main_handler(
        arguments['--thread'],
        arguments['-u'],
        arguments['-d'] - 1,
        arguments['<Keyword>'],
        arguments['--testself'])<commit_msg>using docopt instead of getopt<commit_after>#!/usr/bin/env
# -*- coding: utf-8 -*-
# __author__ = 'QY-Y'
"""Spider.

Usage:
  spider.py [-u <URL>]
            [<Keyword> <Keyword>...]
            [-d <Depth>]
            [--testself]
            [-f <LogFile>]
            [-l <Level>]
            [--thread <ThreadNumber>]
            [--dbfile <DatabaseFileName>]
  spider.py (-h | --help)
  spider.py --version

Options:
  -h --help     Show this screen.
  --version     Show version.
  -d <Depth>    Depth of the spider.                                [default: 1].
  --testself    Test module                                         [default: False]
  -f <LogFile>  Name/Path of the log file.                          [default: spider.log]
  -l <Level>    How detailed the logging should be(1/less-5/more).  [default: 1]
  --thread <ThreadNumber>                                           [default: 10]
  --dbfile <DatabaseFileName>                                       [default: spider.db]

Description:
  -u <URL> Root URL.
  <Keyword> Keyword list, split by ' '.
"""

import os
import time
import sys
import Queue
import re
from docopt import docopt
import logging
import threading
import signal
import sqlite3
import urllib2
import lxml.html
import hashlib
import gzip
from StringIO import StringIO
from progressbar import *

IS_EXIT = False

def handler(signum, frame):
    # Description : chang global variable IS_EXIT.
    global IS_EXIT
    IS_EXIT = True

def request_url(url):
# Description : send request to url,return contentof response.
# Output :
#     req.get_type()
#     req.get_host()
#     content
    req = urllib2.Request(url)
    req.add_header(
        'User-Agent',
        'Mozilla/5.0 (Windows NT 6.2; rv:16.0) Gecko/20100101 Firefox/48.0')
    req.add_header('Accept-Encoding', 'gzip')
    req.add_header('Accept-Language', 'en-US,en;q=0.5')
    try:
        res = urllib2.urlopen(url, timeout=10)
        content = res.read()
        if res.info().get('Content-Encoding') == 'gzip':
            buf = StringIO(content)
            content = gzip.GzipFile(fileobj=buf).read()
        return req.get_type(), req.get_host(), content
    except Exception as e:
        logging.warning(' open[' + url + '] failed ' + str(e))
        return req.get_type(), req.get_host(), False

def init_data_base(database):
# Description : init database file create table(if necessary)
# Input:
#     dbFile: Name of the Database file
# Output:
#     db: A sqlite3 connection obejct
#     c: A cursor from db
    exist = False
    files = os.listdir('.')
    if database in files:
        exist = True
    db = sqlite3.connect(database, check_same_thread=False)
    c = db.cursor()
    if not exist:
        try:
            c.execute(
                'create table spider(id integer primary key,url text,key text,content text)')
            db.commit()
        except sqlite3.OperationalError:
            logging.cratical(database + 'Error in creating table')
    return db, c

def insert_data(url, key_word, content, saved_count):
# Description : insert data to database
    content = urllib2.quote(str(content))
    try:
        c.execute('insert into spider(url,key,content) values("'
                  + url + '","' + key_word + '","' + content + '")')
        db.commit()
        saved_count.put(url)
    except sqlite3.OperationalError:
        logging.critical(' insert [' + url + '] error')

def argv_check(argvs):
# Description : check user input, exit if illegal
    if (not 0 < argvs['-l'] < 6):
        print '-l must be in [1,2,3,4,5]'
        sys.exit()
    elif argvs['-d'] < 1:
        print '-d must larger than 0'
        sys.exit()
    elif argvs['--thread'] < 1:
        print '--thread must larger than 0'
        sys.exit()
    elif not argvs['-u']:
        print 'Please input root url as -u'
        sys.exit()

# Description : WORKER: search and save pages
class WORKER(threading.Thread):

    def __init__(self, links, keys, rlock, url_set, md5_set, saved_count):
        threading.Thread.__init__(self)
        self.task_queue=links
        self.key_list=keys
        self.rlock=rlock
        self.link=None
        self.depth=None
        self.setDaemon(True)
        self.start()
        self.url_set=url_set
        self.md5_set=md5_set
        self.count=0
        self.saved_count=saved_count

#    Description : 1.get task from task queue.
#                   2.get sub links
#                   3.save pages
    def run(self):
        while True:
            try:
                self.link, self.depth=self.task_queue.get(timeout=2)
            except Exception as e:
                logging.info("get task error")
                break
            try:
                res_host, res_type, data=request_url(self.link)
                if not data:
                    self.task_queue.task_done()
                    continue
                self.count += 1
                md5=hashlib.md5(data).hexdigest()
                if self.depth > 0:
                    self.depth -= 1
                    sub_link_list=self.getLinks(data, res_host, res_type)
                else:
                    sub_link_list=[]
                for sub_link in sub_link_list:
                    if not (sub_link[0] in self.url_set):
                        self.url_set.add((sub_link[0]))
                        self.task_queue.put(sub_link)
                if not (md5 in self.md5_set):
                    self.save(data)
                    self.md5_set.add(md5)
            except Exception as e:
                logging.critical(str(e))
            self.task_queue.task_done()

#    Description : save pages if key word found
    def save(self, data):
        if not data:
            return
        for key in self.key_list:
            if data.find(key) > 0:
                logging.info(" [" + key + "] found in [" + self.link + "]")
                self.rlock.acquire()
                insert_data(self.link, key, data, self.saved_count)
                self.rlock.release()
                break
#    Description : get sub links,returns a list

    def getLinks(self, data, res_type, res_host):
        if not data:
            return []
        host=res_type + '://' + res_host
        try:
            data=data.decode('utf8', 'ignore')
            doc=lxml.html.document_fromstring(data)
        except Exception as e:
            logging.critical(str(e))
        tags=['a', 'iframe', 'frame']
        doc.make_links_absolute(host)
        links=doc.iterlinks()
        new_link_list=[]
        absolute_Link_list=[]
        for l in links:
            if l[0].tag in tags:
                new_link_list.append(l)
        del links
        for link in new_link_list:
            if re.match('http', link[2]):
                absolute_Link_list.append((link[2], self.depth))
        return absolute_Link_list

# Description : manage the task queue and the thread pool,meanwhile show
# progress

class THREAD_POOL:

    def __init__(self, num, event, url, depth, key_word_list):
        self.num=num
        self.event=event
        self.threads=[]
        self.task_queue=Queue.Queue()
        self.task_queue.put((url, depth))
        self.key_list=key_word_list
        self.url_set=set()
        self.url_set.add(url)
        self.md5_set=set()
        self.saved_count=Queue.Queue()
        for i in range(self.num):
            new_thread=WORKER(
                self.task_queue,
                self.key_list,
                rlock,
                self.url_set,
                self.md5_set,
                self.saved_count)
            self.threads.append(new_thread)
        logging.info(" pool init done, " + str(self.num) + " woreker created")
# Description : show progress every 0.5 second and check the global variable IS_EXIT
#           return if IS_EXIT is true or all tasks is done.

    def show_percent(self):
        global IS_EXIT
        while self.task_queue.unfinished_tasks:
            if IS_EXIT:
                try:
                    self.task_queue.empty()
                    self.event.set()
                except Exception as e:
                    logging.critical(str(e))
                return
            width=50
            now=self.task_queue.qsize()
            total=len(self.url_set)
            percent=(float(now) / float(total))
            percent=int((1 - percent) * 100)
            info=('[%%-%ds' % width) % (width * percent / 100 * '=')
            info += '] ' + str(percent) + "%    "
            info += str(len(self.url_set) - self.task_queue.qsize())
            info += '/' + str(len(self.url_set)) + '\r'
            sys.stdout.write((len(info) + 5) * ' ' + '\r')
            sys.stdout.flush()
            sys.stdout.write(info)
            sys.stdout.flush()

            time.sleep(0.5)

        sys.stdout.write('\n')
        print self.saved_count.qsize(), "pages saved. All tasks done at", time.ctime()
        self.event.set()
        return

class testSameDB(threading.Thread):

    def __init__(self, cursor, md5, progress):
        threading.Thread.__init__(self)
        self.c=cursor
        self.count=0
        self.md5=md5
        self.progress=progress
        self.start()

    def run(self):
        while True:
                # get 10,000 contents each time
            self.c.execute('select content from spider limit %s,%s'
                           % (self.count, self.count + 10000))
            self.count += 10000
            contents=self.c.fetchall()
            if len(contents) == 0:
                break
            for c in contents:
                res=hashlib.md5(c[0].encode('utf8'))
                self.md5.add(res.hexdigest())
                self.progress[0]=self.count

# checke if there are duplicate pages in database(by MD5)
def test(dbFile):
    progress=[0]
    md5=set()
    # get count of all pages in database.
    db=sqlite3.connect(dbFile, check_same_thread=False)
    c=db.cursor()
    c.execute('select count(*) from spider')
    total_num=c.fetchall()[0][0]
    if total_num is 0:
        logging.warning("0 content in database")
        return
    t=testSameDB(c, md5, progress)
    t.join()
    pBar=ProgressBar(widgets=[Percentage(), Bar()],
                       maxval=total_num).start()
    while progress[0] < total_num:
        pBar.update(progress[0] + 1)
    pBar.finish()
    if len(md5) == total_num:
        print('no duplicate pages found in database')
    else:
        print('duplicate pages found in database')

def main_handler(thread_num, url, depth, key_word_list, is_test):
    event=threading.Event()
    event.clear()
    pool=THREAD_POOL(
        num=thread_num,
        event=event,
        url=url,
        depth=depth,
        key_word_list=key_word_list)
    pool.show_percent()
    if is_test:
        test(dbFile)

if __name__ == '__main__':
    arguments=docopt(__doc__, version='0.1')
    try:
        arguments['-l']=int(arguments['-l'])
        arguments['-d']=int(arguments['-d'])
        arguments['--thread']=int(arguments['--thread'])

    except:
        print "-l,-d,--thread must be numbers"
        sys.exit()
    argv_check(arguments)
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)
    rlock=threading.RLock()

    db, c=init_data_base(arguments['--dbfile'])
    logLevel={
        1: logging.CRITICAL,
        2: logging.ERROR,
        3: logging.WARNING,
        4: logging.INFO,
        5: logging.DEBUG,
    }
    logging.basicConfig(
        filename=arguments['-f'],
        level=arguments['-l'],
        format='%(asctime)s-%(levelname)s-%(filename)s-'
        '%(threadName)s-%(message)s',
        datefmt='[%d/%b/%Y %H:%M:%S]',
    )
    main_handler(
        arguments['--thread'],
        arguments['-u'],
        arguments['-d'] - 1,
        arguments['<Keyword>'],
        arguments['--testself'])<|endoftext|>


<commit_before>#!/usr/bin/env python

from bs4 import BeautifulSoup
import spotipy
import eyed3
import requests
import pafy
import shutil
import os
import sys
import argparse
#import spotipy.util as util

os.chdir(sys.path[0])

if not os.path.exists("Music"):
	os.makedirs("Music")
open('list.txt', 'a').close()

spotify = spotipy.Spotify()

parser = argparse.ArgumentParser()
parser.add_argument("-n", "--no-convert", help="skip the conversion process and meta-tags",
                        action="store_true")
args = parser.parse_args()
if args.no_convert:
        print("-n, --no-convert  skip the conversion process and meta-tags")

def searchYT(number):
	items = (requests.request(method='GET', url=URL)).text
	items_parse = BeautifulSoup(items, "html.parser")
	first_result = items_parse.find_all(attrs={'class':'yt-uix-tile-link'})[0]['href']
	check = 1
	while not first_result.find('channel') == -1 or not first_result.find('googleads') == -1:
		first_result = items_parse.find_all(attrs={'class':'yt-uix-tile-link'})[check]['href']
		check += 1
	del check
	full_link = "youtube.com" + first_result
	#print(full_link)
	global video
	video = pafy.new(full_link)
	global raw_title
	raw_title = (video.title).encode("utf-8")
	global title
	title = ((video.title).replace("\\", "_").replace("/", "_").replace(":", "_").replace("*", "_").replace("?", "_").replace('"', "_").replace("<", "_").replace(">", "_").replace("|", "_").replace(" ", "_")).encode('utf-8')
	if not number == None:
		print(str(number) + '. ' + (video.title).encode("utf-8"))
	else:
		print(video.title).encode("utf-8")

def checkExists(islist):
	if os.path.exists("Music/" + title + ".m4a.temp"):
		os.remove("Music/" + title + ".m4a.temp")
	global extension
	if args.no_convert:
		extension = '.m4a'
	else:
		if os.path.exists("Music/" + title + ".m4a"):
			os.remove("Music/" + title + ".m4a")
		extension = '.mp3'
	if os.path.isfile("Music/" + title + extension):
		if extension == '.mp3':
			audiofile = eyed3.load("Music/" + title + extension)
			if isSpotify() and not audiofile.tag.title == content['name']:
				os.remove("Music/" + title + extension)
				return False
		if islist:
			trimSong()
			return True
		else:
			prompt = raw_input('Song with same name has already been downloaded. Re-download? (y/n/play): ')
			if prompt == "y":
				os.remove("Music/" + title + extension)
				return False
			elif prompt =="play":
				if not os.name == 'nt':
					os.system('mplayer "' + 'Music/' + title + extension + '"')
				else:
					print('Playing ' + title + extension)
					os.system('start ' + 'Music/' + title + extension)
				return True
			else:
				return True

def getLyrics():
	if not title == '':
		if song == '':
			link = 'https://duckduckgo.com/html/?q=' + raw_title.replace(' ', '+') + '+musixmatch'
		else:
			link = 'https://duckduckgo.com/html/?q=' + (content['artists'][0]['name'] + ' - ' + content['name']).replace(' ', '+') + '+musixmatch'
		page = requests.request(method='GET', url=link).text
		soup = BeautifulSoup(page, 'html.parser')
		link = soup.find('a', {'class':'result__url'})['href']
		page = requests.request(method='GET', url=link).text
		soup = BeautifulSoup(page, 'html.parser')
		for x in soup.find_all('p', {'class':'mxm-lyrics__content'}):
			print(x.get_text()).encode('utf-8')
	else:
		print('No log to read from..')

def fixSong():
	print 'Fixing meta-tags'
        audiofile = eyed3.load("Music/" + title + '.mp3')
        audiofile.tag.artist = content['artists'][0]['name']
        audiofile.tag.album = content['album']['name']
        audiofile.tag.title = content['name']
        albumart = (requests.get(content['album']['images'][0]['url'], stream=True)).raw
        with open('last_albumart.jpg', 'wb') as out_file:
                shutil.copyfileobj(albumart, out_file)
        albumart = open("last_albumart.jpg", "rb").read()
        audiofile.tag.images.set(3,albumart,"image/jpeg")
        audiofile.tag.save(version=(2,3,0))

def playSong():
	if not title == '':
		if not os.name == 'nt':
			os.system('mplayer "' + 'Music/' + title + extension + '"')
		else:
			print('Playing ' + title + '.mp3')
			os.system('start ' + 'Music/' + title + extension)

def convertSong():
	print('Converting ' + title + '.m4a to mp3')
	if not os.name == 'nt':
		os.system('avconv -loglevel 0 -i "' + 'Music/' + title + '.m4a" -ab 192k "' + 'Music/' + title + '.mp3"')
	else:
		os.system('Scripts\\avconv.exe -loglevel 0 -i "' + 'Music/' + title + '.m4a" -ab 192k "' + 'Music/' + title + '.mp3"')
	os.remove('Music/' + title + '.m4a')

def downloadSong():
	a = video.getbestaudio(preftype="m4a")
	a.download(filepath="Music/" + title + ".m4a")

def isSpotify():
	if (len(raw_song) == 22 and raw_song.replace(" ", "%20") == raw_song) or (raw_song.find('spotify') > -1):
		return True
	else:
		return False

def trimSong():
	with open('list.txt', 'r') as fin:
		data = fin.read().splitlines(True)
	with open('list.txt', 'w') as fout:
		fout.writelines(data[1:])

def trackPredict():
	global URL
	if isSpotify():
		global content
		content = spotify.track(raw_song)
		song = (content['artists'][0]['name'] + ' - ' + content['name']).replace(" ", "%20").encode('utf-8')
		URL = "https://www.youtube.com/results?sp=EgIQAQ%253D%253D&q=" + song
	else:
		song = raw_song.replace(" ", "%20")
		URL = "https://www.youtube.com/results?sp=EgIQAQ%253D%253D&q=" + song
		song = ''

def graceQuit():
		print('')
		print('')
		print('Exitting..')
		exit()

title = ''
song = ''

while True:
	x = 0
	y = 0
	try:
		for m in os.listdir('Music/'):
			if m.endswith('.m4a.temp'):
				os.remove('Music/' + m)
		print('')
		print('')
		raw_song = raw_input('>> Enter a song/cmd: ').decode('utf-8').encode('utf-8')
		print('')
		if raw_song == 'exit':
			exit()
		elif raw_song == 'play':
			playSong()

		elif raw_song == 'lyrics':
			getLyrics()

		elif raw_song == 'list':
			f = open('list.txt', 'r').read()
			lines = f.splitlines()
			for raw_song in lines:
				if not len(raw_song) == 0:
					x = x + 1
			print('Total songs in list = ' + str(x) + ' songs')
			for raw_song in lines:
				try:
					if not len(raw_song) == 0:
						trackPredict()
						print('')
						y = y + 1
						searchYT(number=y)
						if not checkExists(islist=True):
							downloadSong()
							print('')
							if not args.no_convert:
								convertSong()
								if isSpotify():
									fixSong()
							trimSong()
					else:
						trimSong()
				except KeyboardInterrupt:
					graceQuit()
				except:
					lines.append(raw_song)
					trimSong()
					with open('list.txt', 'a') as myfile:
						myfile.write(raw_song)
					print('Failed to download song. Will retry after other songs.')
		else:
			try:
				trackPredict()
				searchYT(number=None)
				if not checkExists(islist=False):
					downloadSong()
					print('')
					if not args.no_convert:
						convertSong()
						if isSpotify():
							fixSong()
			except KeyboardInterrupt:
				graceQuit()
	except KeyboardInterrupt:
		graceQuit()
<commit_msg>fix indentation and a parantheses error (#30)<commit_after>#!/usr/bin/env python

from bs4 import BeautifulSoup
import spotipy
import eyed3
import requests
import pafy
import shutil
import os
import sys
import argparse
#import spotipy.util as util

os.chdir(sys.path[0])

if not os.path.exists("Music"):
	os.makedirs("Music")
open('list.txt', 'a').close()

spotify = spotipy.Spotify()

parser = argparse.ArgumentParser()
parser.add_argument("-n", "--no-convert", help="skip the conversion process and meta-tags",
                        action="store_true")
args = parser.parse_args()
if args.no_convert:
        print("-n, --no-convert  skip the conversion process and meta-tags")

def searchYT(number):
	items = (requests.request(method='GET', url=URL)).text
	items_parse = BeautifulSoup(items, "html.parser")
	first_result = items_parse.find_all(attrs={'class':'yt-uix-tile-link'})[0]['href']
	check = 1
	while not first_result.find('channel') == -1 or not first_result.find('googleads') == -1:
		first_result = items_parse.find_all(attrs={'class':'yt-uix-tile-link'})[check]['href']
		check += 1
	del check
	full_link = "youtube.com" + first_result
	#print(full_link)
	global video
	video = pafy.new(full_link)
	global raw_title
	raw_title = (video.title).encode("utf-8")
	global title
	title = ((video.title).replace("\\", "_").replace("/", "_").replace(":", "_").replace("*", "_").replace("?", "_").replace('"', "_").replace("<", "_").replace(">", "_").replace("|", "_").replace(" ", "_")).encode('utf-8')
	if not number == None:
		print(str(number) + '. ' + (video.title).encode("utf-8"))
	else:
		print(video.title).encode("utf-8")

def checkExists(islist):
	if os.path.exists("Music/" + title + ".m4a.temp"):
		os.remove("Music/" + title + ".m4a.temp")
	global extension
	if args.no_convert:
		extension = '.m4a'
	else:
		if os.path.exists("Music/" + title + ".m4a"):
			os.remove("Music/" + title + ".m4a")
		extension = '.mp3'
	if os.path.isfile("Music/" + title + extension):
		if extension == '.mp3':
			audiofile = eyed3.load("Music/" + title + extension)
			if isSpotify() and not audiofile.tag.title == content['name']:
				os.remove("Music/" + title + extension)
				return False
		if islist:
			trimSong()
			return True
		else:
			prompt = raw_input('Song with same name has already been downloaded. Re-download? (y/n/play): ')
			if prompt == "y":
				os.remove("Music/" + title + extension)
				return False
			elif prompt =="play":
				if not os.name == 'nt':
					os.system('mplayer "' + 'Music/' + title + extension + '"')
				else:
					print('Playing ' + title + extension)
					os.system('start ' + 'Music/' + title + extension)
				return True
			else:
				return True

def getLyrics():
	if not title == '':
		if song == '':
			link = 'https://duckduckgo.com/html/?q=' + raw_title.replace(' ', '+') + '+musixmatch'
		else:
			link = 'https://duckduckgo.com/html/?q=' + (content['artists'][0]['name'] + ' - ' + content['name']).replace(' ', '+') + '+musixmatch'
		page = requests.request(method='GET', url=link).text
		soup = BeautifulSoup(page, 'html.parser')
		link = soup.find('a', {'class':'result__url'})['href']
		page = requests.request(method='GET', url=link).text
		soup = BeautifulSoup(page, 'html.parser')
		for x in soup.find_all('p', {'class':'mxm-lyrics__content'}):
			print(x.get_text()).encode('utf-8')
	else:
		print('No log to read from..')

def fixSong():
	print ('Fixing meta-tags')
	audiofile = eyed3.load("Music/" + title + '.mp3')
	audiofile.tag.artist = content['artists'][0]['name']
	audiofile.tag.album = content['album']['name']
	audiofile.tag.title = content['name']
	albumart = (requests.get(content['album']['images'][0]['url'], stream=True)).raw
	with open('last_albumart.jpg', 'wb') as out_file:
		shutil.copyfileobj(albumart, out_file)
	albumart = open("last_albumart.jpg", "rb").read()
	audiofile.tag.images.set(3,albumart,"image/jpeg")
	audiofile.tag.save(version=(2,3,0))

def playSong():
	if not title == '':
		if not os.name == 'nt':
			os.system('mplayer "' + 'Music/' + title + extension + '"')
		else:
			print('Playing ' + title + '.mp3')
			os.system('start ' + 'Music/' + title + extension)

def convertSong():
	print('Converting ' + title + '.m4a to mp3')
	if not os.name == 'nt':
		os.system('avconv -loglevel 0 -i "' + 'Music/' + title + '.m4a" -ab 192k "' + 'Music/' + title + '.mp3"')
	else:
		os.system('Scripts\\avconv.exe -loglevel 0 -i "' + 'Music/' + title + '.m4a" -ab 192k "' + 'Music/' + title + '.mp3"')
	os.remove('Music/' + title + '.m4a')

def downloadSong():
	a = video.getbestaudio(preftype="m4a")
	a.download(filepath="Music/" + title + ".m4a")

def isSpotify():
	if (len(raw_song) == 22 and raw_song.replace(" ", "%20") == raw_song) or (raw_song.find('spotify') > -1):
		return True
	else:
		return False

def trimSong():
	with open('list.txt', 'r') as fin:
		data = fin.read().splitlines(True)
	with open('list.txt', 'w') as fout:
		fout.writelines(data[1:])

def trackPredict():
	global URL
	if isSpotify():
		global content
		content = spotify.track(raw_song)
		song = (content['artists'][0]['name'] + ' - ' + content['name']).replace(" ", "%20").encode('utf-8')
		URL = "https://www.youtube.com/results?sp=EgIQAQ%253D%253D&q=" + song
	else:
		song = raw_song.replace(" ", "%20")
		URL = "https://www.youtube.com/results?sp=EgIQAQ%253D%253D&q=" + song
		song = ''

def graceQuit():
		print('')
		print('')
		print('Exitting..')
		exit()

title = ''
song = ''

while True:
	x = 0
	y = 0
	try:
		for m in os.listdir('Music/'):
			if m.endswith('.m4a.temp'):
				os.remove('Music/' + m)
		print('')
		print('')
		raw_song = raw_input('>> Enter a song/cmd: ').decode('utf-8').encode('utf-8')
		print('')
		if raw_song == 'exit':
			exit()
		elif raw_song == 'play':
			playSong()

		elif raw_song == 'lyrics':
			getLyrics()

		elif raw_song == 'list':
			f = open('list.txt', 'r').read()
			lines = f.splitlines()
			for raw_song in lines:
				if not len(raw_song) == 0:
					x = x + 1
			print('Total songs in list = ' + str(x) + ' songs')
			for raw_song in lines:
				try:
					if not len(raw_song) == 0:
						trackPredict()
						print('')
						y = y + 1
						searchYT(number=y)
						if not checkExists(islist=True):
							downloadSong()
							print('')
							if not args.no_convert:
								convertSong()
								if isSpotify():
									fixSong()
							trimSong()
					else:
						trimSong()
				except KeyboardInterrupt:
					graceQuit()
				except:
					lines.append(raw_song)
					trimSong()
					with open('list.txt', 'a') as myfile:
						myfile.write(raw_song)
					print('Failed to download song. Will retry after other songs.')
		else:
			try:
				trackPredict()
				searchYT(number=None)
				if not checkExists(islist=False):
					downloadSong()
					print('')
					if not args.no_convert:
						convertSong()
						if isSpotify():
							fixSong()
			except KeyboardInterrupt:
				graceQuit()
	except KeyboardInterrupt:
		graceQuit()
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
$Id$

Copyright (c) 2006-2011 sqlmap developers (http://sqlmap.sourceforge.net/)
See the file 'doc/COPYING' for copying permission
"""

import codecs
import os
import sys
import time
import traceback
import warnings

warnings.filterwarnings(action="ignore", message=".*was already imported", category=UserWarning)
warnings.filterwarnings(action="ignore", category=DeprecationWarning)

try:
    import psyco
    psyco.full()
    psyco.profile()
except ImportError:
    pass

from lib.controller.controller import start
from lib.core.common import banner
from lib.core.common import dataToStdout
from lib.core.common import getUnicode
from lib.core.common import setPaths
from lib.core.common import weAreFrozen
from lib.core.data import cmdLineOptions
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.data import paths
from lib.core.common import unhandledExceptionMessage
from lib.core.exception import exceptionsTuple
from lib.core.exception import sqlmapSilentQuitException
from lib.core.exception import sqlmapUserQuitException
from lib.core.option import init
from lib.core.profiling import profile
from lib.core.testing import smokeTest
from lib.core.testing import liveTest
from lib.core.xmldump import closeDumper
from lib.parse.cmdline import cmdLineParser

def modulePath():
    """
    This will get us the program's directory, even if we are frozen
    using py2exe
    """

    if weAreFrozen():
        return os.path.dirname(getUnicode(sys.executable, sys.getfilesystemencoding()))
    else:
        return os.path.dirname(os.path.realpath(__file__))

def main():
    """
    Main function of sqlmap when running from command line.
    """
    paths.SQLMAP_ROOT_PATH = modulePath()
    setPaths()
    banner()

    # Store original command line options for possible later restoration
    cmdLineOptions.update(cmdLineParser().__dict__)

    dataToStdout("[*] starting at: %s\n\n" % time.strftime("%X"), forceOutput=True)

    try:
        init(cmdLineOptions)
        if conf.profile:
            profile()
        elif conf.smokeTest:
            smokeTest()
        elif conf.liveTest:
            liveTest()
        else:
            start()

    except sqlmapUserQuitException:
        errMsg = "user quit"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except sqlmapSilentQuitException:
        closeDumper(False)

    except exceptionsTuple, e:
        e = getUnicode(e)
        logger.critical(e)
        closeDumper(False, e)

    except KeyboardInterrupt, _:
        print
        errMsg = "user aborted"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except EOFError, _:
        print
        errMsg = "exit"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except:
        print
        errMsg = unhandledExceptionMessage()
        logger.critical(errMsg)
        traceback.print_exc()
        closeDumper(False, errMsg)

    else:
        closeDumper(True)

    finally:
        dataToStdout("\n[*] shutting down at: %s\n\n" % time.strftime("%X"), forceOutput=True)

        kb.threadContinue = False
        kb.threadException = True

if __name__ == "__main__":
    main()
<commit_msg>cosmetics<commit_after>#!/usr/bin/env python

"""
$Id$

Copyright (c) 2006-2011 sqlmap developers (http://sqlmap.sourceforge.net/)
See the file 'doc/COPYING' for copying permission
"""

import codecs
import os
import sys
import time
import traceback
import warnings

warnings.filterwarnings(action="ignore", message=".*was already imported", category=UserWarning)
warnings.filterwarnings(action="ignore", category=DeprecationWarning)

try:
    import psyco
    psyco.full()
    psyco.profile()
except ImportError:
    pass

from lib.controller.controller import start
from lib.core.common import banner
from lib.core.common import dataToStdout
from lib.core.common import getUnicode
from lib.core.common import setPaths
from lib.core.common import weAreFrozen
from lib.core.data import cmdLineOptions
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.data import paths
from lib.core.common import unhandledExceptionMessage
from lib.core.exception import exceptionsTuple
from lib.core.exception import sqlmapSilentQuitException
from lib.core.exception import sqlmapUserQuitException
from lib.core.option import init
from lib.core.profiling import profile
from lib.core.testing import smokeTest
from lib.core.testing import liveTest
from lib.core.xmldump import closeDumper
from lib.parse.cmdline import cmdLineParser

def modulePath():
    """
    This will get us the program's directory, even if we are frozen
    using py2exe
    """

    if weAreFrozen():
        return os.path.dirname(getUnicode(sys.executable, sys.getfilesystemencoding()))
    else:
        return os.path.dirname(os.path.realpath(__file__))

def main():
    """
    Main function of sqlmap when running from command line.
    """

    paths.SQLMAP_ROOT_PATH = modulePath()
    setPaths()
    banner()

    # Store original command line options for possible later restoration
    cmdLineOptions.update(cmdLineParser().__dict__)

    dataToStdout("[*] starting at: %s\n\n" % time.strftime("%X"), forceOutput=True)

    try:
        init(cmdLineOptions)
        if conf.profile:
            profile()
        elif conf.smokeTest:
            smokeTest()
        elif conf.liveTest:
            liveTest()
        else:
            start()

    except sqlmapUserQuitException:
        errMsg = "user quit"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except sqlmapSilentQuitException:
        closeDumper(False)

    except exceptionsTuple, e:
        e = getUnicode(e)
        logger.critical(e)
        closeDumper(False, e)

    except KeyboardInterrupt, _:
        print
        errMsg = "user aborted"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except EOFError, _:
        print
        errMsg = "exit"
        logger.error(errMsg)
        closeDumper(False, errMsg)

    except:
        print
        errMsg = unhandledExceptionMessage()
        logger.critical(errMsg)
        traceback.print_exc()
        closeDumper(False, errMsg)

    else:
        closeDumper(True)

    finally:
        dataToStdout("\n[*] shutting down at: %s\n\n" % time.strftime("%X"), forceOutput=True)

        kb.threadContinue = False
        kb.threadException = True

if __name__ == "__main__":
    main()
<|endoftext|>