<commit_before>"""Sketch analyzer plugin for the Safe Browsing API."""
from __future__ import unicode_literals

import fnmatch
import logging
import re

import pkg_resources
import requests
from flask import current_app

from timesketch.lib.analyzers import interface
from timesketch.lib.analyzers import manager

class SafeBrowsingSketchPlugin(interface.BaseSketchAnalyzer):
    """Sketch analyzer for Safe Browsing."""

    NAME = 'safebrowsing'

    # Safe Browsing API v4, threatMatches.find endpoint.
    _SAFE_BROWSING_THREATMATCHING_ENDPOINT = (
        'https://safebrowsing.googleapis.com/v4/threatMatches:find')

    # Maximal number of URLs in a single Lookup request as per the
    # Safe Browsing API documentation.
    _SAFE_BROWSING_BULK_LIMIT = 500

    # An optional file containing URL wildcards to be allow listed
    # in a YAML file.
    _URL_ALLOW_LIST_CONFIG = 'safebrowsing_allowlist.yaml'

    # The keys to be added to the TS event from the ThreatMatch object
    # we get from Safe Browsing API.
    _SAFEBROWSING_ENTRY_KEEP = frozenset([
        'platformType',
        'threatType',
    ])

    # Used to find proper URLs in the 'url' entries of TS events.
    _URL_BEGINNING_RE = re.compile(r'(http(s|):\/\/\S*)')

    def __init__(self, index_name, sketch_id):
        """Initialize The Sketch Analyzer.

        Args:
            index_name: Elasticsearch index name
            sketch_id: Sketch ID
        """
        super(SafeBrowsingSketchPlugin, self).__init__(index_name, sketch_id)

        self._safebrowsing_api_key = current_app.config.get(
            'SAFEBROWSING_API_KEY')

        self._google_client_id = current_app.config.get(
            'SAFEBROWSING_CLIENT_ID',
            'Timesketch',
        )

        self._google_client_version = current_app.config.get(
            'SAFEBROWSING_CLIENT_VERSION',
            pkg_resources.get_distribution('timesketch').version,
        )

    def _is_url_allowlisted(self, url, allowlist):
        """Does a fnmatch against the allowlist.

        Args:
            url: The url
            allowlist: The allowlist, list-like
        Returns:
            Boolean with the result
        """

        for url_pattern in allowlist:
            if fnmatch.fnmatchcase(url, url_pattern):
                return True

        return False

    def _do_safebrowsing_lookup(self, urls, platforms, types):
        """URL lookup against the Safe Browsing API.

        Args:
            urls: URLs
            platforms: platformTypes field of threatInfo
            types: threatTypes field of threatInfo
        Returns:
            Dict of URLs with the hits
        """
        results = {}

        api_client = {
            'clientId': self._google_client_id,
            'clientVersion': self._google_client_version,
        }

        for index in range(0, len(urls), self._SAFE_BROWSING_BULK_LIMIT):
            body = {
                'client': api_client,
                'threatInfo': {
                    'platformTypes': platforms,
                    'threatTypes': types,
                    'threatEntryTypes': ['URL'],
                    'threatEntries': [
                        {'url': url}
                        for url in urls[
                            index:index + self._SAFE_BROWSING_BULK_LIMIT]
                    ],
                },
            }

            response = requests.post(
                self._SAFE_BROWSING_THREATMATCHING_ENDPOINT,
                params={'key': self._safebrowsing_api_key},
                json=body,
            )

            response.raise_for_status()

            result = response.json()

            if not result:
                continue

            if 'matches' not in result:
                continue

            for match in result.get('matches'):
                result_url = match.get('threat', {}).get('url')

                if not result_url:
                    continue

                # Removing all key/values that are not defined in
                # the _SAFEBROWSING_ENTRY_KEEP.
                threat_result = match.copy()
                for key in match.keys():
                    if key not in self._SAFEBROWSING_ENTRY_KEEP:
                        threat_result.pop(key)

                results[result_url] = threat_result

        return results

    def _sanitize_url(self, url_entry):
        """Finds http[s]:// in 'url_entry' and returns its content from there.

        Args:
            url_entry: a URL, with some other characters before and after

        Returns:
            String with the URL or empty string if not found
        """
        m = self._URL_BEGINNING_RE.search(url_entry)

        if m:
            return m.group(1)

        return ''

    def run(self):
        """Entry point for the analyzer.

        Returns:
            String with summary of the analyzer result
        """
        # Exit ASAP if the API key is missing.
        if not self._safebrowsing_api_key:
            return 'Safe Browsing API requires an API key!'

        query = (
            '{"query": { "bool": { "should": [ '
            '{ "exists" : { "field" : "url" }} ] } } }')

        return_fields = ['url']

        events = self.event_stream(
            query_dsl=query,
            return_fields=return_fields,
        )

        urls = {}

        for event in events:
            url = self._sanitize_url(event.source.get('url'))

            if not url:
                continue

            urls.setdefault(url, []).append(event)

        # Exit early if there are no URLs in the data set to analyze.
        if not urls:
            return 'No URLs to analyze.'

        url_allowlisted = 0

        url_allowlist = set(
            interface.get_yaml_config(
                self._URL_ALLOW_LIST_CONFIG,
            ),
        )

        if not url_allowlist:
            domain_analyzer_allowlisted = current_app.config.get(
                'DOMAIN_ANALYZER_EXCLUDE_DOMAINS',
                [],
            )
            for domain in domain_analyzer_allowlisted:
                url_allowlist.add('*.%s/*' % domain)

        logging.info(
            '{0:d} entries on the allowlist.'.format(len(url_allowlist)),
        )

        safebrowsing_platforms = current_app.config.get(
            'SAFEBROWSING_PLATFORMS',
            ['ANY_PLATFORM'],
        )

        safebrowsing_types = current_app.config.get(
            'SAFEBROWSING_THREATTYPES',
            ['MALWARE'],
        )

        lookup_urls = []

        for url in urls:
            if self._is_url_allowlisted(url, url_allowlist):
                url_allowlisted += 1
                continue

            lookup_urls.append(url)

        try:
            safebrowsing_results = self._do_safebrowsing_lookup(
                lookup_urls,
                safebrowsing_platforms,
                safebrowsing_types,
            )
        except requests.HTTPError:
            return 'Couldn\'t reach the Safe Browsing API.'

        for url, events in urls.items():
            if url not in safebrowsing_results:
                continue

            safebrowsing_result = safebrowsing_results[url]
            for event in events:
                tags = ['google-safebrowsing-url']

                threat_type = safebrowsing_result.get('threatType')

                if threat_type:
                    tags.append(
                        'google-safebrowsing-%s' % threat_type.lower(),
                    )

                event.add_tags(tags)

                threat_attributes = []
                for item in safebrowsing_result.items():
                    threat_attributes.append('%s: %s' % item)

                event.add_attributes(
                    {
                        'google-safebrowsing-threat': ', '.join(
                            threat_attributes),
                    },
                )
                event.commit()

        return (
            '{0:d} Safe Browsing result(s) on {1:d} URL(s), '
            '{2:d} allow listed.').format(
                len(safebrowsing_results),
                len(urls),
                url_allowlisted,
            )

manager.AnalysisManager.register_analyzer(SafeBrowsingSketchPlugin)
<commit_msg>Logging HTTP errors<commit_after>"""Sketch analyzer plugin for the Safe Browsing API."""
from __future__ import unicode_literals

import fnmatch
import logging
import re

import pkg_resources
import requests
from flask import current_app

from timesketch.lib.analyzers import interface
from timesketch.lib.analyzers import manager

