<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    ThinkOpen Solutions Brasil
#    Copyright (C) Thinkopen Solutions <http://www.tkobr.com>.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.osv import osv, fields
from openerp import api
from openerp import SUPERUSER_ID
import logging
_logger = logging.getLogger(__name__)

class res_partner(osv.osv):
    _name = "res.partner"
    _inherit = "res.partner"
    _description = "inherited partner class"

    def create(self, cr, uid, vals, context=False):
        mail_obj = self.pool.get('res.partner.email')
        res = super(res_partner, self).create(cr, uid, vals, context=context)
        email_ids = mail_obj.search(cr, uid, [('res_partner_id', '=', res)])
        if email_ids:
            mail_obj.write(cr, uid, email_ids[0], {'is_active': True})
            vals.update({'email': mail_obj.browse(
                cr, uid, email_ids[0]).email})
        return res

    def write(self, cr, uid, ids, vals, context=None):
        res = super(
            res_partner,
            self).write(
            cr,
            uid,
            ids,
            vals,
            context=context)
        email_ids = []
        if isinstance(ids, int):
            ids = [ids]
        mail_obj = self.pool.get('res.partner.email')
        if len(ids):
            email_ids = mail_obj.search(
                cr, uid, [('res_partner_id', '=', ids[0])])
            if len(email_ids) == 1:
                mail_obj.write(cr, uid, email_ids[0], {'is_active': True})
        return res

    def _get_email_id(self, cr, uid, ids, name, args, context=False):
        res = {}
        email_obj = self.pool.get('res.partner.email')
        for record in self.browse(cr, uid, ids):
            email_ids = email_obj.search(
                cr, uid, [
                    ('res_partner_id', '=', record.id), ('is_active', '=', True)])
            if email_ids:
                email = email_obj.browse(cr, uid, email_ids[0]).email
                res[record.id] = email
            else:
                res[record.id] = False
        return res

    def _set_email_id(self, cr, uid, ids, name, value, args, context=None):
        mail_obj = self.pool.get('res.partner.email')
        if value and ids:
            for record in self.browse(cr, uid, ids):
                mail_ids = mail_obj.search(
                    cr, uid, [
                        ('res_partner_id', '=', record.id), ('email', 'ilike', value)])
                if mail_ids:
                    previous_mail_ids = mail_obj.search(
                        cr, uid, [('res_partner_id', '=', record.id)])
                    mail_obj.write(
                        cr, uid, previous_mail_ids, {
                            'is_active': False})
                    mail_obj.write(cr, uid, mail_ids, {'is_active': True})
                if not mail_ids:
                    active_mail_ids = mail_obj.search(
                        cr, uid, [('res_partner_id', '=', record.id), ('is_active', '=', True)])
                    if len(active_mail_ids):
                        mail_obj.write(
                            cr, uid, active_mail_ids, {
                                'email': value})
                    else:
                        # mail_obj.write(cr, uid, previous_mail_ids, {'is_active' : False})
                        mail_obj.create(cr, uid, {'res_partner_id': record.id,
                                                  'email': value,
                                                  'is_active': True})
        if not value:
            for record in self.browse(cr, uid, ids):
                mail_id = mail_obj.search(
                    cr, uid, [
                        ('res_partner_id', '=', record.id), ('is_active', '=', True)])
                if len(mail_id):
                    mail_obj.unlink(cr, uid, mail_id)
        return True

    def _get_partner(self, cr, uid, ids, context=None):
        result = {}
        for part in self.pool.get('res.partner.email').browse(
                cr, uid, ids, context=context):
            result[part.res_partner_id.id] = True
        return result.keys()

    def _get_mail_ids(self, cr, uid, ids, fields, arg, context=None):
        res = {}
        for mail in ids:
            res[mail] = self.pool.get('res.partner.email').search(
                cr, uid, [('res_partner_id', '=', mail)])
        return res

    _columns = {
        'email': fields.function(
            _get_email_id, type='char', fnct_inv=_set_email_id, string='Email', store={
                'res.partner': (
                    lambda self, cr, uid, ids, c={}: ids, ['email_ids'], 10), 'res.partner.email': (
                    _get_partner, [
                        'email', 'is_active'], 10), }), 'email_ids': fields.one2many(
                            'res.partner.email', 'res_partner_id', 'Emails'), 'email_ids_readonly': fields.function(
                                _get_mail_ids, type='one2many', relation='res.partner.email', string='Emails')}

    def _create_multiple_emails_at_first_install(
            self, cr, SUPERUSER_ID, ids=None, context=None):
        partner_ids = self.search(cr, SUPERUSER_ID, [('email', '!=', False)])
        part_email_obj = self.pool.get('res.partner.email')
        for partner_id in partner_ids:
            _logger.info(
                "Setting up multiple email '%s' record for client %s ." %
                (partner.email, partner.name))
            partner = self.browse(cr, SUPERUSER_ID, partner_id)
            part_email_obj.create(cr, SUPERUSER_ID, {
                'email': partner.email,
                'is_active': True,
                'res_partner_id': partner_id,
            })
        return True
<commit_msg>place creation of multi email at first install inside a try<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    ThinkOpen Solutions Brasil
#    Copyright (C) Thinkopen Solutions <http://www.tkobr.com>.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.osv import osv, fields
from openerp import api
from openerp import SUPERUSER_ID
import logging
_logger = logging.getLogger(__name__)

class res_partner(osv.osv):
    _name = "res.partner"
    _inherit = "res.partner"
    _description = "inherited partner class"

    def create(self, cr, uid, vals, context=False):
        mail_obj = self.pool.get('res.partner.email')
        res = super(res_partner, self).create(cr, uid, vals, context=context)
        email_ids = mail_obj.search(cr, uid, [('res_partner_id', '=', res)])
        if email_ids:
            mail_obj.write(cr, uid, email_ids[0], {'is_active': True})
            vals.update({'email': mail_obj.browse(
                cr, uid, email_ids[0]).email})
        return res

    def write(self, cr, uid, ids, vals, context=None):
        res = super(
            res_partner,
            self).write(
            cr,
            uid,
            ids,
            vals,
            context=context)
        email_ids = []
        if isinstance(ids, int):
            ids = [ids]
        mail_obj = self.pool.get('res.partner.email')
        if len(ids):
            email_ids = mail_obj.search(
                cr, uid, [('res_partner_id', '=', ids[0])])
            if len(email_ids) == 1:
                mail_obj.write(cr, uid, email_ids[0], {'is_active': True})
        return res

    def _get_email_id(self, cr, uid, ids, name, args, context=False):
        res = {}
        email_obj = self.pool.get('res.partner.email')
        for record in self.browse(cr, uid, ids):
            email_ids = email_obj.search(
                cr, uid, [
                    ('res_partner_id', '=', record.id), ('is_active', '=', True)])
            if email_ids:
                email = email_obj.browse(cr, uid, email_ids[0]).email
                res[record.id] = email
            else:
                res[record.id] = False
        return res

    def _set_email_id(self, cr, uid, ids, name, value, args, context=None):
        mail_obj = self.pool.get('res.partner.email')
        if value and ids:
            for record in self.browse(cr, uid, ids):
                mail_ids = mail_obj.search(
                    cr, uid, [
                        ('res_partner_id', '=', record.id), ('email', 'ilike', value)])
                if mail_ids:
                    previous_mail_ids = mail_obj.search(
                        cr, uid, [('res_partner_id', '=', record.id)])
                    mail_obj.write(
                        cr, uid, previous_mail_ids, {
                            'is_active': False})
                    mail_obj.write(cr, uid, mail_ids, {'is_active': True})
                if not mail_ids:
                    active_mail_ids = mail_obj.search(
                        cr, uid, [('res_partner_id', '=', record.id), ('is_active', '=', True)])
                    if len(active_mail_ids):
                        mail_obj.write(
                            cr, uid, active_mail_ids, {
                                'email': value})
                    else:
                        # mail_obj.write(cr, uid, previous_mail_ids, {'is_active' : False})
                        mail_obj.create(cr, uid, {'res_partner_id': record.id,
                                                  'email': value,
                                                  'is_active': True})
        if not value:
            for record in self.browse(cr, uid, ids):
                mail_id = mail_obj.search(
                    cr, uid, [
                        ('res_partner_id', '=', record.id), ('is_active', '=', True)])
                if len(mail_id):
                    mail_obj.unlink(cr, uid, mail_id)
        return True

    def _get_partner(self, cr, uid, ids, context=None):
        result = {}
        for part in self.pool.get('res.partner.email').browse(
                cr, uid, ids, context=context):
            result[part.res_partner_id.id] = True
        return result.keys()

    def _get_mail_ids(self, cr, uid, ids, fields, arg, context=None):
        res = {}
        for mail in ids:
            res[mail] = self.pool.get('res.partner.email').search(
                cr, uid, [('res_partner_id', '=', mail)])
        return res

    _columns = {
        'email': fields.function(
            _get_email_id, type='char', fnct_inv=_set_email_id, string='Email', store={
                'res.partner': (
                    lambda self, cr, uid, ids, c={}: ids, ['email_ids'], 10), 'res.partner.email': (
                    _get_partner, [
                        'email', 'is_active'], 10), }), 'email_ids': fields.one2many(
                            'res.partner.email', 'res_partner_id', 'Emails'), 'email_ids_readonly': fields.function(
                                _get_mail_ids, type='one2many', relation='res.partner.email', string='Emails')}

    def _create_multiple_emails_at_first_install(
            self, cr, SUPERUSER_ID, ids=None, context=None):
        partner_ids = self.search(cr, SUPERUSER_ID, [('email', '!=', False)])
        part_email_obj = self.pool.get('res.partner.email')
        for partner_id in partner_ids:
            _logger.info(
                "Setting up multiple email '%s' record for client %s ." %
                (partner.email, partner.name))
            partner = self.browse(cr, SUPERUSER_ID, partner_id)
            try:
                part_email_obj.create(cr, SUPERUSER_ID, {
                    'email': partner.email,
                    'is_active': True,
                    'res_partner_id': partner_id,
                })
            except:
                _logger.warning(
                    "Unable to setup multiple email '%s' record for client %s ." %
                    (partner.email, partner.name))
        return True
<|endoftext|>


<commit_before>import asyncio
import logging
import time
import os
import base64
from collections import deque
from autobahn.asyncio.websocket import WebSocketClientProtocol, WebSocketClientFactory
from wstan.relay import RelayMixin
from wstan import parse_socks_addr, loop, config, can_return_error_page, gen_error_page

# noinspection PyAttributeOutsideInit
class CustomWSClientProtocol(WebSocketClientProtocol):
    """Add auto-ping switch (dirty way) and let us manually start handshaking."""
    # this framework mix camel and underline naming style, nice!
    def __init__(self):
        super().__init__()
        self.customUriPath = '/'
        self.delayedHandshake = asyncio.Future()

    def enableAutoPing(self, interval):
        try:
            self.autoPingInterval = interval
            self.autoPingPendingCall = loop.call_later(interval, self._sendAutoPing)
        except AttributeError:
            logging.warning('failed to enable auto-ping, maybe library changed its internal method')

    def disableAutoPing(self):
        try:
            self.autoPingInterval = 0
            if self.autoPingPendingCall:
                self.autoPingPendingCall.cancel()
                self.autoPingPendingCall = None
        except AttributeError:
            logging.warning('failed to disable auto-ping')

    def startHandshake(self):
        """Delay handshake because some states must be set before handshake (so
        they can't be set in factory)."""
        self.delayedHandshake.set_result(None)

    def restartHandshake(self):
        """Customize handshake HTTP header."""
        asyncio.wait_for(self.delayedHandshake, None)
        assert self.websocket_key
        request = [
            'GET %s HTTP/1.1' % self.customUriPath,
            'Host: %s:%d' % (self.factory.host, self.factory.port),
            'User-Agent: wwww',
            'Sec-WebSocket-Key: %s' % self.websocket_key.decode(),
            'Sec-WebSocket-Version: %d' % self.SPEC_TO_PROTOCOL_VERSION[self.version],
            'Pragma: no-cache',
            'Cache-Control: no-cache',
            'Connection: Upgrade',
            'Upgrade: WebSocket',
            '',
            ''  # ends with \r\n\r\n
        ]
        self.http_request_data = '\r\n'.join(request).encode('utf8')
        self.sendData(self.http_request_data)
        if self.debug:
            self.log.debug(request)

class WSTunClientProtocol(CustomWSClientProtocol, RelayMixin):
    POOL_MAX_SIZE = 16
    TUN_MAX_IDLE_TIMEOUT = 35  # in seconds. close tunnel on timeout
    TUN_PING_INTERVAL = 5  # only tunnels in pool do auto-ping
    TUN_OPEN_TIMEOUT = 5  # time to wait after TCP established and before succeeded WS handshake
    POOL_NOM_SIZE, TUN_MIN_IDLE_TIMEOUT = round(POOL_MAX_SIZE / 2), round(TUN_MAX_IDLE_TIMEOUT / 2)
    pool = deque()

    def __init__(self):
        super().__init__()
        self.lastIdleTime = None
        self.checkTimeoutTask = None
        self.tunOpen = asyncio.Future()
        self.inPool = False
        self.canReturnErrorPage = False
        nonce = os.urandom(16)
        if not config.tun_ssl:
            self.initCipher(nonce, encryptor=True)
        self.websocket_key = base64.b64encode(nonce)

    if TUN_MAX_IDLE_TIMEOUT <= 0:
        def resetTunnel(self, reason=''):
            self.sendClose(1000)

        def onResetTunnel(self):
            self.sendClose(1000)

    def succeedReset(self):
        super().succeedReset()
        self.lastIdleTime = time.time()
        WSTunClientProtocol.addToPool(self)

    def onOpen(self):
        self.tunOpen.set_result(None)
        self.lastIdleTime = time.time()
        if not config.debug:
            self.customUriPath = None  # save memory
        if not config.tun_ssl:
            # SHA-1 has 20 bytes
            self.initCipher(base64.b64decode(self.http_headers['sec-websocket-accept'])[:16],
                            decryptor=True)

    def onMessage(self, dat, isBinary):
        if not isBinary:
            logging.error('non binary ws message received')
            return self.sendClose(3000)

        cmd = ord(self.decrypt(dat[:1]))
        if cmd == self.CMD_RST:
            msg = self.parseResetMessage(dat)
            if not msg.startswith('  '):
                logging.info('tunnel abnormal reset: %s' % msg)
                if self.canReturnErrorPage:
                    self._writer.write(gen_error_page("can't connect to destination", msg))
            self.onResetTunnel()
        elif cmd == self.CMD_DAT:
            dat = self.decrypt(dat[1:])
            if self.tunState != self.TUN_STATE_USING:
                # why this happens?
                return
            self.canReturnErrorPage = False
            self._writer.write(dat)
        else:
            logging.error('wrong command')

    def onClose(self, *args, **kwargs):
        super().onClose(*args, **kwargs)
        if self.inPool:
            self.pool.remove(self)

    @classmethod
    @asyncio.coroutine
    def _checkTimeout(cls, tun):
        while tun.state == cls.STATE_OPEN:
            timeout = (cls.TUN_MAX_IDLE_TIMEOUT if len(cls.pool) <= cls.POOL_NOM_SIZE else
                       cls.TUN_MIN_IDLE_TIMEOUT)
            yield from asyncio.sleep(timeout)
            if (tun.tunState == cls.TUN_STATE_IDLE and
               (time.time() - tun.lastIdleTime) > timeout):
                tun.sendClose(1000)

    @classmethod
    def addToPool(cls, tun):
        assert tun.tunState == cls.TUN_STATE_IDLE
        if len(cls.pool) >= cls.POOL_MAX_SIZE:
            tun.sendClose(1000)
        else:
            assert not tun.checkTimeoutTask
            tun.checkTimeoutTask = asyncio.async(cls._checkTimeout(tun))
            tun.inPool = True
            tun.enableAutoPing(cls.TUN_PING_INTERVAL)
            cls.pool.append(tun)

    @classmethod
    @asyncio.coroutine
    def getOrCreate(cls, addrHeader, dat):
        if WSTunClientProtocol.pool:
            logging.debug('reuse tunnel from pool (total %s)' % len(WSTunClientProtocol.pool))
            tun = WSTunClientProtocol.pool.popleft()
            tun.checkTimeoutTask.cancel()
            tun.checkTimeoutTask = None
            tun.inPool = False
            tun.disableAutoPing()
            tun.sendMessage(tun.makeRelayHeader(addrHeader, dat), True)
        else:
            tun = (yield from loop.create_connection(
                factory, config.uri_addr, config.uri_port, ssl=config.tun_ssl))[1]
            # lower latency by sending relay header and data in ws handshake
            tun.customUriPath = '/' + base64.urlsafe_b64encode(tun.makeRelayHeader(addrHeader, dat)).decode()
            tun.restartHandshake()
            try:
                yield from asyncio.wait_for(tun.tunOpen, cls.TUN_OPEN_TIMEOUT)
            except asyncio.TimeoutError:
                tun.dropConnection()
                raise
        return tun

factory = WebSocketClientFactory(config.uri)
factory.protocol = WSTunClientProtocol
factory.useragent = ''
factory.autoPingTimeout = 5
factory.openHandshakeTimeout = 10

@asyncio.coroutine
def socks5_tcp_handler(reader, writer):
    # handle auth method selection
    dat = yield from reader.read(257)
    if len(dat) < 2 or dat[0] != 0x05 or len(dat) != dat[1] + 2:
        logging.warning('local SOCKS v5 server got unknown request')
        return writer.close()
    writer.write(b'\x05\x00')  # \x00 == NO AUTHENTICATION REQUIRED

    # handle relay request
    try:
        dat = yield from reader.read(262)
    except ConnectionError:
        return writer.close()
    try:
        cmd, addr_header = dat[1], dat[2:]
        target_addr, target_port = parse_socks_addr(addr_header)
    except (ValueError, IndexError):
        logging.warning('invalid SOCKS v5 relay request')
        return writer.close()
    logging.info('requesting %s:%d' % (target_addr, target_port))
    if cmd != 0x01:  # CONNECT
        writer.write(b'\x05\x07\x00\x01' + b'\x00' * 6)  # \x07 == COMMAND NOT SUPPORTED
        return writer.close()

    # By accepting request before connected to target delay can be lowered (of a round-trip).
    # But SOCKS client can't get real reason when error happens (Firefox always display
    # connection reset error). Dirty solution: generate a HTML page when a HTTP request failed
    writer.write(b'\x05\x00\x00\x01' + b'\x01' * 6)  # \x00 == SUCCEEDED
    try:
        dat = yield from reader.read(2048)
    except ConnectionError:
        dat = None
    if not dat:
        return writer.close()

    canErr = can_return_error_page(dat)
    try:
        tun = yield from WSTunClientProtocol.getOrCreate(addr_header, dat)
    except Exception as e:
        if canErr:
            writer.write(gen_error_page("can't connect to wstan server", str(e)))
        return writer.close()
    tun.canReturnErrorPage = canErr
    tun.setProxy(reader, writer)

def main():
    server = loop.run_until_complete(
        asyncio.start_server(socks5_tcp_handler, config.addr, config.port))
    print('wstan client -- SOCKS v5 server listen on %s:%d' % (config.addr, config.port))
    try:
        loop.run_forever()
    finally:
        server.close()
        loop.close()
<commit_msg>Enhance error report<commit_after>import asyncio
import logging
import time
import os
import base64
from collections import deque
from autobahn.asyncio.websocket import WebSocketClientProtocol, WebSocketClientFactory
from wstan.relay import RelayMixin
from wstan import parse_socks_addr, loop, config, can_return_error_page, gen_error_page

# noinspection PyAttributeOutsideInit
class CustomWSClientProtocol(WebSocketClientProtocol):
    """Add auto-ping switch (dirty way) and let us manually start handshaking."""
    # this framework mix camel and underline naming style, nice!
    def __init__(self):
        super().__init__()
        self.customUriPath = '/'
        self.delayedHandshake = asyncio.Future()

    def enableAutoPing(self, interval):
        try:
            self.autoPingInterval = interval
            self.autoPingPendingCall = loop.call_later(interval, self._sendAutoPing)
        except AttributeError:
            logging.warning('failed to enable auto-ping, maybe library changed its internal method')

    def disableAutoPing(self):
        try:
            self.autoPingInterval = 0
            if self.autoPingPendingCall:
                self.autoPingPendingCall.cancel()
                self.autoPingPendingCall = None
        except AttributeError:
            logging.warning('failed to disable auto-ping')

    def startHandshake(self):
        """Delay handshake because some states must be set before handshake (so
        they can't be set in factory)."""
        self.delayedHandshake.set_result(None)

    def restartHandshake(self):
        """Customize handshake HTTP header."""
        asyncio.wait_for(self.delayedHandshake, None)
        assert self.websocket_key
        request = [
            'GET %s HTTP/1.1' % self.customUriPath,
            'Host: %s:%d' % (self.factory.host, self.factory.port),
            'User-Agent: wwww',
            'Sec-WebSocket-Key: %s' % self.websocket_key.decode(),
            'Sec-WebSocket-Version: %d' % self.SPEC_TO_PROTOCOL_VERSION[self.version],
            'Pragma: no-cache',
            'Cache-Control: no-cache',
            'Connection: Upgrade',
            'Upgrade: WebSocket',
            '',
            ''  # ends with \r\n\r\n
        ]
        self.http_request_data = '\r\n'.join(request).encode('utf8')
        self.sendData(self.http_request_data)
        if self.debug:
            self.log.debug(request)