class SafeBrowsingSketchPlugin(interface.BaseSketchAnalyzer):
    """Sketch analyzer for Safe Browsing."""

    NAME = 'safebrowsing'

    # Safe Browsing API v4, threatMatches.find endpoint.
    _SAFE_BROWSING_THREATMATCHING_ENDPOINT = (
        'https://safebrowsing.googleapis.com/v4/threatMatches:find')

    # Maximal number of URLs in a single Lookup request as per the
    # Safe Browsing API documentation.
    _SAFE_BROWSING_BULK_LIMIT = 500

    # An optional file containing URL wildcards to be allow listed
    # in a YAML file.
    _URL_ALLOW_LIST_CONFIG = 'safebrowsing_allowlist.yaml'

    # The keys to be added to the TS event from the ThreatMatch object
    # we get from Safe Browsing API.
    _SAFEBROWSING_ENTRY_KEEP = frozenset([
        'platformType',
        'threatType',
    ])

    # Used to find proper URLs in the 'url' entries of TS events.
    _URL_BEGINNING_RE = re.compile(r'(http(s|):\/\/\S*)')

    def __init__(self, index_name, sketch_id):
        """Initialize The Sketch Analyzer.

        Args:
            index_name: Elasticsearch index name
            sketch_id: Sketch ID
        """
        super(SafeBrowsingSketchPlugin, self).__init__(index_name, sketch_id)

        self._safebrowsing_api_key = current_app.config.get(
            'SAFEBROWSING_API_KEY')

        self._google_client_id = current_app.config.get(
            'SAFEBROWSING_CLIENT_ID',
            'Timesketch',
        )

        self._google_client_version = current_app.config.get(
            'SAFEBROWSING_CLIENT_VERSION',
            pkg_resources.get_distribution('timesketch').version,
        )

    def _is_url_allowlisted(self, url, allowlist):
        """Does a fnmatch against the allowlist.

        Args:
            url: The url
            allowlist: The allowlist, list-like
        Returns:
            Boolean with the result
        """

        for url_pattern in allowlist:
            if fnmatch.fnmatchcase(url, url_pattern):
                return True

        return False

    def _do_safebrowsing_lookup(self, urls, platforms, types):
        """URL lookup against the Safe Browsing API.

        Args:
            urls: URLs
            platforms: platformTypes field of threatInfo
            types: threatTypes field of threatInfo
        Returns:
            Dict of URLs with the hits
        """
        results = {}

        api_client = {
            'clientId': self._google_client_id,
            'clientVersion': self._google_client_version,
        }

        for index in range(0, len(urls), self._SAFE_BROWSING_BULK_LIMIT):
            body = {
                'client': api_client,
                'threatInfo': {
                    'platformTypes': platforms,
                    'threatTypes': types,
                    'threatEntryTypes': ['URL'],
                    'threatEntries': [
                        {'url': url}
                        for url in urls[
                            index:index + self._SAFE_BROWSING_BULK_LIMIT]
                    ],
                },
            }

            response = requests.post(
                self._SAFE_BROWSING_THREATMATCHING_ENDPOINT,
                params={'key': self._safebrowsing_api_key},
                json=body,
            )

            try:
                response.raise_for_status()
            except requests.exceptions.HTTPError as e:
                logging.error(e)
                continue

            result = response.json()

            if not result:
                continue

            if 'matches' not in result:
                continue

            for match in result.get('matches'):
                result_url = match.get('threat', {}).get('url')

                if not result_url:
                    continue

                # Removing all key/values that are not defined in
                # the _SAFEBROWSING_ENTRY_KEEP.
                threat_result = match.copy()
                for key in match.keys():
                    if key not in self._SAFEBROWSING_ENTRY_KEEP:
                        threat_result.pop(key)

                results[result_url] = threat_result

        return results

    def _sanitize_url(self, url_entry):
        """Finds http[s]:// in 'url_entry' and returns its content from there.

        Args:
            url_entry: a URL, with some other characters before and after

        Returns:
            String with the URL or empty string if not found
        """
        m = self._URL_BEGINNING_RE.search(url_entry)

        if m:
            return m.group(1)

        return ''

    def run(self):
        """Entry point for the analyzer.

        Returns:
            String with summary of the analyzer result
        """
        # Exit ASAP if the API key is missing.
        if not self._safebrowsing_api_key:
            return 'Safe Browsing API requires an API key!'

        query = (
            '{"query": { "bool": { "should": [ '
            '{ "exists" : { "field" : "url" }} ] } } }')

        return_fields = ['url']

        events = self.event_stream(
            query_dsl=query,
            return_fields=return_fields,
        )

        urls = {}

        for event in events:
            url = self._sanitize_url(event.source.get('url'))

            if not url:
                continue

            urls.setdefault(url, []).append(event)

        # Exit early if there are no URLs in the data set to analyze.
        if not urls:
            return 'No URLs to analyze.'

        url_allowlisted = 0

        url_allowlist = set(
            interface.get_yaml_config(
                self._URL_ALLOW_LIST_CONFIG,
            ),
        )

        if not url_allowlist:
            domain_analyzer_allowlisted = current_app.config.get(
                'DOMAIN_ANALYZER_EXCLUDE_DOMAINS',
                [],
            )
            for domain in domain_analyzer_allowlisted:
                url_allowlist.add('*.%s/*' % domain)

        logging.info(
            '{0:d} entries on the allowlist.'.format(len(url_allowlist)),
        )

        safebrowsing_platforms = current_app.config.get(
            'SAFEBROWSING_PLATFORMS',
            ['ANY_PLATFORM'],
        )

        safebrowsing_types = current_app.config.get(
            'SAFEBROWSING_THREATTYPES',
            ['MALWARE'],
        )

        lookup_urls = []

        for url in urls:
            if self._is_url_allowlisted(url, url_allowlist):
                url_allowlisted += 1
                continue

            lookup_urls.append(url)

        try:
            safebrowsing_results = self._do_safebrowsing_lookup(
                lookup_urls,
                safebrowsing_platforms,
                safebrowsing_types,
            )
        except requests.HTTPError:
            return 'Couldn\'t reach the Safe Browsing API.'

        for url, events in urls.items():
            if url not in safebrowsing_results:
                continue

            safebrowsing_result = safebrowsing_results[url]
            for event in events:
                tags = ['google-safebrowsing-url']

                threat_type = safebrowsing_result.get('threatType')

                if threat_type:
                    tags.append(
                        'google-safebrowsing-%s' % threat_type.lower(),
                    )

                event.add_tags(tags)

                threat_attributes = []
                for item in safebrowsing_result.items():
                    threat_attributes.append('%s: %s' % item)

                event.add_attributes(
                    {
                        'google-safebrowsing-threat': ', '.join(
                            threat_attributes),
                    },
                )
                event.commit()

        return (
            '{0:d} Safe Browsing result(s) on {1:d} URL(s), '
            '{2:d} allow listed.').format(
                len(safebrowsing_results),
                len(urls),
                url_allowlisted,
            )

manager.AnalysisManager.register_analyzer(SafeBrowsingSketchPlugin)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# © 2016 Therp BV <http://therp.nl>
# © 2017 ThinkOpen Solutions <https://tkobr.com>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).

import logging
import os
import shutil
import subprocess
import tempfile
import threading
import time

from odoo.tests import common

ADMIN_USER_ID = common.ADMIN_USER_ID

from odoo import api, fields, models, _
from odoo.exceptions import UserError

_logger = logging.getLogger(__name__)
_MARKER_PHRASE = '[[waiting for OCR]]'
_PDF_OCR_DOCUMENTS_THREADS = []
OCR_LANGUAGE = [('afr', 'Afrikaans'),
                ('amh', 'Amharic'),
                ('ara', 'Arabic'),
                ('asm', 'Assamese'),
                ('aze', 'Azerbaijani'),
                ('aze_cyrl', 'Azerbaijani - Cyrilic'),
                ('bel', 'Belarusian'),
                ('ben', 'Bengali'),
                ('bod', 'Tibetan'),
                ('bos', 'Bosnian'),
                ('bul', 'Bulgarian'),
                ('cat', 'Catalan; Valencian'),
                ('ceb', 'Cebuano'),
                ('ces', 'Czech'),
                ('chi_sim', 'Chinese - Simplified'),
                ('chi_tra', 'Chinese - Traditional'),
                ('chr', 'Cherokee'),
                ('cym', 'Welsh'),
                ('dan', 'Danish'),
                ('dan_frak', 'Danish - Fraktur'),
                ('deu', 'German'),
                ('deu_frak', 'German - Fraktur'),
                ('dzo', 'Dzongkha'),
                ('ell', 'Greek, Modern (1453-)'),
                ('eng', 'English'),
                ('enm', 'English, Middle (1100-1500)'),
                ('epo', 'Esperanto'),
                ('equ', 'Math / equation detection module'),
                ('est', 'Estonian'),
                ('eus', 'Basque'),
                ('fas', 'Persian'),
                ('fin', 'Finnish'),
                ('fra', 'French'),
                ('frk', 'Frankish'),
                ('frm', 'French, Middle (ca.1400-1600)'),
                ('gle', 'Irish'),
                ('glg', 'Galician'),
                ('grc', 'Greek, Ancient (to 1453)'),
                ('guj', 'Gujarati'),
                ('hat', 'Haitian; Haitian Creole'),
                ('heb', 'Hebrew'),
                ('hin', 'Hindi'),
                ('hrv', 'Croatian'),
                ('hun', 'Hungarian'),
                ('iku', 'Inuktitut'),
                ('ind', 'Indonesian'),
                ('isl', 'Icelandic'),
                ('ita', 'Italian'),
                ('ita_old', 'Italian - Old'),
                ('jav', 'Javanese'),
                ('jpn', 'Japanese'),
                ('kan', 'Kannada'),
                ('kat', 'Georgian'),
                ('kat_old', 'Georgian - Old'),
                ('kaz', 'Kazakh'),
                ('khm', 'Central Khmer'),
                ('kir', 'Kirghiz; Kyrgyz'),
                ('kor', 'Korean'),
                ('kur', 'Kurdish'),
                ('lao', 'Lao'),
                ('lat', 'Latin'),
                ('lav', 'Latvian'),
                ('lit', 'Lithuanian'),
                ('mal', 'Malayalam'),
                ('mar', 'Marathi'),
                ('mkd', 'Macedonian'),
                ('mlt', 'Maltese'),
                ('msa', 'Malay'),
                ('mya', 'Burmese'),
                ('nep', 'Nepali'),
                ('nld', 'Dutch; Flemish'),
                ('nor', 'Norwegian'),
                ('ori', 'Oriya'),
                ('osd', 'Orientation and script detection module'),
                ('pan', 'Panjabi; Punjabi'),
                ('pol', 'Polish'),
                ('por', 'Portuguese'),
                ('pus', 'Pushto; Pashto'),
                ('ron', 'Romanian; Moldavian; Moldovan'),
                ('rus', 'Russian'),
                ('san', 'Sanskrit'),
                ('sin', 'Sinhala; Sinhalese'),
                ('slk', 'Slovak'),
                ('slk_frak', 'Slovak - Fraktur'),
                ('slv', 'Slovenian'),
                ('spa', 'Spanish; Castilian'),
                ('spa_old', 'Spanish; Castilian - Old'),
                ('sqi', 'Albanian'),
                ('srp', 'Serbian'),
                ('srp_latn', 'Serbian - Latin'),
                ('swa', 'Swahili'),
                ('swe', 'Swedish'),
                ('syr', 'Syriac'),
                ('tam', 'Tamil'),
                ('tel', 'Telugu'),
                ('tgk', 'Tajik'),
                ('tgl', 'Tagalog'),
                ('tha', 'Thai'),
                ('tir', 'Tigrinya'),
                ('tur', 'Turkish'),
                ('uig', 'Uighur; Uyghur'),
                ('ukr', 'Ukrainian'),
                ('urd', 'Urdu'),
                ('uzb', 'Uzbek'),
                ('uzb_cyrl', 'Uzbek - Cyrilic'),
                ('vie', 'Vietnamese'),
                ('yid', 'Yiddish'), ]

def ncpus():
    # for Linux, Unix and MacOS
    if hasattr(os, "sysconf"):
        if os.sysconf_names.has_key("SC_NPROCESSORS_ONLN"):
            # Linux and Unix
            ncpus = os.sysconf("SC_NPROCESSORS_ONLN")
            if isinstance(ncpus, int) and ncpus > 0:
                return ncpus
        else:
            # MacOS X
            return int(os.popen2("sysctl -n hw.ncpu")[1].read())
    # for Windows
    if os.environ.has_key("NUMBER_OF_PROCESSORS"):
        ncpus = int(os.environ["NUMBER_OF_PROCESSORS"])
        if ncpus > 0:
            return ncpus
    # return the default value
    return 1