class WSTunClientProtocol(CustomWSClientProtocol, RelayMixin):
    POOL_MAX_SIZE = 16
    TUN_MAX_IDLE_TIMEOUT = 35  # in seconds. close tunnel on timeout
    TUN_PING_INTERVAL = 5  # only tunnels in pool do auto-ping
    TUN_OPEN_TIMEOUT = 5  # time to wait after TCP established and before succeeded WS handshake
    POOL_NOM_SIZE, TUN_MIN_IDLE_TIMEOUT = round(POOL_MAX_SIZE / 2), round(TUN_MAX_IDLE_TIMEOUT / 2)
    pool = deque()

    def __init__(self):
        super().__init__()
        self.lastIdleTime = None
        self.checkTimeoutTask = None
        self.tunOpen = asyncio.Future()
        self.inPool = False
        self.canReturnErrorPage = False
        nonce = os.urandom(16)
        if not config.tun_ssl:
            self.initCipher(nonce, encryptor=True)
        self.websocket_key = base64.b64encode(nonce)

    if TUN_MAX_IDLE_TIMEOUT <= 0:
        def resetTunnel(self, reason=''):
            self.sendClose(1000)

        def onResetTunnel(self):
            self.sendClose(1000)

    def succeedReset(self):
        super().succeedReset()
        self.lastIdleTime = time.time()
        WSTunClientProtocol.addToPool(self)

    def onOpen(self):
        self.tunOpen.set_result(None)
        self.lastIdleTime = time.time()
        if not config.debug:
            self.customUriPath = None  # save memory
        if not config.tun_ssl:
            # SHA-1 has 20 bytes
            self.initCipher(base64.b64decode(self.http_headers['sec-websocket-accept'])[:16],
                            decryptor=True)

    def onMessage(self, dat, isBinary):
        if not isBinary:
            logging.error('non binary ws message received')
            return self.sendClose(3000)

        cmd = ord(self.decrypt(dat[:1]))
        if cmd == self.CMD_RST:
            msg = self.parseResetMessage(dat)
            if not msg.startswith('  '):
                logging.info('tunnel abnormal reset: %s' % msg)
                if self.canReturnErrorPage:
                    self._writer.write(gen_error_page("can't connect to destination", msg))
            self.onResetTunnel()
        elif cmd == self.CMD_DAT:
            dat = self.decrypt(dat[1:])
            if self.tunState != self.TUN_STATE_USING:
                # why this happens?
                return
            self.canReturnErrorPage = False
            self._writer.write(dat)
        else:
            logging.error('wrong command')

    def onClose(self, *args, **kwargs):
        super().onClose(*args, **kwargs)
        if self.inPool:
            self.pool.remove(self)

    @classmethod
    @asyncio.coroutine
    def _checkTimeout(cls, tun):
        while tun.state == cls.STATE_OPEN:
            timeout = (cls.TUN_MAX_IDLE_TIMEOUT if len(cls.pool) <= cls.POOL_NOM_SIZE else
                       cls.TUN_MIN_IDLE_TIMEOUT)
            yield from asyncio.sleep(timeout)
            if (tun.tunState == cls.TUN_STATE_IDLE and
               (time.time() - tun.lastIdleTime) > timeout):
                tun.sendClose(1000)

    @classmethod
    def addToPool(cls, tun):
        assert tun.tunState == cls.TUN_STATE_IDLE
        if len(cls.pool) >= cls.POOL_MAX_SIZE:
            tun.sendClose(1000)
        else:
            assert not tun.checkTimeoutTask
            tun.checkTimeoutTask = asyncio.async(cls._checkTimeout(tun))
            tun.inPool = True
            tun.enableAutoPing(cls.TUN_PING_INTERVAL)
            cls.pool.append(tun)

    @classmethod
    @asyncio.coroutine
    def getOrCreate(cls, addrHeader, dat):
        if WSTunClientProtocol.pool:
            logging.debug('reuse tunnel from pool (total %s)' % len(WSTunClientProtocol.pool))
            tun = WSTunClientProtocol.pool.popleft()
            tun.checkTimeoutTask.cancel()
            tun.checkTimeoutTask = None
            tun.inPool = False
            tun.disableAutoPing()
            tun.sendMessage(tun.makeRelayHeader(addrHeader, dat), True)
        else:
            tun = (yield from loop.create_connection(
                factory, config.uri_addr, config.uri_port, ssl=config.tun_ssl))[1]
            # lower latency by sending relay header and data in ws handshake
            tun.customUriPath = '/' + base64.urlsafe_b64encode(tun.makeRelayHeader(addrHeader, dat)).decode()
            tun.restartHandshake()
            try:
                yield from asyncio.wait_for(tun.tunOpen, cls.TUN_OPEN_TIMEOUT)
            except asyncio.TimeoutError:
                tun.dropConnection()
                raise ConnectionRefusedError(tun.wasNotCleanReason)
        return tun

factory = WebSocketClientFactory(config.uri)
factory.protocol = WSTunClientProtocol
factory.useragent = ''
factory.autoPingTimeout = 5
factory.openHandshakeTimeout = 10

@asyncio.coroutine
def socks5_tcp_handler(reader, writer):
    # handle auth method selection
    dat = yield from reader.read(257)
    if len(dat) < 2 or dat[0] != 0x05 or len(dat) != dat[1] + 2:
        logging.warning('local SOCKS v5 server got unknown request')
        return writer.close()
    writer.write(b'\x05\x00')  # \x00 == NO AUTHENTICATION REQUIRED

    # handle relay request
    try:
        dat = yield from reader.read(262)
    except ConnectionError:
        return writer.close()
    try:
        cmd, addr_header = dat[1], dat[2:]
        target_addr, target_port = parse_socks_addr(addr_header)
    except (ValueError, IndexError):
        logging.warning('invalid SOCKS v5 relay request')
        return writer.close()
    logging.info('requesting %s:%d' % (target_addr, target_port))
    if cmd != 0x01:  # CONNECT
        writer.write(b'\x05\x07\x00\x01' + b'\x00' * 6)  # \x07 == COMMAND NOT SUPPORTED
        return writer.close()

    # By accepting request before connected to target delay can be lowered (of a round-trip).
    # But SOCKS client can't get real reason when error happens (Firefox always display
    # connection reset error). Dirty solution: generate a HTML page when a HTTP request failed
    writer.write(b'\x05\x00\x00\x01' + b'\x01' * 6)  # \x00 == SUCCEEDED
    try:
        dat = yield from reader.read(2048)
    except ConnectionError:
        dat = None
    if not dat:
        return writer.close()

    canErr = can_return_error_page(dat)
    try:
        tun = yield from WSTunClientProtocol.getOrCreate(addr_header, dat)
    except Exception as e:
        if canErr:
            writer.write(gen_error_page("can't connect to wstan server", str(e)))
        return writer.close()
    tun.canReturnErrorPage = canErr
    tun.setProxy(reader, writer)

def main():
    server = loop.run_until_complete(
        asyncio.start_server(socks5_tcp_handler, config.addr, config.port))
    print('wstan client -- SOCKS v5 server listen on %s:%d' % (config.addr, config.port))
    try:
        loop.run_forever()
    finally:
        server.close()
        loop.close()
<|endoftext|>


<commit_before>__all__ = (
    'BaseForm',
    'Form',
)

class BaseForm(object):
    """
    Base Form Class.  Provides core behaviour like field construction,
    validation, and data and error proxying.
    """

    def __init__(self, fields, prefix=''):
        """
        :param unbound_fields:
            A dict which maps field names to UnboundField instances.
        :param prefix:
            If provided, all fields will have their name prefixed with the
            value.
        """
        if prefix:
            prefix += '-'

        self._errors = None
        self._fields = {}

        for name, unbound_field in fields.iteritems():
            field = unbound_field.bind(form=self, name=name, prefix=prefix)
            self._fields[name] = field

    def __iter__(self):
        """ Iterate form fields in arbitrary order"""
        return self._fields.itervalues()

    def __contains__(self, item):
        """ Returns `True` if the named field is a member of this form. """
        return (item in self._fields)

    def __getitem__(self, name):
        """ Dict-style access to this form's fields."""
        return self._fields[name]

    def __setitem__(self, name, value):
        self._fields[name] = value

    def __delitem__(self, name):
        del self._fields[name]

    def __getattr__(self, name):
        try:
            return self._fields[name]
        except KeyError:
            raise AttributeError('Form has no field %r' % name)

    def __delattr__(self, name):
        try:
            del self._fields[name]
        except KeyError:
            super(BaseForm, self).__delattr__(name)

    def populate_obj(self, obj):
        """
        Populates the attributes of the passed `obj` with data from the form's
        fields.

        :note: This is a destructive operation; Any attribute with the same name
               as a field will be overridden. Use with caution.
        """
        for name, field in self._fields.iteritems():
            field.populate_obj(obj, name)

    def process(self, formdata=None, obj=None, **kwargs):
        """
        :param formdata:
            Used to pass data coming from the enduser, usually `request.POST` or
            equivalent.
        :param obj:
            If `formdata` has no data for a field, the form will try to get it
            from the passed object.
        :param `**kwargs`:
            If neither `formdata` or `obj` contains a value for a field, the
            form will assign the value of a matching keyword argument to the
            field, if provided.
        """
        if not formdata:
            # XXX This is only because Field.process checks for None, which it
            # really shouldn't
            formdata = None
        for name, field, in self._fields.iteritems():
            if hasattr(obj, name):
                field.process(formdata, getattr(obj, name))
            elif name in kwargs:
                field.process(formdata, kwargs[name])
            else:
                field.process(formdata)

    def validate(self):
        """
        Validates the form by calling `validate` on each field, passing any
        extra `Form.validate_<fieldname>` validators to the field validator.

        Returns `True` if no errors occur.
        """
        self._errors = None
        success = True
        for name, field in self._fields.iteritems():
            extra = []
            inline = getattr(self.__class__, 'validate_%s' % name, None)
            if inline is not None:
                extra.append(inline)
            if not field.validate(self, extra):
                success = False
        return success

    @property
    def data(self):
        return dict((name, f.data) for name, f in self._fields.iteritems())

    @property
    def errors(self):
        if self._errors is None:
            self._errors = dict((name, f.errors) for name, f in self._fields.iteritems() if f.errors)
        return self._errors

class FormMeta(type):
    """
    The metaclass for `Form` and any subclasses of `Form`.

    `FormMeta`'s responsibility is to create the `_unbound_fields` list, which
    is a list of `UnboundField` instances sorted by their order of
    instantiation.  The list is created at the first instantiation of the form.
    If any fields are added/removed from the form, the list is cleared to be
    re-generated on the next instantiaton.

    Any properties which begin with an underscore or are not `UnboundField`
    instances are ignored by the metaclass.
    """
    def __init__(cls, name, bases, attrs):
        type.__init__(cls, name, bases, attrs)
        cls._unbound_fields = None

    def __call__(cls, *args, **kwargs):
        """
        Construct a new `Form` instance, creating `_unbound_fields` on the
        class if it is empty.
        """
        if cls._unbound_fields is None:
            fields = []
            for name in dir(cls):
                if not name.startswith('_'):
                    unbound_field = getattr(cls, name)
                    if hasattr(unbound_field, '_formfield'):
                        fields.append((name, unbound_field))
            # We keep the name as the second element of the sort
            # to ensure a stable sort.
            fields.sort(key=lambda x: (x[1].creation_counter, x[0]))
            cls._unbound_fields = fields
        return type.__call__(cls, *args, **kwargs)

    def __setattr__(cls, name, value):
        """
        Add an attribute to the class, clearing `_unbound_fields` if needed.
        """
        if not name.startswith('_') and hasattr(value, '_formfield'):
            cls._unbound_fields = None
        type.__setattr__(cls, name, value)

    def __delattr__(cls, name):
        """
        Remove an attribute from the class, clearing `_unbound_fields` if
        needed.
        """
        if not name.startswith('_'):
            cls._unbound_fields = None
        type.__delattr__(cls, name)

class Form(BaseForm):
    """
    Declarative Form base class. Extends BaseForm's core behaviour allowing
    fields to be defined on Form subclasses as class attributes.

    In addition, form and instance input data are taken at construction time
    and passed to `process()`.
    """
    __metaclass__ = FormMeta

    def __init__(self, formdata=None, obj=None, prefix='', **kwargs):
        """
        :param formdata:
            Used to pass data coming from the enduser, usually `request.POST` or
            equivalent.
        :param obj:
            If `formdata` has no data for a field, the form will try to get it
            from the passed object.
        :param prefix:
            If provided, all fields will have their name prefixed with the
            value.
        :param `**kwargs`:
            If neither `formdata` or `obj` contains a value for a field, the
            form will assign the value of a matching keyword argument to the
            field, if provided.
        """
        super(Form, self).__init__(dict(self._unbound_fields), prefix=prefix)

        for name in self._fields:
            # Set all the fields to attributes so that they obscure the class
            # attributes with the same names.
            setattr(self, name, self._fields[name])

        self.process(formdata, obj, **kwargs)

    def __iter__(self):
        """ Iterate form fields in their order of definition on the form. """
        for name, _ in self._unbound_fields:
            yield self._fields[name]

    def __delattr__(self, name):
        super(Form, self).__delattr__(name)
        setattr(self, name, None)
<commit_msg>BaseForm now won't suffix prefixes with a dash if they already end with punctuation. This is to aid integration with other frameworks. Also cleaned up some docs.<commit_after>__all__ = (
    'BaseForm',
    'Form',
)

class BaseForm(object):
    """
    Base Form Class.  Provides core behaviour like field construction,
    validation, and data and error proxying.
    """

    def __init__(self, fields, prefix=''):
        """
        :param fields:
            A dict of partially-constructed fields.
        :param prefix:
            If provided, all fields will have their name prefixed with the
            value.
        """
        if prefix and prefix[-1] not in '-_;:/.':
            prefix += '-'

        self._errors = None
        self._fields = {}

        for name, unbound_field in fields.iteritems():
            field = unbound_field.bind(form=self, name=name, prefix=prefix)
            self._fields[name] = field

    def __iter__(self):
        """ Iterate form fields in arbitrary order """
        return self._fields.itervalues()

    def __contains__(self, item):
        """ Returns `True` if the named field is a member of this form. """
        return (item in self._fields)

    def __getitem__(self, name):
        """ Dict-style access to this form's fields."""
        return self._fields[name]

    def __setitem__(self, name, value):
        self._fields[name] = value

    def __delitem__(self, name):
        del self._fields[name]

    def __getattr__(self, name):
        try:
            return self._fields[name]
        except KeyError:
            raise AttributeError('Form has no field %r' % name)

    def __delattr__(self, name):
        try:
            del self._fields[name]
        except KeyError:
            super(BaseForm, self).__delattr__(name)

    def populate_obj(self, obj):
        """
        Populates the attributes of the passed `obj` with data from the form's
        fields.

        :note: This is a destructive operation; Any attribute with the same name
               as a field will be overridden. Use with caution.
        """
        for name, field in self._fields.iteritems():
            field.populate_obj(obj, name)

    def process(self, formdata=None, obj=None, **kwargs):
        """
        Take form, object data, and keyword arg input and have the fields
        process them.

        :param formdata:
            Used to pass data coming from the enduser, usually `request.POST` or
            equivalent.
        :param obj:
            If `formdata` has no data for a field, the form will try to get it
            from the passed object.
        :param `**kwargs`:
            If neither `formdata` or `obj` contains a value for a field, the
            form will assign the value of a matching keyword argument to the
            field, if provided.
        """
        if not formdata:
            # XXX This is only because Field.process checks for None, which it
            # really shouldn't
            formdata = None
        for name, field, in self._fields.iteritems():
            if hasattr(obj, name):
                field.process(formdata, getattr(obj, name))
            elif name in kwargs:
                field.process(formdata, kwargs[name])
            else:
                field.process(formdata)

    def validate(self):
        """
        Validates the form by calling `validate` on each field, passing any
        extra `Form.validate_<fieldname>` validators to the field validator.

        Returns `True` if no errors occur.
        """
        self._errors = None
        success = True
        for name, field in self._fields.iteritems():
            extra = []
            inline = getattr(self.__class__, 'validate_%s' % name, None)
            if inline is not None:
                extra.append(inline)
            if not field.validate(self, extra):
                success = False
        return success

    @property
    def data(self):
        return dict((name, f.data) for name, f in self._fields.iteritems())

    @property
    def errors(self):
        if self._errors is None:
            self._errors = dict((name, f.errors) for name, f in self._fields.iteritems() if f.errors)
        return self._errors

class FormMeta(type):
    """
    The metaclass for `Form` and any subclasses of `Form`.

    `FormMeta`'s responsibility is to create the `_unbound_fields` list, which
    is a list of `UnboundField` instances sorted by their order of
    instantiation.  The list is created at the first instantiation of the form.
    If any fields are added/removed from the form, the list is cleared to be
    re-generated on the next instantiaton.

    Any properties which begin with an underscore or are not `UnboundField`
    instances are ignored by the metaclass.
    """
    def __init__(cls, name, bases, attrs):
        type.__init__(cls, name, bases, attrs)
        cls._unbound_fields = None

    def __call__(cls, *args, **kwargs):
        """
        Construct a new `Form` instance, creating `_unbound_fields` on the
        class if it is empty.
        """
        if cls._unbound_fields is None:
            fields = []
            for name in dir(cls):
                if not name.startswith('_'):
                    unbound_field = getattr(cls, name)
                    if hasattr(unbound_field, '_formfield'):
                        fields.append((name, unbound_field))
            # We keep the name as the second element of the sort
            # to ensure a stable sort.
            fields.sort(key=lambda x: (x[1].creation_counter, x[0]))
            cls._unbound_fields = fields
        return type.__call__(cls, *args, **kwargs)

    def __setattr__(cls, name, value):
        """
        Add an attribute to the class, clearing `_unbound_fields` if needed.
        """
        if not name.startswith('_') and hasattr(value, '_formfield'):
            cls._unbound_fields = None
        type.__setattr__(cls, name, value)

    def __delattr__(cls, name):
        """
        Remove an attribute from the class, clearing `_unbound_fields` if
        needed.
        """
        if not name.startswith('_'):
            cls._unbound_fields = None
        type.__delattr__(cls, name)

class Form(BaseForm):
    """
    Declarative Form base class. Extends BaseForm's core behaviour allowing
    fields to be defined on Form subclasses as class attributes.

    In addition, form and instance input data are taken at construction time
    and passed to `process()`.
    """
    __metaclass__ = FormMeta

    def __init__(self, formdata=None, obj=None, prefix='', **kwargs):
        """
        :param formdata:
            Used to pass data coming from the enduser, usually `request.POST` or
            equivalent.
        :param obj:
            If `formdata` has no data for a field, the form will try to get it
            from the passed object.
        :param prefix:
            If provided, all fields will have their name prefixed with the
            value.
        :param `**kwargs`:
            If neither `formdata` or `obj` contains a value for a field, the
            form will assign the value of a matching keyword argument to the
            field, if provided.
        """
        super(Form, self).__init__(dict(self._unbound_fields), prefix=prefix)

        for name in self._fields:
            # Set all the fields to attributes so that they obscure the class
            # attributes with the same names.
            setattr(self, name, self._fields[name])

        self.process(formdata, obj, **kwargs)

    def __iter__(self):
        """ Iterate form fields in their order of definition on the form. """
        for name, _ in self._unbound_fields:
            yield self._fields[name]

    def __delattr__(self, name):
        super(Form, self).__delattr__(name)
        setattr(self, name, None)
<|endoftext|>


<commit_before>import asyncio
import logging
from copy import copy
from enum import Enum
from datetime import datetime, timezone

from .holder import TaskHolder
from .task import register

log = logging.getLogger(__name__)

class JoinStatus(Enum):

    RUNNING = 'running'
    TIMEOUT = 'timeout'
    DONE = 'done'