_SEMAPHORES_DOC_POOL = threading.BoundedSemaphore(ncpus())
_SEMAPHORES_PAGE_POOL = threading.BoundedSemaphore(ncpus())

class IrAttachment(models.Model):
    _inherit = 'ir.attachment'

    language = fields.Selection(OCR_LANGUAGE, 'Language',
                                default=lambda self:
                                self.env['ir.config_parameter'].get_param(
                                    'document_ocr.language', 'eng'))
    # We need to redefine index_content field to be able to update it
    # on the onchange_language() in form
    index_content = fields.Text('Indexed Content',
                                readonly=False,
                                prefetch=False)
    index_content_rel = fields.Text(related='index_content',
                                    string='Indexed Content Rel')
    processing_time = fields.Char('Processing Time',
                                  readonly=True,
                                  copy=False,
                                  help='Processing time.\n'
                                       '(00:00:00 means less than one second)')

    @api.onchange('language')
    def onchange_language(self):
        process = subprocess.Popen(['tesseract', '--list-langs'],
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
        if self.language not in stderr.split('\n')[1:-1]:
            raise UserError(_(
                "Language not installed."
                " Please ask your system administrator to"
                " install tesseract '%s' language." %
                self.language))
        if self.store_fname:
            bin_data = self._file_read(self.store_fname)
        else:
            bin_data = self.db_datas
        if bin_data:
            index_content = self._index(
                bin_data.decode('base64'), self.datas_fname, self.mimetype)
            return {'value': {
                'index_content': index_content}}
        return {'value': {}}

    @api.model
    def _index(self, bin_data, datas_fname, mimetype):
        content = super(IrAttachment, self)._index(
            bin_data, datas_fname, mimetype)
        if not content or content == 'image':
            has_synchr_param = self.env['ir.config_parameter'].get_param(
                'document_ocr.synchronous', 'False') == 'True'
            has_force_flag = self.env.context.get('document_ocr_force')
            synchr = has_synchr_param or has_force_flag
            if synchr:
                content = self._index_ocr(bin_data)
                self.env.cr.commit()
            else:
                content = _MARKER_PHRASE
        return content

    def _index_ocr(self, bin_data):
        if self.datas_fname:
            process = subprocess.Popen(
                ['tesseract', 'stdin', 'stdout', '-l', self.language],
                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = process.communicate(bin_data)
            if stderr:
                _logger.warning('Error during OCR: %s', stderr)
            return stdout
        else:
            _logger.warning('OCR IMAGE "%s", no image to process...',
                            self.name)
        return False

    def _ocr_image_thread(self, i, t, image):
        global ocr_images_text
        with _SEMAPHORES_PAGE_POOL:
            with threading.Lock():
                _logger.info('OCR PDF INFO "%s" image %d/%d to text...',
                             self.name, i, t)
                ocr_images_text[self.id][i] = self._index_ocr(image)

    def _index_doc_pdf_thread(self, bin_data):
        global ocr_images_text
        with _SEMAPHORES_DOC_POOL:
            with threading.Lock():
                try:
                    ocr_images_text[self.id] = {}
                    tmpdir = tempfile.mkdtemp()
                    _logger.info('OCR PDF INFO "%s"...', self.name)
                    time_start = time.time()
                    stdout, stderr = subprocess.Popen(
                        ['pdftotext', '-layout', '-nopgbrk', '-', '-'],
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE).communicate(bin_data)
                    if stderr:
                        _logger.warning('OCR PDF ERROR to text: %s',
                                        stderr)
                    buf = stdout
                    # OCR PDF Images
                    stdout, stderr = subprocess.Popen(
                        ['pdfimages', '-p', '-', tmpdir + '/ocr'],
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE).communicate(bin_data)
                    if stderr:
                        _logger.warning('OCR PDF WARNING Images: %s',
                                        stderr)
                    # OCR every image greater than 50Kb
                    filelist = sorted([(file) for file
                                       in os.listdir(tmpdir)
                                       if os.path.getsize(
                            os.path.join(tmpdir, file)) > 50000])
                    filelist_size = len(filelist)
                    count = 1
                    workers = []
                    for pdf in filelist:
                        img_file = os.path.join(tmpdir, pdf)
                        image = open(img_file, 'rb').read()
                        t = threading.Thread(target=self._ocr_image_thread,
                                             name=u'ocr_image_' + str(count),
                                             args=(count,
                                                   filelist_size,
                                                   image))
                        t.start()
                        count += 1
                        workers.append(t)
                    for t in workers:
                        t.join()
                    index_content = buf
                    for text in sorted(ocr_images_text[self.id]):
                        try:
                            index_content = \
                                u'%s\n%s' % (
                                    index_content,
                                    ocr_images_text[self.id][text].decode('utf8'))
                        except:
                            index_content = \
                                u'%s\n%s' % (
                                    index_content.decode('utf8'),
                                    ocr_images_text[self.id][text])
                    ocr_images_text.pop(self.id)  # release memory
                    m, s = divmod((time.time() - time_start), 60)
                    h, m = divmod(m, 60)
                    self.index_content = index_content
                    self.processing_time = "%02d:%02d:%02d" % (h, m, s)
                    shutil.rmtree(tmpdir)
                except Exception, e:
                    ocr_images_text.pop(self.id)  # release memory
                    shutil.rmtree(tmpdir)
                    _logger.warning('OCR PDF DOC ERROR: %s', e[0])
        return True

    def _index_pdf(self, bin_data):
        global ocr_images_text
        buf = _MARKER_PHRASE
        has_synchr_param = self.env['ir.config_parameter'].get_param(
            'document_ocr.synchronous', 'False') == 'True'
        has_force_flag = self.env.context.get('document_ocr_force')
        synchr = has_synchr_param or has_force_flag
        try:
            if ocr_images_text:
                pass
        except:
            ocr_images_text = {}
        if synchr:
            t = threading.Thread(target=self._index_doc_pdf_thread,
                                 name=u'index_pdf_' + str(self.id),
                                 args=[bin_data])
            t.start()
        else:
            buf = _MARKER_PHRASE
        return buf

    @api.model
    def _ocr_cron(self):
        for this in self.with_context(document_ocr_force=True).search(
                [('index_content', '=', _MARKER_PHRASE)]):
            if not this.datas:
                continue
            index_content = this._index(
                this.datas.decode('base64'), this.datas_fname, this.mimetype)
            this.write({
                'index_content': index_content,
            })
<commit_msg>Remove _PDF_OCR_DOCUMENTS_THREADS<commit_after># -*- coding: utf-8 -*-
# © 2016 Therp BV <http://therp.nl>
# © 2017 ThinkOpen Solutions <https://tkobr.com>
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl.html).

import logging
import os
import shutil
import subprocess
import tempfile
import threading
import time

from odoo.tests import common

ADMIN_USER_ID = common.ADMIN_USER_ID

from odoo import api, fields, models, _
from odoo.exceptions import UserError

_logger = logging.getLogger(__name__)
_MARKER_PHRASE = '[[waiting for OCR]]'
OCR_LANGUAGE = [('afr', 'Afrikaans'),
                ('amh', 'Amharic'),
                ('ara', 'Arabic'),
                ('asm', 'Assamese'),
                ('aze', 'Azerbaijani'),
                ('aze_cyrl', 'Azerbaijani - Cyrilic'),
                ('bel', 'Belarusian'),
                ('ben', 'Bengali'),
                ('bod', 'Tibetan'),
                ('bos', 'Bosnian'),
                ('bul', 'Bulgarian'),
                ('cat', 'Catalan; Valencian'),
                ('ceb', 'Cebuano'),
                ('ces', 'Czech'),
                ('chi_sim', 'Chinese - Simplified'),
                ('chi_tra', 'Chinese - Traditional'),
                ('chr', 'Cherokee'),
                ('cym', 'Welsh'),
                ('dan', 'Danish'),
                ('dan_frak', 'Danish - Fraktur'),
                ('deu', 'German'),
                ('deu_frak', 'German - Fraktur'),
                ('dzo', 'Dzongkha'),
                ('ell', 'Greek, Modern (1453-)'),
                ('eng', 'English'),
                ('enm', 'English, Middle (1100-1500)'),
                ('epo', 'Esperanto'),
                ('equ', 'Math / equation detection module'),
                ('est', 'Estonian'),
                ('eus', 'Basque'),
                ('fas', 'Persian'),
                ('fin', 'Finnish'),
                ('fra', 'French'),
                ('frk', 'Frankish'),
                ('frm', 'French, Middle (ca.1400-1600)'),
                ('gle', 'Irish'),
                ('glg', 'Galician'),
                ('grc', 'Greek, Ancient (to 1453)'),
                ('guj', 'Gujarati'),
                ('hat', 'Haitian; Haitian Creole'),
                ('heb', 'Hebrew'),
                ('hin', 'Hindi'),
                ('hrv', 'Croatian'),
                ('hun', 'Hungarian'),
                ('iku', 'Inuktitut'),
                ('ind', 'Indonesian'),
                ('isl', 'Icelandic'),
                ('ita', 'Italian'),
                ('ita_old', 'Italian - Old'),
                ('jav', 'Javanese'),
                ('jpn', 'Japanese'),
                ('kan', 'Kannada'),
                ('kat', 'Georgian'),
                ('kat_old', 'Georgian - Old'),
                ('kaz', 'Kazakh'),
                ('khm', 'Central Khmer'),
                ('kir', 'Kirghiz; Kyrgyz'),
                ('kor', 'Korean'),
                ('kur', 'Kurdish'),
                ('lao', 'Lao'),
                ('lat', 'Latin'),
                ('lav', 'Latvian'),
                ('lit', 'Lithuanian'),
                ('mal', 'Malayalam'),
                ('mar', 'Marathi'),
                ('mkd', 'Macedonian'),
                ('mlt', 'Maltese'),
                ('msa', 'Malay'),
                ('mya', 'Burmese'),
                ('nep', 'Nepali'),
                ('nld', 'Dutch; Flemish'),
                ('nor', 'Norwegian'),
                ('ori', 'Oriya'),
                ('osd', 'Orientation and script detection module'),
                ('pan', 'Panjabi; Punjabi'),
                ('pol', 'Polish'),
                ('por', 'Portuguese'),
                ('pus', 'Pushto; Pashto'),
                ('ron', 'Romanian; Moldavian; Moldovan'),
                ('rus', 'Russian'),
                ('san', 'Sanskrit'),
                ('sin', 'Sinhala; Sinhalese'),
                ('slk', 'Slovak'),
                ('slk_frak', 'Slovak - Fraktur'),
                ('slv', 'Slovenian'),
                ('spa', 'Spanish; Castilian'),
                ('spa_old', 'Spanish; Castilian - Old'),
                ('sqi', 'Albanian'),
                ('srp', 'Serbian'),
                ('srp_latn', 'Serbian - Latin'),
                ('swa', 'Swahili'),
                ('swe', 'Swedish'),
                ('syr', 'Syriac'),
                ('tam', 'Tamil'),
                ('tel', 'Telugu'),
                ('tgk', 'Tajik'),
                ('tgl', 'Tagalog'),
                ('tha', 'Thai'),
                ('tir', 'Tigrinya'),
                ('tur', 'Turkish'),
                ('uig', 'Uighur; Uyghur'),
                ('ukr', 'Ukrainian'),
                ('urd', 'Urdu'),
                ('uzb', 'Uzbek'),
                ('uzb_cyrl', 'Uzbek - Cyrilic'),
                ('vie', 'Vietnamese'),
                ('yid', 'Yiddish'), ]

def ncpus():
    # for Linux, Unix and MacOS
    if hasattr(os, "sysconf"):
        if os.sysconf_names.has_key("SC_NPROCESSORS_ONLN"):
            # Linux and Unix
            ncpus = os.sysconf("SC_NPROCESSORS_ONLN")
            if isinstance(ncpus, int) and ncpus > 0:
                return ncpus
        else:
            # MacOS X
            return int(os.popen2("sysctl -n hw.ncpu")[1].read())
    # for Windows
    if os.environ.has_key("NUMBER_OF_PROCESSORS"):
        ncpus = int(os.environ["NUMBER_OF_PROCESSORS"])
        if ncpus > 0:
            return ncpus
    # return the default value
    return 1

_SEMAPHORES_DOC_POOL = threading.BoundedSemaphore(ncpus())
_SEMAPHORES_PAGE_POOL = threading.BoundedSemaphore(ncpus())

class IrAttachment(models.Model):
    _inherit = 'ir.attachment'

    language = fields.Selection(OCR_LANGUAGE, 'Language',
                                default=lambda self:
                                self.env['ir.config_parameter'].get_param(
                                    'document_ocr.language', 'eng'))
    # We need to redefine index_content field to be able to update it
    # on the onchange_language() in form
    index_content = fields.Text('Indexed Content',
                                readonly=False,
                                prefetch=False)
    index_content_rel = fields.Text(related='index_content',
                                    string='Indexed Content Rel')
    processing_time = fields.Char('Processing Time',
                                  readonly=True,
                                  copy=False,
                                  help='Processing time.\n'
                                       '(00:00:00 means less than one second)')

    @api.onchange('language')
    def onchange_language(self):
        process = subprocess.Popen(['tesseract', '--list-langs'],
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
        if self.language not in stderr.split('\n')[1:-1]:
            raise UserError(_(
                "Language not installed."
                " Please ask your system administrator to"
                " install tesseract '%s' language." %
                self.language))
        if self.store_fname:
            bin_data = self._file_read(self.store_fname)
        else:
            bin_data = self.db_datas
        if bin_data:
            index_content = self._index(
                bin_data.decode('base64'), self.datas_fname, self.mimetype)
            return {'value': {
                'index_content': index_content}}
        return {'value': {}}

    @api.model
    def _index(self, bin_data, datas_fname, mimetype):
        content = super(IrAttachment, self)._index(
            bin_data, datas_fname, mimetype)
        if not content or content == 'image':
            has_synchr_param = self.env['ir.config_parameter'].get_param(
                'document_ocr.synchronous', 'False') == 'True'
            has_force_flag = self.env.context.get('document_ocr_force')
            synchr = has_synchr_param or has_force_flag
            if synchr:
                content = self._index_ocr(bin_data)
                self.env.cr.commit()
            else:
                content = _MARKER_PHRASE
        return content

    def _index_ocr(self, bin_data):
        if self.datas_fname:
            process = subprocess.Popen(
                ['tesseract', 'stdin', 'stdout', '-l', self.language],
                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = process.communicate(bin_data)
            if stderr:
                _logger.warning('Error during OCR: %s', stderr)
            return stdout
        else:
            _logger.warning('OCR IMAGE "%s", no image to process...',
                            self.name)
        return False

    def _ocr_image_thread(self, i, t, image):
        global ocr_images_text
        with _SEMAPHORES_PAGE_POOL:
            with threading.Lock():
                _logger.info('OCR PDF INFO "%s" image %d/%d to text...',
                             self.name, i, t)
                ocr_images_text[self.id][i] = self._index_ocr(image)

    def _index_doc_pdf_thread(self, bin_data):
        global ocr_images_text
        with _SEMAPHORES_DOC_POOL:
            with threading.Lock():
                try:
                    ocr_images_text[self.id] = {}
                    tmpdir = tempfile.mkdtemp()
                    _logger.info('OCR PDF INFO "%s"...', self.name)
                    time_start = time.time()
                    stdout, stderr = subprocess.Popen(
                        ['pdftotext', '-layout', '-nopgbrk', '-', '-'],
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE).communicate(bin_data)
                    if stderr:
                        _logger.warning('OCR PDF ERROR to text: %s',
                                        stderr)
                    buf = stdout
                    # OCR PDF Images
                    stdout, stderr = subprocess.Popen(
                        ['pdfimages', '-p', '-', tmpdir + '/ocr'],
                        stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE).communicate(bin_data)
                    if stderr:
                        _logger.warning('OCR PDF WARNING Images: %s',
                                        stderr)
                    # OCR every image greater than 50Kb
                    filelist = sorted([(file) for file
                                       in os.listdir(tmpdir)
                                       if os.path.getsize(
                            os.path.join(tmpdir, file)) > 50000])
                    filelist_size = len(filelist)
                    count = 1
                    workers = []
                    for pdf in filelist:
                        img_file = os.path.join(tmpdir, pdf)
                        image = open(img_file, 'rb').read()
                        t = threading.Thread(target=self._ocr_image_thread,
                                             name=u'ocr_image_' + str(count),
                                             args=(count,
                                                   filelist_size,
                                                   image))
                        t.start()
                        count += 1
                        workers.append(t)
                    for t in workers:
                        t.join()
                    index_content = buf
                    for text in sorted(ocr_images_text[self.id]):
                        try:
                            index_content = \
                                u'%s\n%s' % (
                                    index_content,
                                    ocr_images_text[self.id][text].decode('utf8'))
                        except:
                            index_content = \
                                u'%s\n%s' % (
                                    index_content.decode('utf8'),
                                    ocr_images_text[self.id][text])
                    ocr_images_text.pop(self.id)  # release memory
                    m, s = divmod((time.time() - time_start), 60)
                    h, m = divmod(m, 60)
                    self.index_content = index_content
                    self.processing_time = "%02d:%02d:%02d" % (h, m, s)
                    shutil.rmtree(tmpdir)
                except Exception, e:
                    ocr_images_text.pop(self.id)  # release memory
                    shutil.rmtree(tmpdir)
                    _logger.warning('OCR PDF DOC ERROR: %s', e[0])
        return True

    def _index_pdf(self, bin_data):
        global ocr_images_text
        buf = _MARKER_PHRASE
        has_synchr_param = self.env['ir.config_parameter'].get_param(
            'document_ocr.synchronous', 'False') == 'True'
        has_force_flag = self.env.context.get('document_ocr_force')
        synchr = has_synchr_param or has_force_flag
        try:
            if ocr_images_text:
                pass
        except:
            ocr_images_text = {}
        if synchr:
            t = threading.Thread(target=self._index_doc_pdf_thread,
                                 name=u'index_pdf_' + str(self.id),
                                 args=[bin_data])
            t.start()
        else:
            buf = _MARKER_PHRASE
        return buf

    @api.model
    def _ocr_cron(self):
        for this in self.with_context(document_ocr_force=True).search(
                [('index_content', '=', _MARKER_PHRASE)]):
            if not this.datas:
                continue
            index_content = this._index(
                this.datas.decode('base64'), this.datas_fname, this.mimetype)
            this.write({
                'index_content': index_content,
            })
<|endoftext|>


<commit_before>from zeit.content.image.testing import create_image_group_with_master_image
import gocept.testing.assertion
import mock
import pkg_resources
import zeit.cms.testing
import zeit.content.image.testing
import zeit.edit.interfaces
import zeit.edit.rule
import zope.component

class ImageGroupHelperMixin(object):

    def add_imagegroup(self):
        b = self.browser
        b.open('http://localhost/++skin++cms/repository/2006/')
        menu = b.getControl(name='add_menu')
        menu.displayValue = ['Image group']
        b.open(menu.value[0])

    def set_imagegroup_filename(self, filename):
        self.browser.getControl('File name').value = filename

    def set_imagegroup_title(self, title):
        self.browser.getControl('Image title').value = title

    def set_display_type(self, display_type):
        self.browser.getControl('Display Type').displayValue = [display_type]

    def fill_copyright_information(self):
        b = self.browser
        b.getControl(name='form.copyrights.0..combination_00').value = (
            'ZEIT ONLINE')
        b.getControl(name='form.copyrights.0..combination_01').value = (
            'http://www.zeit.de/')

    def _upload_image(self, field, filename):
        self.browser.getControl(name='form.{}'.format(field)).add_file(
            pkg_resources.resource_stream(
                'zeit.content.image.browser',
                'testdata/{}'.format(filename)),
            'image/jpeg', filename)

    def upload_primary_image(self, filename):
        self._upload_image('master_image_blobs.0.', filename)

    def upload_secondary_image(self, filename):
        self._upload_image('master_image_blobs.1.', filename)

    def upload_tertiary_image(self, filename):
        self._upload_image('master_image_blobs.2.', filename)

    def save_imagegroup(self):
        self.browser.getControl(name='form.actions.add').click()

class ImageGroupGhostTest(
        zeit.cms.testing.BrowserTestCase,
        ImageGroupHelperMixin):

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_adding_imagegroup_adds_a_ghost(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('new-hampshire')
        self.set_imagegroup_title('New Hampshire')
        self.upload_primary_image('new-hampshire-artikel.jpg')
        self.fill_copyright_information()
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            with zeit.cms.testing.interaction():
                wc = zeit.cms.checkout.interfaces.IWorkingcopy(None)
                self.assertEqual(1, len(wc))

class ImageGroupPublishTest(zeit.cms.testing.BrowserTestCase):
    """Integration test for zeit.workflow.browser.publish.Publish.

    Checks that adapter to use ValidatingWorkflow was set up correctly.

    """

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_validation_errors_are_displayed_during_publish(self):
        rm = zope.component.getUtility(zeit.edit.interfaces.IRulesManager)
        rules = [rm.create_rule(['error_if(True, "Custom Error")'], 0)]
        with mock.patch.object(zeit.edit.rule.RulesManager, 'rules', rules):
            b = self.browser
            b.open('http://localhost/++skin++vivi/repository/2007/03/group'
                   '/@@publish.html')
        self.assertEllipsis('...Custom Error...', b.contents)

class ImageGroupBrowserTest(
        zeit.cms.testing.BrowserTestCase,
        ImageGroupHelperMixin,
        gocept.testing.assertion.String):

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_traversing_thumbnail_yields_images(self):
        with zeit.cms.testing.site(self.getRootFolder()):
            with zeit.cms.testing.interaction():
                create_image_group_with_master_image()
        b = self.browser
        b.open('http://localhost/++skin++vivi/repository'
               '/group/thumbnails/square/@@raw')
        self.assertEqual('image/jpeg', b.headers['Content-Type'])

    def test_primary_master_image_is_marked_for_desktop_viewport(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('image-group')
        self.upload_primary_image('opernball.jpg')
        self.fill_copyright_information()
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['2006']['image-group']
        self.assertEqual(1, len(group.master_images))
        self.assertEqual('desktop', group.master_images[0][0])
        self.assertEqual('opernball.jpg', group.master_images[0][1])

    def test_secondary_master_image_is_marked_for_mobile_viewport(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('image-group')
        self.browser.getControl('Add motif').click()
        self.upload_primary_image('opernball.jpg')
        self.upload_secondary_image('new-hampshire-artikel.jpg')
        self.fill_copyright_information()
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['2006']['image-group']
        self.assertEqual(2, len(group.master_images))
        self.assertEqual('mobile', group.master_images[1][0])
        self.assertEqual(
            'new-hampshire-artikel.jpg', group.master_images[1][1])

    def test_tertiary_master_image_has_no_viewport(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('image-group')
        self.browser.getControl('Add motif').click()
        self.browser.getControl('Add motif').click()
        self.upload_primary_image('opernball.jpg')
        self.upload_secondary_image('new-hampshire-artikel.jpg')
        self.upload_tertiary_image('obama-clinton-120x120.jpg')
        self.fill_copyright_information()
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['2006']['image-group']
        self.assertEqual(2, len(group.master_images))
        self.assertEqual(3, len(group.keys()))
        self.assertIn('obama-clinton-120x120.jpg', group.keys())

    def test_thumbnail_source_is_created_on_add(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('image-group')
        with zeit.cms.testing.site(self.getRootFolder()):
            img = zeit.cms.interfaces.ICMSContent(
                'http://xml.zeit.de/2006/DSC00109_2.JPG')
        self.browser.getControl(name='form.master_image_blobs.0.').add_file(
            img.open(), 'image/jpeg', 'master-image.jpg')
        self.fill_copyright_information()
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['2006']['image-group']
        self.assertIn('thumbnail-source-master-image.jpg', group)

    def test_display_type_imagegroup_opens_variant_html_on_save(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('imagegroup')
        self.fill_copyright_information()
        self.set_display_type('Bildergruppe')
        self.save_imagegroup()
        self.assertEndsWith('@@variant.html', self.browser.url)

    def test_display_type_infographic_opens_view_html_on_save(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('infographic')
        self.fill_copyright_information()
        self.set_display_type('Infografik')
        self.save_imagegroup()
        self.assertEndsWith('@@view.html', self.browser.url)

    def test_display_type_imagegroup_shows_edit_tab(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('infographic')
        self.fill_copyright_information()
        self.set_display_type('Bildergruppe')
        self.save_imagegroup()
        self.assertIn('repository/2006/infographic/@@variant.html',
                      self.browser.contents)

    def test_display_type_infographic_hides_edit_tab(self):
        self.add_imagegroup()
        self.set_imagegroup_filename('infographic')
        self.fill_copyright_information()
        self.set_display_type('Infografik')
        self.save_imagegroup()
        self.assertNotIn('repository/2006/infographic/@@variant.html',
                         self.browser.contents)

class ThumbnailTest(zeit.cms.testing.FunctionalTestCase):

    layer = zeit.content.image.testing.ZCML_LAYER

    def setUp(self):
        from zeit.content.image.browser.imagegroup import Thumbnail
        super(ThumbnailTest, self).setUp()
        self.group = create_image_group_with_master_image()
        self.thumbnail = Thumbnail()
        self.thumbnail.context = self.group

    def test_defaults_to_master_image(self):
        self.assertEqual(
            self.group['master-image.jpg'], self.thumbnail._find_image())

    def test_uses_materialized_image_if_present(self):
        from zeit.content.image.testing import create_local_image
        self.group['image-540x304.jpg'] = create_local_image(
            'obama-clinton-120x120.jpg')
        self.assertEqual(
            self.group['image-540x304.jpg'], self.thumbnail._find_image())
<commit_msg>ZON-3363: Refactor ``ImageGroupHelperMixin`` to make tests more readable.<commit_after>from zeit.content.image.testing import create_image_group_with_master_image
import gocept.testing.assertion
import mock
import pkg_resources
import zeit.cms.testing
import zeit.content.image.testing
import zeit.edit.interfaces
import zeit.edit.rule
import zope.component

class ImageGroupHelperMixin(object):

    def add_imagegroup(self, filename='imagegroup', fill_copyright=True):
        b = self.browser
        b.open('http://localhost/++skin++cms/repository/')
        menu = b.getControl(name='add_menu')
        menu.displayValue = ['Image group']
        b.open(menu.value[0])

        # Those information are required but rarely used in tests.
        self.set_filename(filename)
        if fill_copyright:
            self.fill_copyright_information()

    def add_motif(self):
        self.browser.getControl('Add motif').click()

    def set_filename(self, filename):
        self.browser.getControl('File name').value = filename

    def set_title(self, title):
        self.browser.getControl('Image title').value = title

    def set_display_type(self, display_type):
        self.browser.getControl('Display Type').displayValue = [display_type]

    def fill_copyright_information(self):
        b = self.browser
        b.getControl(name='form.copyrights.0..combination_00').value = (
            'ZEIT ONLINE')
        b.getControl(name='form.copyrights.0..combination_01').value = (
            'http://www.zeit.de/')

    def _upload_image(self, field, filename):
        self.browser.getControl(name='form.{}'.format(field)).add_file(
            pkg_resources.resource_stream(
                'zeit.content.image.browser',
                'testdata/{}'.format(filename)),
            'image/jpeg', filename)

    def upload_primary_image(self, filename):
        self._upload_image('master_image_blobs.0.', filename)

    def upload_secondary_image(self, filename):
        self._upload_image('master_image_blobs.1.', filename)

    def upload_tertiary_image(self, filename):
        self._upload_image('master_image_blobs.2.', filename)

    def save_imagegroup(self):
        self.browser.getControl(name='form.actions.add').click()

class ImageGroupGhostTest(
        zeit.cms.testing.BrowserTestCase,
        ImageGroupHelperMixin):

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_adding_imagegroup_adds_a_ghost(self):
        self.add_imagegroup()
        self.set_title('New Hampshire')
        self.upload_primary_image('new-hampshire-artikel.jpg')
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            with zeit.cms.testing.interaction():
                wc = zeit.cms.checkout.interfaces.IWorkingcopy(None)
                self.assertEqual(1, len(wc))

class ImageGroupPublishTest(zeit.cms.testing.BrowserTestCase):
    """Integration test for zeit.workflow.browser.publish.Publish.

    Checks that adapter to use ValidatingWorkflow was set up correctly.

    """

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_validation_errors_are_displayed_during_publish(self):
        rm = zope.component.getUtility(zeit.edit.interfaces.IRulesManager)
        rules = [rm.create_rule(['error_if(True, "Custom Error")'], 0)]
        with mock.patch.object(zeit.edit.rule.RulesManager, 'rules', rules):
            b = self.browser
            b.open('http://localhost/++skin++vivi/repository/2007/03/group'
                   '/@@publish.html')
        self.assertEllipsis('...Custom Error...', b.contents)

class ImageGroupBrowserTest(
        zeit.cms.testing.BrowserTestCase,
        ImageGroupHelperMixin,
        gocept.testing.assertion.String):

    layer = zeit.content.image.testing.ZCML_LAYER

    def test_traversing_thumbnail_yields_images(self):
        with zeit.cms.testing.site(self.getRootFolder()):
            with zeit.cms.testing.interaction():
                create_image_group_with_master_image()
        b = self.browser
        b.open('http://localhost/++skin++vivi/repository'
               '/group/thumbnails/square/@@raw')
        self.assertEqual('image/jpeg', b.headers['Content-Type'])

    def test_primary_master_image_is_marked_for_desktop_viewport(self):
        self.add_imagegroup()
        self.upload_primary_image('opernball.jpg')
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['imagegroup']
        self.assertEqual(1, len(group.master_images))
        self.assertEqual('desktop', group.master_images[0][0])
        self.assertEqual('opernball.jpg', group.master_images[0][1])

    def test_secondary_master_image_is_marked_for_mobile_viewport(self):
        self.add_imagegroup()
        self.add_motif()
        self.upload_primary_image('opernball.jpg')
        self.upload_secondary_image('new-hampshire-artikel.jpg')
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['imagegroup']
        self.assertEqual(2, len(group.master_images))
        self.assertEqual('mobile', group.master_images[1][0])
        self.assertEqual(
            'new-hampshire-artikel.jpg', group.master_images[1][1])

    def test_tertiary_master_image_has_no_viewport(self):
        self.add_imagegroup()
        self.add_motif()
        self.add_motif()
        self.upload_primary_image('opernball.jpg')
        self.upload_secondary_image('new-hampshire-artikel.jpg')
        self.upload_tertiary_image('obama-clinton-120x120.jpg')
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['imagegroup']
        self.assertEqual(2, len(group.master_images))
        self.assertEqual(3, len(group.keys()))
        self.assertIn('obama-clinton-120x120.jpg', group.keys())

    def test_thumbnail_source_is_created_on_add(self):
        self.add_imagegroup()
        with zeit.cms.testing.site(self.getRootFolder()):
            img = zeit.cms.interfaces.ICMSContent(
                'http://xml.zeit.de/2006/DSC00109_2.JPG')
        self.browser.getControl(name='form.master_image_blobs.0.').add_file(
            img.open(), 'image/jpeg', 'master-image.jpg')
        self.save_imagegroup()

        with zeit.cms.testing.site(self.getRootFolder()):
            group = self.repository['imagegroup']
        self.assertIn('thumbnail-source-master-image.jpg', group)

    def test_display_type_imagegroup_opens_variant_html_on_save(self):
        self.add_imagegroup()
        self.set_display_type('Bildergruppe')
        self.save_imagegroup()
        self.assertEndsWith('@@variant.html', self.browser.url)

    def test_display_type_infographic_opens_view_html_on_save(self):
        self.add_imagegroup()
        self.set_display_type('Infografik')
        self.save_imagegroup()
        self.assertEndsWith('@@view.html', self.browser.url)

    def test_display_type_imagegroup_shows_edit_tab(self):
        self.add_imagegroup()
        self.set_display_type('Bildergruppe')
        self.save_imagegroup()
        self.assertIn(
            'repository/imagegroup/@@variant.html', self.browser.contents)

    def test_display_type_infographic_hides_edit_tab(self):
        self.add_imagegroup()
        self.set_display_type('Infografik')
        self.save_imagegroup()
        self.assertNotIn(
            'repository/imagegroup/@@variant.html', self.browser.contents)

class ThumbnailTest(zeit.cms.testing.FunctionalTestCase):

    layer = zeit.content.image.testing.ZCML_LAYER

    def setUp(self):
        from zeit.content.image.browser.imagegroup import Thumbnail
        super(ThumbnailTest, self).setUp()
        self.group = create_image_group_with_master_image()
        self.thumbnail = Thumbnail()
        self.thumbnail.context = self.group

    def test_defaults_to_master_image(self):
        self.assertEqual(
            self.group['master-image.jpg'], self.thumbnail._find_image())

    def test_uses_materialized_image_if_present(self):
        from zeit.content.image.testing import create_local_image
        self.group['image-540x304.jpg'] = create_local_image(
            'obama-clinton-120x120.jpg')
        self.assertEqual(
            self.group['image-540x304.jpg'], self.thumbnail._find_image())
<|endoftext|>


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

#   Copyright 2014 Jason F Nicholls
#
#   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.

#   This Python Script is intended to be used in conjunction with PiMame
#   More specifically Mame4All
#   However you can use it for what ever purpose you see a need
#	Basically it reads the GPIO board of the Pi, and waits for
#   Events, if these events occur it reads 2 MCP23017 ICs for input
#   You can extend the Pi and use 8 ICs just extend this code
#   Note that ALL the pins of the ICs are setup as output.
#
#   For more information see picade.blogspot.com
#
#   Hope you have as much fun with your Pi as I have had with mine.

import RPi.GPIO as GPIO
import uinput
import smbus
import time

# Custom Key Configuration, these keys will be mapped to your controls, below are the default MAME Controls.
DEFINED_KEYS = {
    'MCPA_A_0': uinput.KEY_5,      	 # Coin 1
    'MCPA_A_1': uinput.KEY_6,   	 # Coin 2
    'MCPA_A_2': uinput.KEY_1,  		 # Start 1
    'MCPA_A_3': uinput.KEY_2,  		 # Start 2
    'MCPA_A_4': uinput.KEY_ESC, 	 # Escape - exit the game
    'MCPA_A_5': uinput.KEY_P,  		 # Pause
    'MCPA_A_6': uinput.KEY_ENTER, 	 # Enter
    'MCPA_A_7': uinput.KEY_TAB, 	 # Tab - Access in game menu
    'MCPA_B_0': uinput.KEY_UP,  	 # Player 1 Up
    'MCPA_B_1': uinput.KEY_DOWN, 	 # Player 1 Down
    'MCPA_B_2': uinput.KEY_LEFT,	 # Player 1 Left
    'MCPA_B_3': uinput.KEY_RIGHT, 	 # Player 1 Right
    'MCPA_B_4': uinput.KEY_R,		 # Player 2 Up
    'MCPA_B_5': uinput.KEY_F,		 # Player 2 Down
    'MCPA_B_6': uinput.KEY_D,		 # Player 2 Left
    'MCPA_B_7': uinput.KEY_G,		 # Player 2 Right
    'MCPB_A_0': uinput.KEY_LEFTCTRL,     # Player 1 Fire 1
    'MCPB_A_1': uinput.KEY_LEFTALT,	 # Player 1 Fire 2
    'MCPB_A_2': uinput.KEY_SPACE,	 # Player 1 Fire 3
    'MCPB_A_3': uinput.KEY_LEFTSHIFT,    # Player 1 Fire 4
    'MCPB_A_4': uinput.KEY_Z,	 	 # Player 1 Fire 5
    'MCPB_A_5': uinput.KEY_X,	         # Player 1 Fire 6
    'MCPB_A_6': uinput.KEY_SPACE,	 # Not used
    'MCPB_A_7': uinput.KEY_SPACE,        # Not used
    'MCPB_B_0': uinput.KEY_A,		 # Player 2 Fire 1
    'MCPB_B_1': uinput.KEY_S,		 # Player 2 Fire 2
    'MCPB_B_2': uinput.KEY_Q,		 # Player 2 Fire 3
    'MCPB_B_3': uinput.KEY_W,		 # Player 2 Fire 4
    'MCPB_B_4': uinput.KEY_E,		 # Player 2 Fire 5 // Normally Joystick 4 and is not set by default
    'MCPB_B_5': uinput.KEY_C,		 # Player 2 Fire 6 // Normally Joystick 5 and is not set by default
    'MCPB_B_6': uinput.KEY_SPACE,	 # Not used
    'MCPB_B_7': uinput.KEY_SPACE,	 # Not used
  }

# Setup Variables and Constants
bus = smbus.SMBus(1) # Rev 2 Pi uses 1, Rev 1 Pi uses 0

# Device addresses refer to the A0-A2 Pins of the MCP23017 See the online manual
DEVICE_B = 0x20 # Device address (A0-A2 ALL GND)
DEVICE_A = 0x21 # Device B address (A0-A1 GND A2 VCC)

# The addresses below are all defined by the MCP23017 for more information see the online manual
# http://ww1.microchip.com/downloads/en/DeviceDoc/21952b.pdf
IODIRA   = 0x00 # Pin direction register
IODIRB	 = 0x01
IOPOLA   = 0x02 # Polarity
IOPOLB   = 0x03
GPINTENA = 0x04 # Interrupt on Change Manager
GPINTENB = 0x05
DEFVALA  = 0x06 # Default Value of the Register
DEFVALB  = 0x07
INTCONA  = 0x08 # Interrupt Control Register
INTCONB  = 0x09
IOCON    = 0x0A # Expander Configuration Bank, Mirrow, Sequential, Slew State, Hardware Address, Open Drain Int, Int Active High or Low
#ICON    = 0x0B # Not required
GPPUA    = 0x0C # Pull Up Register
GPPUB    = 0x0D
INTFA    = 0x0E # Flags the pin that caused the interrupt
INTFB    = 0x0F
INTCAPA  = 0x10 # The value of the pin at the time of the interrupt
INTCAPB  = 0x11
GPIOA    = 0x12 # Register for inputs
GPIOB    = 0x13
OLATA    = 0x14 #
OLATB    = 0x15

# Raspberry PI Pins will be used for interrupts
INT_PIN_B = 17
INT_PIN_A = 18

INT_PIN_D = 23
INT_PIN_C = 22

PIN_SHUTDOWN = 24 # We dont really need this, but I want to use this later to shut down the pi

DOWN_KEY  = 1 # Constant representing a push key value - required by uinput
UP_KEY    = 0 # Constant representing a release key value - required by uinput

# Initialize PI Pins
GPIO.setmode(GPIO.BCM)
# We are not using resistors externally so setup the internal pull up resistors for the Pi GPIO Pins, this means the button will pull it down from 3.3V to 0
GPIO.setup(INT_PIN_A, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(INT_PIN_B, GPIO.IN, pull_up_down=GPIO.PUD_UP)

GPIO.setup(INT_PIN_C, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(INT_PIN_D, GPIO.IN, pull_up_down=GPIO.PUD_UP)

# We could have done the same above, just put this in to show the difference, here the base is 0 and it will be pulled up to 3.3V is the button is pressed.
GPIO.setup(PIN_SHUTDOWN, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)

# Get a handle to the keyboard Python uinput, defining all the keys we will be using
keyboard_device = uinput.Device([uinput.KEY_5,uinput.KEY_6, uinput.KEY_1,
    uinput.KEY_2, uinput.KEY_ESC, uinput.KEY_P,
    uinput.KEY_ENTER, uinput.KEY_TAB, uinput.KEY_UP, uinput.KEY_DOWN,
    uinput.KEY_LEFT, uinput.KEY_RIGHT, uinput.KEY_R, uinput.KEY_F,
    uinput.KEY_D, uinput.KEY_G, uinput.KEY_LEFTCTRL, uinput.KEY_LEFTALT,
    uinput.KEY_SPACE, uinput.KEY_LEFTSHIFT, uinput.KEY_Z, uinput.KEY_X,
    uinput.KEY_A, uinput.KEY_S, uinput.KEY_Q, uinput.KEY_W, uinput.KEY_E, uinput.KEY_C])

# Method to Setup each of the MCP23017 Devices
def setup_mcp(device_address):
  #Bank 0 Configuration with Interrupt Mirror Disabled, Active Int with Active Low
  bus.write_byte_data(device_address, IOCON, 0b00100000)
  # Turn on the PULL UP Register for the first line GPA0
  bus.write_byte_data(device_address, GPPUA, 0xFF)
  bus.write_byte_data(device_address, GPPUB, 0xFF)
  # Reverse the polarity
  bus.write_byte_data(device_address, IOPOLA, 0xFF)
  bus.write_byte_data(device_address, IOPOLB, 0xFF)
  # All pins as INPUT
  bus.write_byte_data(device_address, IODIRA, 0xFF)
  bus.write_byte_data(device_address, IODIRB, 0xFF)
  # Turn on Interrupts for the pin
  bus.write_byte_data(device_address, GPINTENA, 0xFF)
  bus.write_byte_data(device_address, GPINTENB, 0xFF)
  # Set the Int on Change to 0 so that the previous value is compared
  bus.write_byte_data(device_address, INTCONA, 0x00)
  bus.write_byte_data(device_address, INTCONB, 0x00)

#Translates the device, the control line and the button into a string for lookup
#def getCode(device, intf, switch):
#   if (device == DEVICE_A):
#      code = "MCPA_"
#   else:
#      code = "MCPB_"
#   if (intf == INTFA):
#      code = code + "A_"
#   else:
#      code = code + "B_"
#
#   if (switch & 0x01):
#      code = code + "7"
#   if (switch & 0x02):
#      code = code + "6"
#   if (switch & 0x04):
#      code = code + "5"
#   if (switch & 0x08):
#      code = code + "4"
#   if (switch & 0x10):
#      code = code + "3"
#   if (switch & 0x20):
#      code = code + "2"
#   if (switch & 0x40):
#      code = code + "1"
#   if (switch & 0x80):
#      code = code + "0"
#   return code

def setKeys(device, intf, buttonPressed, buttonState):
  if (device == DEVICE_A):
    code = "MCPA_"
  else:
    code = "MCPB_"
  if (intf == INTFA):
    code = code + "A_"
  else:
    code = code + "B_"

  for x in range(0,7):
    y = 0x80 >> x
    if (buttonPressed & y):
       keyboard_device.emit(DEFINED_KEYS[code + str(x)], buttonState & y)

#Findout which button was pressed and trigger the input
def checkButton(device, intf, intcap):
  buttonPressed = bus.read_byte_data(device, intf)
  buttonValue = bus.read_byte_data(device, intcap)

  setKeys(device, intf, buttonPressed, buttonValue)

  #if (buttonPressed > 0) :
  #  upDownValue = UP_KEY
  #  if (buttonValue & buttonPressed > 0):
  #    upDownValue = DOWN_KEY
  #  keyboard_device.emit(DEFINED_KEYS[getCode(device, intf, buttonPressed)], upDownValue)

#Event Callbacks
def eventOnPinB(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_B, INTFA, INTCAPA)
def eventOnPinD(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_B, INTFB, INTCAPB)
#Event Callback for Pin A
def eventOnPinA(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_A, INTFA, INTCAPA)
def eventOnPinC(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_A, INTFB, INTCAPB)

# Initialise the devices
setup_mcp(DEVICE_A)
setup_mcp(DEVICE_B)

# Clear any data on the interrupt lines
bus.read_byte_data(DEVICE_A, INTCAPA)
bus.read_byte_data(DEVICE_A, INTCAPB)
bus.read_byte_data(DEVICE_B, INTCAPA)
bus.read_byte_data(DEVICE_B, INTCAPB)

# Add event detection
# Not that the bounce time is set to 0 and we handle the time in the callback
# The reason is that if an even happens within the time, it will be ignored
# But the MCP23017 will flag the interrupt even though the GPIO Event Detect will ignore
# the event. No new event will ever be processed because the event INTs must be cleared
# hence making it 0, means we will always read it, but we wait in the call back just incase
GPIO.add_event_detect(INT_PIN_A, GPIO.BOTH, callback=eventOnPinA, bouncetime=0)
GPIO.add_event_detect(INT_PIN_B, GPIO.BOTH, callback=eventOnPinB, bouncetime=0)
GPIO.add_event_detect(INT_PIN_C, GPIO.BOTH, callback=eventOnPinC, bouncetime=0)
GPIO.add_event_detect(INT_PIN_D, GPIO.BOTH, callback=eventOnPinD, bouncetime=0)

try:
  GPIO.wait_for_edge(PIN_SHUTDOWN, GPIO.RISING) # Wait for shutdown on pin 24
#  while 1:
#    time.sleep(0.02)
except KeyboardInterrupt:
  GPIO.cleanup()

GPIO.cleanup()

print "Done"

<commit_msg>Fixed Controls<commit_after>#! /usr/bin/env python

#   Copyright 2014 Jason F Nicholls
#
#   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.

#   This Python Script is intended to be used in conjunction with PiMame
#   More specifically Mame4All
#   However you can use it for what ever purpose you see a need
#	Basically it reads the GPIO board of the Pi, and waits for
#   Events, if these events occur it reads 2 MCP23017 ICs for input
#   You can extend the Pi and use 8 ICs just extend this code
#   Note that ALL the pins of the ICs are setup as output.
#
#   For more information see picade.blogspot.com
#
#   Hope you have as much fun with your Pi as I have had with mine.

import RPi.GPIO as GPIO
import uinput
import smbus
import time

# Custom Key Configuration, these keys will be mapped to your controls, below are the default MAME Controls.
DEFINED_KEYS = {
    'MCPA_A_0': uinput.KEY_5,      	 # Coin 1
    'MCPA_A_1': uinput.KEY_6,   	 # Coin 2
    'MCPA_A_2': uinput.KEY_1,  		 # Start 1
    'MCPA_A_3': uinput.KEY_2,  		 # Start 2
    'MCPA_A_4': uinput.KEY_ESC, 	 # Escape - exit the game
    'MCPA_A_5': uinput.KEY_P,  		 # Pause
    'MCPA_A_6': uinput.KEY_ENTER, 	 # Enter
    'MCPA_A_7': uinput.KEY_TAB, 	 # Tab - Access in game menu
    'MCPA_B_0': uinput.KEY_G,  	 	 # Player 2 Right
    'MCPA_B_1': uinput.KEY_D, 	 	 # Player 2 Left
    'MCPA_B_2': uinput.KEY_F,	 	 # Player 2 Down
    'MCPA_B_3': uinput.KEY_R, 	 	 # Player 2 UP
    'MCPA_B_4': uinput.KEY_RIGHT,	 # Player 1 Right
    'MCPA_B_5': uinput.KEY_LEFT,	 # Player 1 Left
    'MCPA_B_6': uinput.KEY_DOWN,	 # Player 1 Down
    'MCPA_B_7': uinput.KEY_UP,		 # Player 1 Up
    'MCPB_A_0': uinput.KEY_LEFTCTRL,     # Player 1 Fire 1
    'MCPB_A_1': uinput.KEY_LEFTALT,	 # Player 1 Fire 2
    'MCPB_A_2': uinput.KEY_SPACE,	 # Player 1 Fire 3
    'MCPB_A_3': uinput.KEY_LEFTSHIFT,    # Player 1 Fire 4
    'MCPB_A_4': uinput.KEY_Z,	 	 # Player 1 Fire 5
    'MCPB_A_5': uinput.KEY_X,	         # Player 1 Fire 6
    'MCPB_A_6': uinput.KEY_SPACE,	 # Not used
    'MCPB_A_7': uinput.KEY_SPACE,        # Not used
    'MCPB_B_0': uinput.KEY_SPACE,	 # Not Used
    'MCPB_B_1': uinput.KEY_SPACE,	 # Not Used
    'MCPB_B_2': uinput.KEY_C,		 # C Player 2 Fire 6 // Normally Joystick 4 and is not set by default
    'MCPB_B_3': uinput.KEY_E,		 # E Player 2 Fire 5 // Normally Joystick 5 and is not set by default
    'MCPB_B_4': uinput.KEY_W,		 # W Player 2 Fire 4 // Normally Joystick 4 and is not set by default
    'MCPB_B_5': uinput.KEY_Q,		 # Q Player 2 Fire 3 // Normally Joystick 5 and is not set by default
    'MCPB_B_6': uinput.KEY_S,	 	 # S Player 2 Fire 2
    'MCPB_B_7': uinput.KEY_A,	 	 # A Player 2 Fire 1
  }

# Setup Variables and Constants
bus = smbus.SMBus(1) # Rev 2 Pi uses 1, Rev 1 Pi uses 0

# Device addresses refer to the A0-A2 Pins of the MCP23017 See the online manual
DEVICE_B = 0x20 # Device address (A0-A2 ALL GND)
DEVICE_A = 0x21 # Device B address (A0-A1 GND A2 VCC)

# The addresses below are all defined by the MCP23017 for more information see the online manual
# http://ww1.microchip.com/downloads/en/DeviceDoc/21952b.pdf
IODIRA   = 0x00 # Pin direction register
IODIRB	 = 0x01
IOPOLA   = 0x02 # Polarity
IOPOLB   = 0x03
GPINTENA = 0x04 # Interrupt on Change Manager
GPINTENB = 0x05
DEFVALA  = 0x06 # Default Value of the Register
DEFVALB  = 0x07
INTCONA  = 0x08 # Interrupt Control Register
INTCONB  = 0x09
IOCON    = 0x0A # Expander Configuration Bank, Mirrow, Sequential, Slew State, Hardware Address, Open Drain Int, Int Active High or Low
#ICON    = 0x0B # Not required
GPPUA    = 0x0C # Pull Up Register
GPPUB    = 0x0D
INTFA    = 0x0E # Flags the pin that caused the interrupt
INTFB    = 0x0F
INTCAPA  = 0x10 # The value of the pin at the time of the interrupt
INTCAPB  = 0x11
GPIOA    = 0x12 # Register for inputs
GPIOB    = 0x13
OLATA    = 0x14 #
OLATB    = 0x15

# Raspberry PI Pins will be used for interrupts
INT_PIN_B = 17
INT_PIN_A = 18

INT_PIN_D = 23
INT_PIN_C = 22

PIN_SHUTDOWN = 24 # We dont really need this, but I want to use this later to shut down the pi

DOWN_KEY  = 1 # Constant representing a push key value - required by uinput
UP_KEY    = 0 # Constant representing a release key value - required by uinput

# Initialize PI Pins
GPIO.setmode(GPIO.BCM)
# We are not using resistors externally so setup the internal pull up resistors for the Pi GPIO Pins, this means the button will pull it down from 3.3V to 0
GPIO.setup(INT_PIN_A, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(INT_PIN_B, GPIO.IN, pull_up_down=GPIO.PUD_UP)

GPIO.setup(INT_PIN_C, GPIO.IN, pull_up_down=GPIO.PUD_UP)
GPIO.setup(INT_PIN_D, GPIO.IN, pull_up_down=GPIO.PUD_UP)

# We could have done the same above, just put this in to show the difference, here the base is 0 and it will be pulled up to 3.3V is the button is pressed.
GPIO.setup(PIN_SHUTDOWN, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)

# Get a handle to the keyboard Python uinput, defining all the keys we will be using
keyboard_device = uinput.Device([uinput.KEY_5,uinput.KEY_6, uinput.KEY_1,
    uinput.KEY_2, uinput.KEY_ESC, uinput.KEY_P,
    uinput.KEY_ENTER, uinput.KEY_TAB, uinput.KEY_UP, uinput.KEY_DOWN,
    uinput.KEY_LEFT, uinput.KEY_RIGHT, uinput.KEY_R, uinput.KEY_F,
    uinput.KEY_D, uinput.KEY_G, uinput.KEY_LEFTCTRL, uinput.KEY_LEFTALT,
    uinput.KEY_SPACE, uinput.KEY_LEFTSHIFT, uinput.KEY_Z, uinput.KEY_X,
    uinput.KEY_A, uinput.KEY_S, uinput.KEY_Q, uinput.KEY_W, uinput.KEY_E, uinput.KEY_C])

# Method to Setup each of the MCP23017 Devices
def setup_mcp(device_address):
  #Bank 0 Configuration with Interrupt Mirror Disabled, Active Int with Active Low
  bus.write_byte_data(device_address, IOCON, 0b00100000)
  # Turn on the PULL UP Register for the first line GPA0
  bus.write_byte_data(device_address, GPPUA, 0xFF)
  bus.write_byte_data(device_address, GPPUB, 0xFF)
  # Reverse the polarity
  bus.write_byte_data(device_address, IOPOLA, 0xFF)
  bus.write_byte_data(device_address, IOPOLB, 0xFF)
  # All pins as INPUT
  bus.write_byte_data(device_address, IODIRA, 0xFF)
  bus.write_byte_data(device_address, IODIRB, 0xFF)
  # Turn on Interrupts for the pin
  bus.write_byte_data(device_address, GPINTENA, 0xFF)
  bus.write_byte_data(device_address, GPINTENB, 0xFF)
  # Set the Int on Change to 0 so that the previous value is compared
  bus.write_byte_data(device_address, INTCONA, 0x00)
  bus.write_byte_data(device_address, INTCONB, 0x00)

#Translates the device, the control line and the button into a string for lookup
#def getCode(device, intf, switch):
#   if (device == DEVICE_A):
#      code = "MCPA_"
#   else:
#      code = "MCPB_"
#   if (intf == INTFA):
#      code = code + "A_"
#   else:
#      code = code + "B_"
#
#   if (switch & 0x01):
#      code = code + "7"
#   if (switch & 0x02):
#      code = code + "6"
#   if (switch & 0x04):
#      code = code + "5"
#   if (switch & 0x08):
#      code = code + "4"
#   if (switch & 0x10):
#      code = code + "3"
#   if (switch & 0x20):
#      code = code + "2"
#   if (switch & 0x40):
#      code = code + "1"
#   if (switch & 0x80):
#      code = code + "0"
#   return code

def setKeys(device, intf, buttonPressed, buttonState):
  if (device == DEVICE_A):
    code = "MCPA_"
  else:
    code = "MCPB_"
  if (intf == INTFA):
    code = code + "A_"
  else:
    code = code + "B_"

  for x in range(0,7):
    y = 0x80 >> x
    if (buttonPressed & y):
       keyboard_device.emit(DEFINED_KEYS[code + str(x)], buttonState & y)

#Findout which button was pressed and trigger the input
def checkButton(device, intf, intcap):
  buttonPressed = bus.read_byte_data(device, intf)
  buttonValue = bus.read_byte_data(device, intcap)

  setKeys(device, intf, buttonPressed, buttonValue)

  #if (buttonPressed > 0) :
  #  upDownValue = UP_KEY
  #  if (buttonValue & buttonPressed > 0):
  #    upDownValue = DOWN_KEY
  #  keyboard_device.emit(DEFINED_KEYS[getCode(device, intf, buttonPressed)], upDownValue)

#Event Callbacks
def eventOnPinB(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_B, INTFA, INTCAPA)
def eventOnPinD(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_B, INTFB, INTCAPB)
#Event Callback for Pin A
def eventOnPinA(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_A, INTFA, INTCAPA)
def eventOnPinC(channel):
  time.sleep(0.03) #Sleep to cancel out bounce
  checkButton(DEVICE_A, INTFB, INTCAPB)

# Initialise the devices
setup_mcp(DEVICE_A)
setup_mcp(DEVICE_B)

# Clear any data on the interrupt lines
bus.read_byte_data(DEVICE_A, INTCAPA)
bus.read_byte_data(DEVICE_A, INTCAPB)
bus.read_byte_data(DEVICE_B, INTCAPA)
bus.read_byte_data(DEVICE_B, INTCAPB)

# Add event detection
# Not that the bounce time is set to 0 and we handle the time in the callback
# The reason is that if an even happens within the time, it will be ignored
# But the MCP23017 will flag the interrupt even though the GPIO Event Detect will ignore
# the event. No new event will ever be processed because the event INTs must be cleared
# hence making it 0, means we will always read it, but we wait in the call back just incase
GPIO.add_event_detect(INT_PIN_A, GPIO.BOTH, callback=eventOnPinA, bouncetime=0)
GPIO.add_event_detect(INT_PIN_B, GPIO.BOTH, callback=eventOnPinB, bouncetime=0)
GPIO.add_event_detect(INT_PIN_C, GPIO.BOTH, callback=eventOnPinC, bouncetime=0)
GPIO.add_event_detect(INT_PIN_D, GPIO.BOTH, callback=eventOnPinD, bouncetime=0)

try:
  GPIO.wait_for_edge(PIN_SHUTDOWN, GPIO.RISING) # Wait for shutdown on pin 24
#  while 1:
#    time.sleep(0.02)
except KeyboardInterrupt:
  GPIO.cleanup()

GPIO.cleanup()

print "Done"

<|endoftext|>


<commit_before>"""A collection of utility functions."""

import click

_LOG_LEVEL = 0
LOG_PREFIXES = {
    0: click.style('>', fg='white', bold=True),
    1: click.style('>>', fg='cyan', bold=True),
    2: click.style('>>>', fg='blue', bold=True),
}

def set_log_level(log_level):
    """Set the log level to use.

    Log levels are provided as an integer, and the following levels are available:
    0: The default log level
    1: Info level logging
    2: Debug level logging
    """
    if log_level not in [0, 1, 2]:
        return False

    global _LOG_LEVEL  # pylint: disable = global-statement
    _LOG_LEVEL = log_level
    return True

def get_log_level():
    """Get the current log level."""
    return _LOG_LEVEL

def _print(message, log_level, prefix, **kwargs):
    """Print a message to the console, depending on log level."""
    if log_level <= get_log_level():
        if prefix:
            click.secho('{} '.format(LOG_PREFIXES[log_level]), nl=False)

        return click.secho(message, **kwargs)

def log(message, prefix=True, **kwargs):
    """Print a default level log message."""
    return _print(message, log_level=0, prefix=prefix, **kwargs)

def info(message, prefix=True, **kwargs):
    """Print an info level log message."""
    return _print(message, log_level=1, prefix=prefix, **kwargs)

def debug(message, prefix=True, **kwargs):
    """Print a debug level log message."""
    return _print(message, log_level=2, prefix=prefix, **kwargs)
<commit_msg>Added caching of a central Docker client to the utils module<commit_after>"""A collection of utility functions."""

import click
from docker.client import Client
from docker.utils import kwargs_from_env

_LOG_LEVEL = 0
LOG_PREFIXES = {
    0: click.style('>', fg='white', bold=True),
    1: click.style('>>', fg='cyan', bold=True),
    2: click.style('>>>', fg='blue', bold=True),
}
_DOCKER_CLIENT = None

def set_log_level(log_level):
    """Set the log level to use.

    Log levels are provided as an integer, and the following levels are available:
    0: The default log level
    1: Info level logging
    2: Debug level logging
    """
    if log_level not in [0, 1, 2]:
        return False

    global _LOG_LEVEL  # pylint: disable = global-statement
    _LOG_LEVEL = log_level
    return True

def get_log_level():
    """Get the current log level."""
    return _LOG_LEVEL

def _print(message, log_level, prefix, **kwargs):
    """Print a message to the console, depending on log level."""
    if log_level <= get_log_level():
        if prefix:
            click.secho('{} '.format(LOG_PREFIXES[log_level]), nl=False)

        return click.secho(message, **kwargs)

def log(message, prefix=True, **kwargs):
    """Print a default level log message."""
    return _print(message, log_level=0, prefix=prefix, **kwargs)

def info(message, prefix=True, **kwargs):
    """Print an info level log message."""
    return _print(message, log_level=1, prefix=prefix, **kwargs)

def debug(message, prefix=True, **kwargs):
    """Print a debug level log message."""
    return _print(message, log_level=2, prefix=prefix, **kwargs)

def docker_client():
    """Create a Docker client instance or return an existing one."""
    global _DOCKER_CLIENT  # pylint: disable = global-statement

    if _DOCKER_CLIENT:
        debug('Reusing existing Docker client.')
        return _DOCKER_CLIENT
    else:
        debug('Creating Docker client.')
        _DOCKER_CLIENT = Client(**kwargs_from_env())
        return _DOCKER_CLIENT
<|endoftext|>