@register('join', 'execute')
class JoinTask(TaskHolder):

    """
    A join task is an almost-standard task that allows awaiting for multiple
    parents and gathering their results using `data_received()`.
    The `wait_for` config parameter is mandatory and is a list of task IDs.
    """

    __slots__ = ('_data_stash', '_wait_for', '_report', '_task')

    def __init__(self, config):
        super().__init__(config)
        self._data_stash = []
        # Must be either a number of tasks or a list of ids
        # Copy because we alter the list/int during `_step`
        self._wait_for = copy(self.config['wait_for'])

        # Reporting
        self._task = None
        self._report = {'tasks': [], 'status': JoinStatus.RUNNING.value}

    def report(self):
        return self._report

    def add_task_report(self, task_id):
        task_report = {'id': task_id, 'end_time': datetime.now(timezone.utc)}
        self._report['tasks'].append(task_report)
        self._task.dispatch_progress(self._report)

    def _step(self, event):
        """
        Take a new event and consume it with the `wait_for` configuration.
        """
        task_id = event.source._task_template_id
        self.add_task_report(task_id)
        # Check if wait for a number of tasks
        if isinstance(self._wait_for, int):
            self._data_stash.append(event.data.copy())
            self._wait_for -= 1
            return self._wait_for == 0
        # Else, wait for list of ids
        elif isinstance(self._wait_for, list) and task_id in self._wait_for:
            self._data_stash.append(event.data.copy())
            self._wait_for.remove(task_id)
            return len(self._wait_for) == 0

    async def _wait_for_tasks(self):
        while True:
            new_event = await self.queue.get()
            if self._step(new_event) is True:
                return

    async def execute(self, event):
        log.info('Join task waiting for tasks (%s)', self._wait_for)
        self._task = asyncio.Task.current_task()
        # Trigger first step for this event
        self._step(event)
        data = event.data
        data['timeout'] = False
        try:
            await asyncio.wait_for(self._wait_for_tasks(), self._timeout)
        except asyncio.TimeoutError:
            log.warning("Join timed out, still waiting for %s", self._wait_for)
            self._report['status'] = JoinStatus.TIMEOUT.value
            data['timeout'] = True
            self._task.dispatch_progress(self._report)
        else:
            log.debug('All awaited parents joined')
            self._report['status'] = JoinStatus.DONE.value
            self._task.dispatch_progress(self._report)
        # Data contains the outputs of the first task to join
        # Variable `data_stash` contains a list of all parent tasks outputs
        data['data_stash'] = self._data_stash
        return data
<commit_msg>Use timeouts on join task<commit_after>import asyncio
import logging
from copy import copy
from enum import Enum
from datetime import datetime, timezone

from .holder import TaskHolder
from .task import register

log = logging.getLogger(__name__)

class JoinStatus(Enum):

    RUNNING = 'running'
    TIMEOUT = 'timeout'
    DONE = 'done'

@register('join', 'execute')
class JoinTask(TaskHolder):

    """
    A join task is an almost-standard task that allows awaiting for multiple
    parents and gathering their results using `data_received()`.
    The `wait_for` config parameter is mandatory and is a list of task IDs.
    """

    __slots__ = ('_data_stash', '_wait_for', '_report', '_task', 'data')

    def __init__(self, config):
        super().__init__(config)
        self._data_stash = []
        # Must be either a number of tasks or a list of ids
        # Copy because we alter the list/int during `_step`
        self._wait_for = copy(self.config['wait_for'])

        # Reporting
        self._task = None
        self.data = None
        self._report = {'tasks': [], 'status': JoinStatus.RUNNING.value}

    def report(self):
        return self._report

    def add_task_report(self, task_id):
        task_report = {'id': task_id, 'end_time': datetime.now(timezone.utc)}
        self._report['tasks'].append(task_report)
        self._task.dispatch_progress(self._report)

    def _step(self, event):
        """
        Take a new event and consume it with the `wait_for` configuration.
        """
        task_id = event.source._task_template_id
        self.add_task_report(task_id)
        # Check if wait for a number of tasks
        if isinstance(self._wait_for, int):
            self._data_stash.append(event.data.copy())
            self._wait_for -= 1
            return self._wait_for == 0
        # Else, wait for list of ids
        elif isinstance(self._wait_for, list) and task_id in self._wait_for:
            self._data_stash.append(event.data.copy())
            self._wait_for.remove(task_id)
            return len(self._wait_for) == 0

    async def _wait_for_tasks(self):
        while True:
            new_event = await self.queue.get()
            if self._step(new_event) is True:
                return

    async def execute(self, event):
        log.info('Join task waiting for tasks (%s)', self._wait_for)
        self._task = asyncio.Task.current_task()
        # Trigger first step for this event
        self._step(event)
        self.data = event.data
        await self._wait_for_tasks()
        log.debug('All awaited parents joined')
        self._report['status'] = JoinStatus.DONE.value
        self._task.dispatch_progress(self._report)
        # Data contains the outputs of the first task to join
        # Variable `data_stash` contains a list of all parent tasks outputs
        self.data['data_stash'] = self._data_stash
        return self.data

    async def teardown(self):
        self._report['status'] = JoinStatus.TIMEOUT.value
        self._task.dispatch_progress(self._report)
        self.data['data_stash'] = self._data_stash
        return self.data
<|endoftext|>


<commit_before># Copyright (c) 2014 by California Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the California Institute of Technology nor
#    the names of its contributors may be used to endorse or promote
#    products derived from this software without specific prior
#    written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL CALTECH
# OR THE CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
"""Code generation and exporting of controllers from TuLiP

Routines in this module are cross-cutting in the sense that they
concern multiple aspects of solutions created by TuLiP and accordingly
should not be placed under a specific subpackage, like tulip.transys.
"""

import Queue
import re
import time

def write_python_case(filename, *args, **kwargs):
    with open(filename, "w") as f:
        f.write(python_case(*args, **kwargs))

def python_case(M, classname="TulipStrategy"):
    """Export MealyMachine as Python class based on flat if-else block

    Assumptions
    ===========
    - Initial state is "Sinit".

    @type M: L{MealyMachine}
    @rtype: str
    """
    tab = 4*" "
    nl = "\n"
    state_table = dict([(s,i) for (i,s) in enumerate(M.states)])
    if len(M.outputs) == 1:
        output_prefix = ""
        output_suffix = ""
    else:
        output_prefix = "("
        output_suffix = ")"
    output_tuple = output_prefix+", ".join([str(v) for v in M.outputs])+output_suffix

    code = "class "+classname+":" + nl

    code += tab + "\"\"\"" + 2*nl
    code += tab + "Automatically generated by tulip.dumpsmach on "+time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()) + nl
    code += tab + "To learn more about TuLiP, visit http://tulip-control.org" + nl
    code += tab + "\"\"\"" + nl

    code += tab + "def __init__(self):"+nl
    code += 2*tab+"self.state = "+str(state_table["Sinit"]) + 2*nl

    code += tab + "def move(self"
    if len(M.inputs) > 0:
        code += ", "+", ".join(M.inputs)
    code += "):" + nl
    code += 2*tab + "\"\"\"Given inputs, take move and return outputs." + 2*nl
    code += 2*tab + "@return: "+output_tuple + nl
    code += 2*tab + "\"\"\"" + nl
    first = True
    for state in M.states:
        code += 2*tab
        if first:
            first = False
            code += "if"
        else:
            code += "elif"
        code += " self.state == "+str(state_table[state])+":" + nl

        first = True
        for (from_state, to_state, label_dict) in M.transitions(nbunch=(state,),
                                                                data=True):
            assert from_state == state
            code += 3*tab
            if first:
                first = False
                code += "if"
            else:
                code += "elif"
            if len(M.inputs) > 0:
                code += " (" + ") and (".join([str(k)+" == "+str(v) for (k,v) in label_dict.items() if k in M.inputs])+"):" + nl
            else:
                code += " True:" + nl
            code += 4*tab + "self.state = "+str(state_table[to_state]) + nl
            code += 4*tab + (nl+4*tab).join([str(k)+" = "+str(v) for (k,v) in label_dict.items() if k in M.outputs]) + nl

        if len(M.inputs) > 0:
            code += 3*tab + "else:" + nl
            code += 4*tab + "raise Exception(\"Unrecognized input: "
            code += "+\";".join([" "+str(v)+" = \"+str("+str(v)+")" for v in M.inputs])
            code += ")" + nl

    code += 2*tab + "else:" + nl
    code += 3*tab + "raise Exception(\"Unrecognized internal state: \"+str(self.state))" + nl

    code += 2*tab + "return " + output_tuple + nl

    return code
<commit_msg>TRIV: Delete unused imports in dumpsmach<commit_after># Copyright (c) 2014 by California Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the California Institute of Technology nor
#    the names of its contributors may be used to endorse or promote
#    products derived from this software without specific prior
#    written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL CALTECH
# OR THE CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
"""Code generation and exporting of controllers from TuLiP

Routines in this module are cross-cutting in the sense that they
concern multiple aspects of solutions created by TuLiP and accordingly
should not be placed under a specific subpackage, like tulip.transys.
"""

import time

def write_python_case(filename, *args, **kwargs):
    with open(filename, "w") as f:
        f.write(python_case(*args, **kwargs))

def python_case(M, classname="TulipStrategy"):
    """Export MealyMachine as Python class based on flat if-else block

    Assumptions
    ===========
    - Initial state is "Sinit".

    @type M: L{MealyMachine}
    @rtype: str
    """
    tab = 4*" "
    nl = "\n"
    state_table = dict([(s,i) for (i,s) in enumerate(M.states)])
    if len(M.outputs) == 1:
        output_prefix = ""
        output_suffix = ""
    else:
        output_prefix = "("
        output_suffix = ")"
    output_tuple = output_prefix+", ".join([str(v) for v in M.outputs])+output_suffix

    code = "class "+classname+":" + nl

    code += tab + "\"\"\"" + 2*nl
    code += tab + "Automatically generated by tulip.dumpsmach on "+time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()) + nl
    code += tab + "To learn more about TuLiP, visit http://tulip-control.org" + nl
    code += tab + "\"\"\"" + nl

    code += tab + "def __init__(self):"+nl
    code += 2*tab+"self.state = "+str(state_table["Sinit"]) + 2*nl

    code += tab + "def move(self"
    if len(M.inputs) > 0:
        code += ", "+", ".join(M.inputs)
    code += "):" + nl
    code += 2*tab + "\"\"\"Given inputs, take move and return outputs." + 2*nl
    code += 2*tab + "@return: "+output_tuple + nl
    code += 2*tab + "\"\"\"" + nl
    first = True
    for state in M.states:
        code += 2*tab
        if first:
            first = False
            code += "if"
        else:
            code += "elif"
        code += " self.state == "+str(state_table[state])+":" + nl

        first = True
        for (from_state, to_state, label_dict) in M.transitions(nbunch=(state,),
                                                                data=True):
            assert from_state == state
            code += 3*tab
            if first:
                first = False
                code += "if"
            else:
                code += "elif"
            if len(M.inputs) > 0:
                code += " (" + ") and (".join([str(k)+" == "+str(v) for (k,v) in label_dict.items() if k in M.inputs])+"):" + nl
            else:
                code += " True:" + nl
            code += 4*tab + "self.state = "+str(state_table[to_state]) + nl
            code += 4*tab + (nl+4*tab).join([str(k)+" = "+str(v) for (k,v) in label_dict.items() if k in M.outputs]) + nl

        if len(M.inputs) > 0:
            code += 3*tab + "else:" + nl
            code += 4*tab + "raise Exception(\"Unrecognized input: "
            code += "+\";".join([" "+str(v)+" = \"+str("+str(v)+")" for v in M.inputs])
            code += ")" + nl

    code += 2*tab + "else:" + nl
    code += 3*tab + "raise Exception(\"Unrecognized internal state: \"+str(self.state))" + nl

    code += 2*tab + "return " + output_tuple + nl

    return code
<|endoftext|>


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

import os
import re
import shutil
import stat
import subprocess
import sys
import tempfile
import zipfile

def _unzip(name):
  destination = tempfile.mkdtemp()
  with zipfile.ZipFile(name) as f:
    f.extractall(destination)
  return destination

def _class_files(path):
  classes = []
  for root, dirs, files in os.walk(path):
    for file in files:
      if file.endswith('.class'):
        classes.append(os.path.join(root, file))
  return classes

def _javap_public(files):
  return subprocess.check_output(['javap', '-public'] + files)

def _split_info_into_infos(info):
  original = re.findall(r'Compiled from ".*?\.java"\n.*?\n(?:  .*\n)*}\n', info, flags=re.MULTILINE)
  infos = {}
  for info in original:
    lines = str(info).split('\n')[1:]
    header = lines[0]
    if not header.startswith('public '):
      continue
    match = re.match(r'(?:public |abstract |class |interface |final )+([^ ]+) .*', header)
    name = match.group(1)
    infos[name] = '\n'.join(lines)
  return infos

def _write_infos_to_temp(temp_folder, original_stat, infos):
  times = (original_stat[stat.ST_ATIME], original_stat[stat.ST_MTIME])

  for name, info in infos.items():
    out = os.path.join(temp_folder, name)
    with open(out, 'w') as f:
      f.write(info)
    os.utime(out, times)

def _exec_diff(base_path, one_path, two_path):
  one_path = os.path.relpath(one_path, base_path)
  two_path = os.path.relpath(two_path, base_path)
  subprocess.call(['diff', '-U', '0', '-N', one_path, two_path], cwd=base_path)

def process_archive(temp_folder, jar):
  name = os.path.splitext(os.path.basename(jar))[0]
  jar_stat = os.stat(jar)
  archive_folder = os.path.join(temp_folder, name)
  os.mkdir(archive_folder)

  unaar = None
  if jar.endswith('.aar'):
    unaar = _unzip(jar)
    jar = os.path.join(unaar, 'classes.jar')

  unjar = _unzip(jar)
  classes = _class_files(unjar)
  info = _javap_public(classes)
  infos = _split_info_into_infos(info)

  _write_infos_to_temp(archive_folder, jar_stat, infos)

  shutil.rmtree(unjar)
  if unaar is not None:
    shutil.rmtree(unaar)

  return archive_folder

def _main(old_archive, new_archive):
  temp_folder = tempfile.mkdtemp()

  old_data = process_archive(temp_folder, old_archive)
  new_data = process_archive(temp_folder, new_archive)
  _exec_diff(temp_folder, old_data, new_data)

  shutil.rmtree(temp_folder)

if __name__ == '__main__':
  if len(sys.argv) != 3:
    print('Usage: %s old.jar new.jar' % sys.argv[0])
    sys.exit(1)
  _main(sys.argv[1], sys.argv[2])
<commit_msg>Ensure subprocess result is eagerly converted back to a string.<commit_after>#!/usr/bin/env python

import os
import re
import shutil
import stat
import subprocess
import sys
import tempfile
import zipfile

def _unzip(name):
  destination = tempfile.mkdtemp()
  with zipfile.ZipFile(name) as f:
    f.extractall(destination)
  return destination

def _class_files(path):
  classes = []
  for root, dirs, files in os.walk(path):
    for file in files:
      if file.endswith('.class'):
        classes.append(os.path.join(root, file))
  return classes

def _javap_public(files):
  return str(subprocess.check_output(['javap', '-public'] + files))

def _split_info_into_infos(info):
  original = re.findall(r'Compiled from ".*?\.java"\n.*?\n(?:  .*\n)*}\n', info, flags=re.MULTILINE)
  infos = {}
  for info in original:
    lines = str(info).split('\n')[1:]
    header = lines[0]
    if not header.startswith('public '):
      continue
    match = re.match(r'(?:public |abstract |class |interface |final )+([^ ]+) .*', header)
    name = match.group(1)
    infos[name] = '\n'.join(lines)
  return infos

def _write_infos_to_temp(temp_folder, original_stat, infos):
  times = (original_stat[stat.ST_ATIME], original_stat[stat.ST_MTIME])

  for name, info in infos.items():
    out = os.path.join(temp_folder, name)
    with open(out, 'w') as f:
      f.write(info)
    os.utime(out, times)

def _exec_diff(base_path, one_path, two_path):
  one_path = os.path.relpath(one_path, base_path)
  two_path = os.path.relpath(two_path, base_path)
  subprocess.call(['diff', '-U', '0', '-N', one_path, two_path], cwd=base_path)

def process_archive(temp_folder, jar):
  name = os.path.splitext(os.path.basename(jar))[0]
  jar_stat = os.stat(jar)
  archive_folder = os.path.join(temp_folder, name)
  os.mkdir(archive_folder)

  unaar = None
  if jar.endswith('.aar'):
    unaar = _unzip(jar)
    jar = os.path.join(unaar, 'classes.jar')

  unjar = _unzip(jar)
  classes = _class_files(unjar)
  info = _javap_public(classes)
  infos = _split_info_into_infos(info)

  _write_infos_to_temp(archive_folder, jar_stat, infos)

  shutil.rmtree(unjar)
  if unaar is not None:
    shutil.rmtree(unaar)

  return archive_folder

def _main(old_archive, new_archive):
  temp_folder = tempfile.mkdtemp()

  old_data = process_archive(temp_folder, old_archive)
  new_data = process_archive(temp_folder, new_archive)
  _exec_diff(temp_folder, old_data, new_data)

  shutil.rmtree(temp_folder)

if __name__ == '__main__':
  if len(sys.argv) != 3:
    print('Usage: %s old.jar new.jar' % sys.argv[0])
    sys.exit(1)
  _main(sys.argv[1], sys.argv[2])
<|endoftext|>


<commit_before>from collections import namedtuple, defaultdict
from copy import deepcopy
from datetime import datetime
import sys

import json
from django.conf import settings
from kafka import KafkaConsumer

from corehq.util.io import ClosingContextProxy
from corehq.util.json import CommCareJSONEncoder
from dimagi.utils.chunked import chunked
from dimagi.utils.modules import to_function

from pillowtop.exceptions import PillowNotFoundError
from pillowtop.logger import pillow_logging
from pillowtop.dao.exceptions import DocumentMismatchError, DocumentMissingError

def _get_pillow_instance(full_class_str):
    pillow_class = _import_class_or_function(full_class_str)
    if pillow_class is None:
        raise ValueError('No pillow class found for {}'.format(full_class_str))
    return pillow_class()

def _import_class_or_function(full_class_str):
    return to_function(full_class_str, failhard=settings.DEBUG)

def get_all_pillow_classes():
    return [config.get_class() for config in get_all_pillow_configs()]

def get_all_pillow_instances():
    return [config.get_instance() for config in get_all_pillow_configs()]

def get_couch_pillow_instances():
    from pillowtop.feed.couch import CouchChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), CouchChangeFeed)
    ]

def get_kafka_pillow_instances():
    from corehq.apps.change_feed.consumer.feed import KafkaChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), KafkaChangeFeed)
    ]

def get_all_pillow_configs():
    return get_pillow_configs_from_settings_dict(getattr(settings, 'PILLOWTOPS', {}))

def get_pillow_configs_from_settings_dict(pillow_settings_dict):
    """
    The pillow_settings_dict is expected to be a dict mapping groups to list of pillow configs
    """
    for section, list_of_pillows in pillow_settings_dict.items():
        for pillow_config in list_of_pillows:
            yield get_pillow_config_from_setting(section, pillow_config)

class PillowConfig(namedtuple('PillowConfig', ['section', 'name', 'class_name', 'instance_generator', 'params'])):
    """
    Helper object for getting pillow classes/instances from settings
    """

    def __hash__(self):
        return hash(self.name)

    def get_class(self):
        return _import_class_or_function(self.class_name)

    def get_instance(self, **kwargs):
        if self.instance_generator:
            instance_generator_fn = _import_class_or_function(self.instance_generator)
            params = deepcopy(self.params)  # parameters defined in settings file
            params.update(kwargs)  # parameters passed in via run_ptop
            return instance_generator_fn(self.name, **params)
        else:
            return _get_pillow_instance(self.class_name)

def get_pillow_config_from_setting(section, pillow_config_string_or_dict):
    if isinstance(pillow_config_string_or_dict, str):
        return PillowConfig(
            section,
            pillow_config_string_or_dict.rsplit('.', 1)[1],
            pillow_config_string_or_dict,
            None,
            {}
        )
    else:
        assert 'class' in pillow_config_string_or_dict
        class_name = pillow_config_string_or_dict['class']
        return PillowConfig(
            section,
            pillow_config_string_or_dict.get('name', class_name),
            class_name,
            pillow_config_string_or_dict.get('instance', None),
            pillow_config_string_or_dict.get('params', {}),
        )

def get_pillow_by_name(pillow_class_name, instantiate=True, **kwargs):
    config = get_pillow_config_by_name(pillow_class_name)
    return config.get_instance(**kwargs) if instantiate else config.get_class()

def get_pillow_config_by_name(pillow_name):
    all_configs = get_all_pillow_configs()
    for config in all_configs:
        if config.name == pillow_name:
            return config
    raise PillowNotFoundError('No pillow found with name {}'.format(pillow_name))

def force_seq_int(seq):
    if seq is None or seq == '':
        return None
    elif isinstance(seq, dict):
        # multi-topic checkpoints don't support a single sequence id
        return None
    elif isinstance(seq, str):
        return int(seq.split('-')[0])
    else:
        assert isinstance(seq, int), type(seq)
        return seq

def safe_force_seq_int(seq, default=None):
    if isinstance(seq, dict):
        return default
    try:
        return force_seq_int(seq)
    except (AssertionError, ValueError):
        return default

def _get_consumer():
    return ClosingContextProxy(KafkaConsumer(
        client_id='pillowtop_utils',
        bootstrap_servers=settings.KAFKA_BROKERS,
        request_timeout_ms=1000
    ))

def get_all_pillows_json():
    pillow_configs = get_all_pillow_configs()
    consumer = _get_consumer()
    with consumer:
        return [get_pillow_json(pillow_config, consumer) for pillow_config in pillow_configs]

def get_pillow_json(pillow_config, consumer=None):
    assert isinstance(pillow_config, PillowConfig)

    def _is_kafka(checkpoint):
        return checkpoint.sequence_format == 'json'

    pillow = pillow_config.get_instance()
    if consumer is None:
        consumer = _get_consumer()

    checkpoint = pillow.get_checkpoint()
    timestamp = checkpoint.timestamp
    if timestamp:
        time_since_last = datetime.utcnow() - timestamp
        seconds_since_last = time_since_last.total_seconds()
        hours_since_last = seconds_since_last // 3600

        try:
            # remove microsecond portion
            prettified_time_since_last = str(time_since_last)
            prettified_time_since_last = prettified_time_since_last[0:prettified_time_since_last.index('.')]
        except ValueError:
            pass
    else:
        seconds_since_last = 0
        prettified_time_since_last = ''
        hours_since_last = None

    if _is_kafka(checkpoint):
        # breaking composition boundaries so that we can use one Kafka connection
        offsets = {
            "{},{}".format(tp.topic, tp.partition): offset
            for tp, offset in consumer.end_offsets(list(checkpoint.wrapped_sequence.keys())).items()
        }
    else:
        offsets = pillow.get_change_feed().get_latest_offsets_json()

    def _seq_to_int(checkpoint, seq):
        from pillowtop.models import kafka_seq_to_str
        if _is_kafka(checkpoint):
            return json.loads(kafka_seq_to_str(seq))
        else:
            return force_seq_int(seq)

    return {
        'name': pillow_config.name,
        'seq_format': checkpoint.sequence_format,
        'seq': _seq_to_int(checkpoint, checkpoint.wrapped_sequence),
        'offsets': offsets,
        'seconds_since_last': seconds_since_last,
        'time_since_last': prettified_time_since_last,
        'hours_since_last': hours_since_last
    }

ChangeError = namedtuple('ChangeError', 'change exception')

class ErrorCollector(object):
    def __init__(self):
        self.errors = []

    def add_error(self, error):
        self.errors.append(error)

def build_bulk_payload(index_info, changes, doc_transform=None, error_collector=None):
    doc_transform = doc_transform or (lambda x: x)
    payload = []
    for change in changes:
        if change.deleted and change.id:
            payload.append({
                "delete": {
                    "_index": index_info.index,
                    "_type": index_info.type,
                    "_id": change.id
                }
            })
        elif not change.deleted:
            try:
                doc = change.get_document()
                doc = doc_transform(doc)
                payload.append({
                    "index": {
                        "_index": index_info.index,
                        "_type": index_info.type,
                        "_id": doc['_id']
                    }
                })
                payload.append(doc)
            except Exception as e:
                if not error_collector:
                    raise
                error_collector.add_error(ChangeError(change, e))
    return payload

def prepare_bulk_payloads(bulk_changes, max_size, chunk_size=100):
    payloads = [b'']
    for bulk_chunk in chunked(bulk_changes, chunk_size):
        current_payload = payloads[-1]
        json_bulk_chunks = [
            json.dumps(obj, cls=CommCareJSONEncoder).encode('utf-8')
            for obj in bulk_chunk
        ]
        payload_chunk = b'\n'.join(json_bulk_chunks) + b'\n'
        appended_payload = current_payload + payload_chunk
        new_payload_size = sys.getsizeof(appended_payload)
        if new_payload_size > max_size:
            payloads.append(payload_chunk)
        else:
            payloads[-1] = appended_payload

    return [_f for _f in payloads if _f]

def ensure_matched_revisions(change, fetched_document):
    """
    This function ensures that the document fetched from a change matches the
    revision at which it was pushed to kafka at.

    See http://manage.dimagi.com/default.asp?237983 for more details

    :raises: DocumentMismatchError - Raised when the revisions of the fetched document
        and the change metadata do not match
    """

    change_has_rev = change.metadata and change.metadata.document_rev is not None
    doc_has_rev = fetched_document and '_rev' in fetched_document
    if doc_has_rev and change_has_rev:

        doc_rev = fetched_document['_rev']
        change_rev = change.metadata.document_rev
        if doc_rev != change_rev:
            fetched_rev = _convert_rev_to_int(doc_rev)
            stored_rev = _convert_rev_to_int(change_rev)
            if fetched_rev < stored_rev or stored_rev == -1:
                message = "Mismatched revs for {}: Cloudant rev {} vs. Changes feed rev {}".format(
                    change.id,
                    doc_rev,
                    change_rev
                )
                pillow_logging.warning(message)
                raise DocumentMismatchError(message)

def _convert_rev_to_int(rev):
    try:
        return int(rev.split('-')[0])
    except (ValueError, AttributeError):
        return -1

def ensure_document_exists(change):
    """
    This is to ensure that the Couch document exists in the Couch database when the
    change is processed. We only care about the scenario where the document is missing.

    :raises: DocumentMissingError - Raised when the document is missing (never existed)
    """
    change.get_document()
    if change.error_raised is not None and isinstance(change.error_raised, DocumentMissingError):
        raise change.error_raised

def bulk_fetch_changes_docs(changes, domain):
    # break up by doctype
    changes_by_doctype = defaultdict(list)
    for change in changes:
        assert change.metadata.domain == domain
        changes_by_doctype[change.metadata.data_source_name].append(change)

    # query
    docs = []
    for _, _changes in changes_by_doctype.items():
        doc_store = _changes[0].document_store
        doc_ids_to_query = [change.id for change in _changes if change.should_fetch_document()]
        new_docs = list(doc_store.iter_documents(doc_ids_to_query))
        docs_queried_prior = [change.document for change in _changes if not change.should_fetch_document()]
        docs.extend(new_docs + docs_queried_prior)

    # catch missing docs
    retry_changes = set()
    docs_by_id = {doc['_id']: doc for doc in docs}
    for change in changes:
        if change.id not in docs_by_id:
            # we need to capture DocumentMissingError which is not possible in bulk
            #   so let pillow fall back to serial mode to capture the error for missing docs
            retry_changes.add(change)
            continue
        else:
            # set this, so that subsequent doc lookups are avoided
            change.set_document(docs_by_id[change.id])
        try:
            ensure_matched_revisions(change, docs_by_id.get(change.id))
        except DocumentMismatchError:
            retry_changes.add(change)
    return retry_changes, docs
<commit_msg>make domain check optional and don't use assert<commit_after>from collections import namedtuple, defaultdict
from copy import deepcopy
from datetime import datetime
import sys

import json
from django.conf import settings
from kafka import KafkaConsumer

from corehq.util.io import ClosingContextProxy
from corehq.util.json import CommCareJSONEncoder
from dimagi.utils.chunked import chunked
from dimagi.utils.modules import to_function

from pillowtop.exceptions import PillowNotFoundError
from pillowtop.logger import pillow_logging
from pillowtop.dao.exceptions import DocumentMismatchError, DocumentMissingError

def _get_pillow_instance(full_class_str):
    pillow_class = _import_class_or_function(full_class_str)
    if pillow_class is None:
        raise ValueError('No pillow class found for {}'.format(full_class_str))
    return pillow_class()

def _import_class_or_function(full_class_str):
    return to_function(full_class_str, failhard=settings.DEBUG)

def get_all_pillow_classes():
    return [config.get_class() for config in get_all_pillow_configs()]

def get_all_pillow_instances():
    return [config.get_instance() for config in get_all_pillow_configs()]

def get_couch_pillow_instances():
    from pillowtop.feed.couch import CouchChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), CouchChangeFeed)
    ]

def get_kafka_pillow_instances():
    from corehq.apps.change_feed.consumer.feed import KafkaChangeFeed
    return [
        pillow for pillow in get_all_pillow_instances()
        if isinstance(pillow.get_change_feed(), KafkaChangeFeed)
    ]

def get_all_pillow_configs():
    return get_pillow_configs_from_settings_dict(getattr(settings, 'PILLOWTOPS', {}))

def get_pillow_configs_from_settings_dict(pillow_settings_dict):
    """
    The pillow_settings_dict is expected to be a dict mapping groups to list of pillow configs
    """
    for section, list_of_pillows in pillow_settings_dict.items():
        for pillow_config in list_of_pillows:
            yield get_pillow_config_from_setting(section, pillow_config)

class PillowConfig(namedtuple('PillowConfig', ['section', 'name', 'class_name', 'instance_generator', 'params'])):
    """
    Helper object for getting pillow classes/instances from settings
    """

    def __hash__(self):
        return hash(self.name)

    def get_class(self):
        return _import_class_or_function(self.class_name)

    def get_instance(self, **kwargs):
        if self.instance_generator:
            instance_generator_fn = _import_class_or_function(self.instance_generator)
            params = deepcopy(self.params)  # parameters defined in settings file
            params.update(kwargs)  # parameters passed in via run_ptop
            return instance_generator_fn(self.name, **params)
        else:
            return _get_pillow_instance(self.class_name)

def get_pillow_config_from_setting(section, pillow_config_string_or_dict):
    if isinstance(pillow_config_string_or_dict, str):
        return PillowConfig(
            section,
            pillow_config_string_or_dict.rsplit('.', 1)[1],
            pillow_config_string_or_dict,
            None,
            {}
        )
    else:
        assert 'class' in pillow_config_string_or_dict
        class_name = pillow_config_string_or_dict['class']
        return PillowConfig(
            section,
            pillow_config_string_or_dict.get('name', class_name),
            class_name,
            pillow_config_string_or_dict.get('instance', None),
            pillow_config_string_or_dict.get('params', {}),
        )

def get_pillow_by_name(pillow_class_name, instantiate=True, **kwargs):
    config = get_pillow_config_by_name(pillow_class_name)
    return config.get_instance(**kwargs) if instantiate else config.get_class()

def get_pillow_config_by_name(pillow_name):
    all_configs = get_all_pillow_configs()
    for config in all_configs:
        if config.name == pillow_name:
            return config
    raise PillowNotFoundError('No pillow found with name {}'.format(pillow_name))

def force_seq_int(seq):
    if seq is None or seq == '':
        return None
    elif isinstance(seq, dict):
        # multi-topic checkpoints don't support a single sequence id
        return None
    elif isinstance(seq, str):
        return int(seq.split('-')[0])
    else:
        assert isinstance(seq, int), type(seq)
        return seq

def safe_force_seq_int(seq, default=None):
    if isinstance(seq, dict):
        return default
    try:
        return force_seq_int(seq)
    except (AssertionError, ValueError):
        return default

def _get_consumer():
    return ClosingContextProxy(KafkaConsumer(
        client_id='pillowtop_utils',
        bootstrap_servers=settings.KAFKA_BROKERS,
        request_timeout_ms=1000
    ))

def get_all_pillows_json():
    pillow_configs = get_all_pillow_configs()
    consumer = _get_consumer()
    with consumer:
        return [get_pillow_json(pillow_config, consumer) for pillow_config in pillow_configs]

def get_pillow_json(pillow_config, consumer=None):
    assert isinstance(pillow_config, PillowConfig)

    def _is_kafka(checkpoint):
        return checkpoint.sequence_format == 'json'

    pillow = pillow_config.get_instance()
    if consumer is None:
        consumer = _get_consumer()

    checkpoint = pillow.get_checkpoint()
    timestamp = checkpoint.timestamp
    if timestamp:
        time_since_last = datetime.utcnow() - timestamp
        seconds_since_last = time_since_last.total_seconds()
        hours_since_last = seconds_since_last // 3600

        try:
            # remove microsecond portion
            prettified_time_since_last = str(time_since_last)
            prettified_time_since_last = prettified_time_since_last[0:prettified_time_since_last.index('.')]
        except ValueError:
            pass
    else:
        seconds_since_last = 0
        prettified_time_since_last = ''
        hours_since_last = None

    if _is_kafka(checkpoint):
        # breaking composition boundaries so that we can use one Kafka connection
        offsets = {
            "{},{}".format(tp.topic, tp.partition): offset
            for tp, offset in consumer.end_offsets(list(checkpoint.wrapped_sequence.keys())).items()
        }
    else:
        offsets = pillow.get_change_feed().get_latest_offsets_json()

    def _seq_to_int(checkpoint, seq):
        from pillowtop.models import kafka_seq_to_str
        if _is_kafka(checkpoint):
            return json.loads(kafka_seq_to_str(seq))
        else:
            return force_seq_int(seq)

    return {
        'name': pillow_config.name,
        'seq_format': checkpoint.sequence_format,
        'seq': _seq_to_int(checkpoint, checkpoint.wrapped_sequence),
        'offsets': offsets,
        'seconds_since_last': seconds_since_last,
        'time_since_last': prettified_time_since_last,
        'hours_since_last': hours_since_last
    }

ChangeError = namedtuple('ChangeError', 'change exception')

class ErrorCollector(object):
    def __init__(self):
        self.errors = []

    def add_error(self, error):
        self.errors.append(error)

def build_bulk_payload(index_info, changes, doc_transform=None, error_collector=None):
    doc_transform = doc_transform or (lambda x: x)
    payload = []
    for change in changes:
        if change.deleted and change.id:
            payload.append({
                "delete": {
                    "_index": index_info.index,
                    "_type": index_info.type,
                    "_id": change.id
                }
            })
        elif not change.deleted:
            try:
                doc = change.get_document()
                doc = doc_transform(doc)
                payload.append({
                    "index": {
                        "_index": index_info.index,
                        "_type": index_info.type,
                        "_id": doc['_id']
                    }
                })
                payload.append(doc)
            except Exception as e:
                if not error_collector:
                    raise
                error_collector.add_error(ChangeError(change, e))
    return payload

def prepare_bulk_payloads(bulk_changes, max_size, chunk_size=100):
    payloads = [b'']
    for bulk_chunk in chunked(bulk_changes, chunk_size):
        current_payload = payloads[-1]
        json_bulk_chunks = [
            json.dumps(obj, cls=CommCareJSONEncoder).encode('utf-8')
            for obj in bulk_chunk
        ]
        payload_chunk = b'\n'.join(json_bulk_chunks) + b'\n'
        appended_payload = current_payload + payload_chunk
        new_payload_size = sys.getsizeof(appended_payload)
        if new_payload_size > max_size:
            payloads.append(payload_chunk)
        else:
            payloads[-1] = appended_payload

    return [_f for _f in payloads if _f]

def ensure_matched_revisions(change, fetched_document):
    """
    This function ensures that the document fetched from a change matches the
    revision at which it was pushed to kafka at.

    See http://manage.dimagi.com/default.asp?237983 for more details

    :raises: DocumentMismatchError - Raised when the revisions of the fetched document
        and the change metadata do not match
    """

    change_has_rev = change.metadata and change.metadata.document_rev is not None
    doc_has_rev = fetched_document and '_rev' in fetched_document
    if doc_has_rev and change_has_rev:

        doc_rev = fetched_document['_rev']
        change_rev = change.metadata.document_rev
        if doc_rev != change_rev:
            fetched_rev = _convert_rev_to_int(doc_rev)
            stored_rev = _convert_rev_to_int(change_rev)
            if fetched_rev < stored_rev or stored_rev == -1:
                message = "Mismatched revs for {}: Cloudant rev {} vs. Changes feed rev {}".format(
                    change.id,
                    doc_rev,
                    change_rev
                )
                pillow_logging.warning(message)
                raise DocumentMismatchError(message)

def _convert_rev_to_int(rev):
    try:
        return int(rev.split('-')[0])
    except (ValueError, AttributeError):
        return -1

def ensure_document_exists(change):
    """
    This is to ensure that the Couch document exists in the Couch database when the
    change is processed. We only care about the scenario where the document is missing.

    :raises: DocumentMissingError - Raised when the document is missing (never existed)
    """
    change.get_document()
    if change.error_raised is not None and isinstance(change.error_raised, DocumentMissingError):
        raise change.error_raised

def bulk_fetch_changes_docs(changes, domain=None):
    # break up by doctype
    changes_by_doctype = defaultdict(list)
    for change in changes:
        if domain and change.metadata.domain == domain:
            raise ValueError("Domain does not match change")
        changes_by_doctype[change.metadata.data_source_name].append(change)

    # query
    docs = []
    for _, _changes in changes_by_doctype.items():
        doc_store = _changes[0].document_store
        doc_ids_to_query = [change.id for change in _changes if change.should_fetch_document()]
        new_docs = list(doc_store.iter_documents(doc_ids_to_query))
        docs_queried_prior = [change.document for change in _changes if not change.should_fetch_document()]
        docs.extend(new_docs + docs_queried_prior)

    # catch missing docs
    retry_changes = set()
    docs_by_id = {doc['_id']: doc for doc in docs}
    for change in changes:
        if change.id not in docs_by_id:
            # we need to capture DocumentMissingError which is not possible in bulk
            #   so let pillow fall back to serial mode to capture the error for missing docs
            retry_changes.add(change)
            continue
        else:
            # set this, so that subsequent doc lookups are avoided
            change.set_document(docs_by_id[change.id])
        try:
            ensure_matched_revisions(change, docs_by_id.get(change.id))
        except DocumentMismatchError:
            retry_changes.add(change)
    return retry_changes, docs
<|endoftext|>


<commit_before>import oauth2_provider.urls
import oidc_provider.urls
from django.conf import settings
from django.conf.urls.static import static
from django.urls import include, path, re_path
from django.contrib import admin
from django.contrib.staticfiles import views as static_views
from django.http import HttpResponse
from django.views.defaults import permission_denied
from rest_framework.routers import SimpleRouter

from devices.api import UserDeviceViewSet
from identities.api import UserIdentityViewSet
from oidc_apis.views import get_api_tokens_view
from tunnistamo import social_auth_urls
from users.views import EmailNeededView, LoginView, LogoutView

from .api import GetJWTView, UserView

def show_login(request):
    html = ""
    if request.user.is_authenticated:
        html += "
%s
" % request.user
        if request.user.ad_groups.exists():
            html += "
AD groups
"
            html += "
"
            for group in request.user.ad_groups.all():
                html += "
%s
" % str(group)
            html += "
"
    else:
        html += "not logged in"
    html += ""
    return HttpResponse(html)

router = SimpleRouter()
router.register('user_identity', UserIdentityViewSet)
router.register('user_device', UserDeviceViewSet)

urlpatterns = [
    path('admin/', admin.site.urls),
    path('api-tokens/', get_api_tokens_view),
    path('accounts/profile/', show_login),
    path('accounts/login/', LoginView.as_view()),
    path('accounts/logout/', LogoutView.as_view()),
    path('accounts/', include(social_auth_urls, namespace='social')),
    path('oauth2/applications/', permission_denied),
    path('oauth2/', include(oauth2_provider.urls, namespace='oauth2_provider')),
    path('openid/', include(oidc_provider.urls, namespace='oidc_provider')),
    re_path(r'^user/(?P<username>[\w.@+-]+)/?$', UserView.as_view()),
    path('user/', UserView.as_view()),
    path('jwt-token/', GetJWTView.as_view()),
    path('login/', LoginView.as_view()),
    path('logout/', LogoutView.as_view()),
    path('email-needed/', EmailNeededView.as_view(), name='email_needed'),
    path('v1/', include((router.urls, 'v1'))),
]

if settings.DEBUG:
    static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)
<commit_msg>Fix flake8 and isort problems<commit_after>import oauth2_provider.urls
import oidc_provider.urls
from django.conf import settings
from django.conf.urls.static import static
from django.contrib import admin
from django.http import HttpResponse
from django.urls import include, path, re_path
from django.views.defaults import permission_denied
from rest_framework.routers import SimpleRouter

from devices.api import UserDeviceViewSet
from identities.api import UserIdentityViewSet
from oidc_apis.views import get_api_tokens_view
from tunnistamo import social_auth_urls
from users.views import EmailNeededView, LoginView, LogoutView

from .api import GetJWTView, UserView

def show_login(request):
    html = ""
    if request.user.is_authenticated:
        html += "
%s
" % request.user
        if request.user.ad_groups.exists():
            html += "
AD groups
"
            html += "
"
            for group in request.user.ad_groups.all():
                html += "
%s
" % str(group)
            html += "
"
    else:
        html += "not logged in"
    html += ""
    return HttpResponse(html)

router = SimpleRouter()
router.register('user_identity', UserIdentityViewSet)
router.register('user_device', UserDeviceViewSet)

urlpatterns = [
    path('admin/', admin.site.urls),
    path('api-tokens/', get_api_tokens_view),
    path('accounts/profile/', show_login),
    path('accounts/login/', LoginView.as_view()),
    path('accounts/logout/', LogoutView.as_view()),
    path('accounts/', include(social_auth_urls, namespace='social')),
    path('oauth2/applications/', permission_denied),
    path('oauth2/', include(oauth2_provider.urls, namespace='oauth2_provider')),
    path('openid/', include(oidc_provider.urls, namespace='oidc_provider')),
    re_path(r'^user/(?P<username>[\w.@+-]+)/?$', UserView.as_view()),
    path('user/', UserView.as_view()),
    path('jwt-token/', GetJWTView.as_view()),
    path('login/', LoginView.as_view()),
    path('logout/', LogoutView.as_view()),
    path('email-needed/', EmailNeededView.as_view(), name='email_needed'),
    path('v1/', include((router.urls, 'v1'))),
]

if settings.DEBUG:
    static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)
<|endoftext|>