<commit_before>#!/usr/bin/env python
import re
import sys

import socket

try:
    from urlparse import urlparse, urlunparse
except ImportError:
    from urllib.parse import urlparse, urlunparse
import logging

try:
    from bs4 import BeautifulSoup
except ImportError:
    sys.stderr.write('You need to get BeautifulSoup installed\n')
    sys.stderr.write('Do it now, as privileged user/root run: pip install beautifulsoup4\now')
    sys.exit(1)

from .proxy import NullProxy, HttpProxy, Socks5Proxy, httplib, socks

__license__ = """
Copyright (c) 2014, {Sandro Gauci|Wendel G. Henrique}
All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * 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.
    * Neither the name of EnableSecurity or Trustwave 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 THE COPYRIGHT HOLDER OR 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.
"""

# unicode mapping borrowed from http://packetstormsecurity.org/web/unicode-fun.txt
# by Gary O'leary-Steele of Sec-1 Ltd
try:
    from urllib import quote, unquote
except ImportError:
    from urllib.parse import quote, unquote
unicodemapping = {' ': '%u0020',
                  '/': '%u2215',
                  '\\': '%u2215',
                  "'": '%u02b9',
                  '"': '%u0022',
                  '>': '%u003e',
                  '<': '%u003c',
                  '#': '%uff03',
                  '!': '%uff01',
                  '$': '%uff04',
                  '*': '%uff0a',
                  '@': '%u0040',
                  '.': '%uff0e',
                  '_': '%uff3f',
                  '(': '%uff08',
                  ')': '%uff09',
                  ',': '%uff0c',
                  '%': '%u0025',
                  '-': '%uff0d',
                  ';': '%uff1b',
                  ':': '%uff1a',
                  '|': '%uff5c',
                  '&': '%uff06',
                  '+': '%uff0b',
                  '=': '%uff1d',
                  'a': '%uff41',
                  'A': '%uff21',
                  'b': '%uff42',
                  'B': '%uff22',
                  'c': '%uff43',
                  'C': '%uff23',
                  'd': '%uff44',
                  'D': '%uff24',
                  'e': '%uff45',
                  'E': '%uff25',
                  'f': '%uff46',
                  'F': '%uff26',
                  'g': '%uff47',
                  'G': '%uff27',
                  'h': '%uff48',
                  'H': '%uff28',
                  'i': '%uff49',
                  'I': '%uff29',
                  'j': '%uff4a',
                  'J': '%uff2a',
                  'k': '%uff4b',
                  'K': '%uff2b',
                  'l': '%uff4c',
                  'L': '%uff2c',
                  'm': '%uff4d',
                  'M': '%uff2d',
                  'n': '%uff4e',
                  'N': '%uff2e',
                  'o': '%uff4f',
                  'O': '%uff2f',
                  'p': '%uff50',
                  'P': '%uff30',
                  'q': '%uff51',
                  'Q': '%uff31',
                  'r': '%uff52',
                  'R': '%uff32',
                  's': '%uff53',
                  'S': '%uff33',
                  't': '%uff54',
                  'T': '%uff34',
                  'u': '%uff55',
                  'U': '%uff35',
                  'v': '%uff56',
                  'V': '%uff36',
                  'w': '%uff57',
                  'W': '%uff37',
                  'x': '%uff58',
                  'X': '%uff38',
                  'y': '%uff59',
                  'Y': '%uff39',
                  'z': '%uff5a',
                  'Z': '%uff3a',
                  '0': '%uff10',
                  '1': '%uff11',
                  '2': '%uff12',
                  '3': '%uff13',
                  '4': '%uff14',
                  '5': '%uff15',
                  '6': '%uff16',
                  '7': '%uff17',
                  '8': '%uff18',
                  '9': '%uff19'}

homoglyphicmapping = {"'": '%ca%bc'}

def oururlparse(target):
    log = logging.getLogger('urlparser')

    ssl = False
    o = urlparse(target)
    if o[0] not in ['http', 'https', '']:
        log.error('scheme %s not supported' % o[0])
        return
    if o[0] == 'https':
        ssl = True
    if len(o[2]) > 0:
        path = o[2]
    else:
        path = '/'
    tmp = o[1].split(':')
    if len(tmp) > 1:
        port = tmp[1]
    else:
        port = None
    hostname = tmp[0]
    query = o[4]
    return (hostname, port, path, query, ssl)

def modifyurl(path, modfunc, log):
    path = path
    log.debug('path is currently %s' % path)
    #s = re.search('(\[.*?\])',path)
    for m in re.findall('(\[.*?\])', path):
        ourstr = m[1:-1]
        newstr = modfunc(ourstr)
        log.debug('String was %s' % ourstr)
        log.debug('String became %s' % newstr)
        path = path.replace(m, newstr)
    log.debug('the path is now %s' % path)
    return path

def modifypath(path, newstrs, log, encode=True):
    log.debug('path is currently %s' % path)
    for m in re.findall('(\[.*?\])', path):
        ourstr = m[1:-1]
        for newstr in newstrs:
            if encode:
                newstr = quote(newstr)
            log.debug('String was %s' % ourstr)
            log.debug('String became %s' % newstr)
            newpath = path.replace(m, newstr).replace(']', '').replace('[', '')
            yield (newpath)

def bruteforceascii(ourstr):
    listourstr = list(ourstr)
    for pos in range(len(ourstr)):
        for i in range(256):
            newlistourstr = listourstr[:]
            newlistourstr[pos] = chr(i)
            yield (quote(''.join(newlistourstr)))

def unicodeurlencode(ourstr):
    newstr = str()
    for character in ourstr:
        if unicodemapping.has_key(character):
            newstr += unicodemapping[character]
        else:
            newstr += character
    return newstr

def nullify(ourstr):
    newstr = str()
    for character in ourstr:
        newstr += character + "\x00"
    return quote(newstr)

def replacechars(ourstr, origchar, newchar):
    newstr = ourstr.replace(origchar, newchar)
    return newstr

def nullifyspaces(ourstr):
    return quote(replacechars(ourstr, ' ', '\x00'))

def slashspaces(ourstr):
    return replacechars(ourstr, ' ', '/')

def tabifyspaces(ourstr):
    return replacechars(ourstr, ' ', '\t')

def crlfspaces(ourstr):
    return replacechars(ourstr, ' ', '\n')

def backslashquotes(ourstr):
    return replacechars(ourstr, "'", "\''")

class waftoolsengine:
    def __init__(self, target='www.microsoft.com', port=80, ssl=False,
                 debuglevel=0, path='/', followredirect=True, extraheaders={},
                 proxy=False):
        """
        target: the hostname or ip of the target server
        port: defaults to 80
        ssl: defaults to false
        """
        self.target = target
        if port is None:
            if ssl:
                port = 443
            else:
                port = 80
        self.port = int(port)
        self.ssl = ssl
        self.debuglevel = debuglevel
        self.cachedresponses = dict()
        self.requestnumber = 0
        self.path = path
        self.redirectno = 0
        self.followredirect = followredirect
        self.crawlpaths = list()
        self.extraheaders = extraheaders

        try:
            self.proxy = self._parse_proxy(proxy) if proxy else NullProxy()
        except Exception as e:
            self.log.critical("Proxy disabled: %s" % e)
            self.proxy = NullProxy()

    def request(self, method='GET', path=None, usecache=True,
                cacheresponse=True, headers=None,
                comingfromredir=False):
        followredirect = self.followredirect
        if comingfromredir:
            self.redirectno += 1
            if self.redirectno >= 5:
                self.log.error('We received way too many redirects.. stopping that')
                followredirect = False
        else:
            self.redirectno = 0
        if path is None:
            path = self.path
        for hdr in self.extraheaders.keys():
            if headers is None:
                headers = {}
            headers[hdr] = self.extraheaders[hdr]
        if headers is not None:
            knownheaders = map(lambda x: x.lower(), headers.keys())
        else:
            knownheaders = {}
            headers = {}
        if not 'user-agent' in knownheaders:
            headers[
                'User-Agent'] = 'Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.5; en-US; rv:1.9.1b1) Gecko/20081007 Firefox/3.0'
        if not 'accept-charset' in knownheaders:
            headers['Accept-Charset'] = 'ISO-8859-1,utf-8;q=0.7,*;q=0.7'
        if not 'accept' in knownheaders:
            headers['Accept'] = '*/*'
        k = str([method, path, headers])
        if usecache:
            if k in self.cachedresponses.keys():
                self.log.debug('Using cached version of %s, %s' % (method, path))
                return self.cachedresponses[k]
            else:
                self.log.debug('%s not found in %s' % (k, self.cachedresponses.keys()))
        r = self._request(method, path, headers)
        if cacheresponse:
            self.cachedresponses[k] = r

        if r:
            response, responsebody = r
            if response.status in [301, 302, 307]:
                if followredirect:
                    if response.getheader('location'):
                        newloc = response.getheader('location')
                        self.log.info('Redirected to %s' % newloc)
                        pret = oururlparse(newloc)
                        if pret is not None:
                            (target, port, path, query, ssl) = pret
                            if not port: port = 80
                            if target == '':
                                target = self.target
                            if port is None:
                                port = self.port
                            if not path.startswith('/'):
                                path = '/' + path
                            if (target, port, ssl) == (self.target, self.port, ssl):
                                r = self.request(method, path, usecache, cacheresponse,
                                                 headers, comingfromredir=True)
                            else:
                                self.log.warn('Tried to redirect to a different server %s' % newloc)
                        else:
                            self.log.warn('%s is not a well formatted url' % response.getheader('location'))
        return r

    def _request(self, method, path, headers):
        original_socket = socket.socket
        try:
            conn_factory, connect_host, connect_port, query_path = \
                    self.proxy.prepare(self.target, self.port, path, self.ssl)

            if sys.hexversion > 0x2060000:
                h = conn_factory(connect_host, connect_port, timeout=4)
            else:
                h = conn_factory(connect_host, connect_port)

            if self.debuglevel <= 10:
                if self.debuglevel > 1:
                    h.set_debuglevel(self.debuglevel)
            try:
                self.log.info('Sending %s %s' % (method, path))
                h.request(method, query_path, headers=headers)
            except socket.error:
                self.log.warn('Could not initialize connection to %s' % self.target)
                return
            self.requestnumber += 1

            response = h.getresponse()
            responsebody = response.read()
            h.close()
            r = response, responsebody
        except (socket.error, socket.timeout, httplib.BadStatusLine):
            self.log.warn('Hey.. they closed our connection!')
            r = None
        finally:
            self.proxy.terminate()

        return r

    def querycrawler(self, path=None, curdepth=0, maxdepth=1):
        self.log.debug('Crawler is visiting %s' % path)
        localcrawlpaths = list()
        if curdepth > maxdepth:
            self.log.info('maximum depth %s reached' % maxdepth)
            return
        r = self.request(path=path)
        if r is None:
            return
        response, responsebody = r
        try:
            soup = BeautifulSoup(responsebody)
        except:
            self.log.warn('could not parse the response body')
            return
        tags = soup('a')
        for tag in tags:
            try:
                href = tag["href"]
                if href is not None:
                    tmpu = urlparse(href)
                    if (tmpu[1] != '') and (self.target != tmpu[1]):
                        # not on the same domain name .. ignore
                        self.log.debug('Ignoring link because it is not on the same site %s' % href)
                        continue
                    if tmpu[0] not in ['http', 'https', '']:
                        self.log.debug('Ignoring link because it is not an http uri %s' % href)
                        continue
                    path = tmpu[2]
                    if not path.startswith('/'):
                        path = '/' + path
                    if len(tmpu[4]) > 0:
                        # found a query .. thats all we need
                        location = urlunparse(('', '', path, tmpu[3], tmpu[4], ''))
                        self.log.info('Found query %s' % location)
                        return href
                    if path not in self.crawlpaths:
                        href = unquote(path)
                        self.log.debug('adding %s for crawling' % href)
                        self.crawlpaths.append(href)
                        localcrawlpaths.append(href)
            except KeyError:
                pass
        for nextpath in localcrawlpaths:
            r = self.querycrawler(path=nextpath, curdepth=curdepth + 1, maxdepth=maxdepth)
            if r:
                return r

    def _parse_proxy(self, proxy):
        parts = urlparse(proxy)
        if not parts.scheme or not parts.netloc:
            raise Exception("Invalid proxy specified, scheme required")

        netloc = parts.netloc.split(":")
        if len(netloc) != 2:
            raise Exception("Proxy port unspecified")

        try:
            if parts.scheme == "socks5":
                if socks is None:
                    raise Exception("socks5 proxy requires PySocks")

                return Socks5Proxy(netloc[0], int(netloc[1]))
            elif parts.scheme == "http":
                return HttpProxy(netloc[0], int(netloc[1]))
            else:
                raise Exception("Unsupported proxy scheme")
        except ValueError:
            raise Exception("Invalid port number")

def scrambledheader(header):
    c = 'connection'
    if len(header) != len(c):
        return False
    if header == c:
        return False
    for character in c:
        if c.count(character) != header.count(character):
            return False
    return True

<commit_msg>reapplied disabling of ssl certificate check<commit_after>#!/usr/bin/env python
import re
import sys

import socket

try:
    from urlparse import urlparse, urlunparse
except ImportError:
    from urllib.parse import urlparse, urlunparse
import logging

try:
    from bs4 import BeautifulSoup
except ImportError:
    sys.stderr.write('You need to get BeautifulSoup installed\n')
    sys.stderr.write('Do it now, as privileged user/root run: pip install beautifulsoup4\now')
    sys.exit(1)

from .proxy import NullProxy, HttpProxy, Socks5Proxy, httplib, socks

__license__ = """
Copyright (c) 2014, {Sandro Gauci|Wendel G. Henrique}
All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * 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.
    * Neither the name of EnableSecurity or Trustwave 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 THE COPYRIGHT HOLDER OR 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.
"""

# unicode mapping borrowed from http://packetstormsecurity.org/web/unicode-fun.txt
# by Gary O'leary-Steele of Sec-1 Ltd
try:
    from urllib import quote, unquote
except ImportError:
    from urllib.parse import quote, unquote
unicodemapping = {' ': '%u0020',
                  '/': '%u2215',
                  '\\': '%u2215',
                  "'": '%u02b9',
                  '"': '%u0022',
                  '>': '%u003e',
                  '<': '%u003c',
                  '#': '%uff03',
                  '!': '%uff01',
                  '$': '%uff04',
                  '*': '%uff0a',
                  '@': '%u0040',
                  '.': '%uff0e',
                  '_': '%uff3f',
                  '(': '%uff08',
                  ')': '%uff09',
                  ',': '%uff0c',
                  '%': '%u0025',
                  '-': '%uff0d',
                  ';': '%uff1b',
                  ':': '%uff1a',
                  '|': '%uff5c',
                  '&': '%uff06',
                  '+': '%uff0b',
                  '=': '%uff1d',
                  'a': '%uff41',
                  'A': '%uff21',
                  'b': '%uff42',
                  'B': '%uff22',
                  'c': '%uff43',
                  'C': '%uff23',
                  'd': '%uff44',
                  'D': '%uff24',
                  'e': '%uff45',
                  'E': '%uff25',
                  'f': '%uff46',
                  'F': '%uff26',
                  'g': '%uff47',
                  'G': '%uff27',
                  'h': '%uff48',
                  'H': '%uff28',
                  'i': '%uff49',
                  'I': '%uff29',
                  'j': '%uff4a',
                  'J': '%uff2a',
                  'k': '%uff4b',
                  'K': '%uff2b',
                  'l': '%uff4c',
                  'L': '%uff2c',
                  'm': '%uff4d',
                  'M': '%uff2d',
                  'n': '%uff4e',
                  'N': '%uff2e',
                  'o': '%uff4f',
                  'O': '%uff2f',
                  'p': '%uff50',
                  'P': '%uff30',
                  'q': '%uff51',
                  'Q': '%uff31',
                  'r': '%uff52',
                  'R': '%uff32',
                  's': '%uff53',
                  'S': '%uff33',
                  't': '%uff54',
                  'T': '%uff34',
                  'u': '%uff55',
                  'U': '%uff35',
                  'v': '%uff56',
                  'V': '%uff36',
                  'w': '%uff57',
                  'W': '%uff37',
                  'x': '%uff58',
                  'X': '%uff38',
                  'y': '%uff59',
                  'Y': '%uff39',
                  'z': '%uff5a',
                  'Z': '%uff3a',
                  '0': '%uff10',
                  '1': '%uff11',
                  '2': '%uff12',
                  '3': '%uff13',
                  '4': '%uff14',
                  '5': '%uff15',
                  '6': '%uff16',
                  '7': '%uff17',
                  '8': '%uff18',
                  '9': '%uff19'}

homoglyphicmapping = {"'": '%ca%bc'}

def oururlparse(target):
    log = logging.getLogger('urlparser')

    ssl = False
    o = urlparse(target)
    if o[0] not in ['http', 'https', '']:
        log.error('scheme %s not supported' % o[0])
        return
    if o[0] == 'https':
        ssl = True
    if len(o[2]) > 0:
        path = o[2]
    else:
        path = '/'
    tmp = o[1].split(':')
    if len(tmp) > 1:
        port = tmp[1]
    else:
        port = None
    hostname = tmp[0]
    query = o[4]
    return (hostname, port, path, query, ssl)

def modifyurl(path, modfunc, log):
    path = path
    log.debug('path is currently %s' % path)
    #s = re.search('(\[.*?\])',path)
    for m in re.findall('(\[.*?\])', path):
        ourstr = m[1:-1]
        newstr = modfunc(ourstr)
        log.debug('String was %s' % ourstr)
        log.debug('String became %s' % newstr)
        path = path.replace(m, newstr)
    log.debug('the path is now %s' % path)
    return path

def modifypath(path, newstrs, log, encode=True):
    log.debug('path is currently %s' % path)
    for m in re.findall('(\[.*?\])', path):
        ourstr = m[1:-1]
        for newstr in newstrs:
            if encode:
                newstr = quote(newstr)
            log.debug('String was %s' % ourstr)
            log.debug('String became %s' % newstr)
            newpath = path.replace(m, newstr).replace(']', '').replace('[', '')
            yield (newpath)

def bruteforceascii(ourstr):
    listourstr = list(ourstr)
    for pos in range(len(ourstr)):
        for i in range(256):
            newlistourstr = listourstr[:]
            newlistourstr[pos] = chr(i)
            yield (quote(''.join(newlistourstr)))

def unicodeurlencode(ourstr):
    newstr = str()
    for character in ourstr:
        if unicodemapping.has_key(character):
            newstr += unicodemapping[character]
        else:
            newstr += character
    return newstr

def nullify(ourstr):
    newstr = str()
    for character in ourstr:
        newstr += character + "\x00"
    return quote(newstr)

def replacechars(ourstr, origchar, newchar):
    newstr = ourstr.replace(origchar, newchar)
    return newstr

def nullifyspaces(ourstr):
    return quote(replacechars(ourstr, ' ', '\x00'))

def slashspaces(ourstr):
    return replacechars(ourstr, ' ', '/')

def tabifyspaces(ourstr):
    return replacechars(ourstr, ' ', '\t')

def crlfspaces(ourstr):
    return replacechars(ourstr, ' ', '\n')

def backslashquotes(ourstr):
    return replacechars(ourstr, "'", "\''")

class waftoolsengine:
    def __init__(self, target='www.microsoft.com', port=80, ssl=False,
                 debuglevel=0, path='/', followredirect=True, extraheaders={},
                 proxy=False):
        """
        target: the hostname or ip of the target server
        port: defaults to 80
        ssl: defaults to false
        """
        self.target = target
        if port is None:
            if ssl:
                port = 443
            else:
                port = 80
        self.port = int(port)
        self.ssl = ssl
        self.debuglevel = debuglevel
        self.cachedresponses = dict()
        self.requestnumber = 0
        self.path = path
        self.redirectno = 0
        self.followredirect = followredirect
        self.crawlpaths = list()
        self.extraheaders = extraheaders

        try:
            self.proxy = self._parse_proxy(proxy) if proxy else NullProxy()
        except Exception as e:
            self.log.critical("Proxy disabled: %s" % e)
            self.proxy = NullProxy()

    def request(self, method='GET', path=None, usecache=True,
                cacheresponse=True, headers=None,
                comingfromredir=False):
        followredirect = self.followredirect
        if comingfromredir:
            self.redirectno += 1
            if self.redirectno >= 5:
                self.log.error('We received way too many redirects.. stopping that')
                followredirect = False
        else:
            self.redirectno = 0
        if path is None:
            path = self.path
        for hdr in self.extraheaders.keys():
            if headers is None:
                headers = {}
            headers[hdr] = self.extraheaders[hdr]
        if headers is not None:
            knownheaders = map(lambda x: x.lower(), headers.keys())
        else:
            knownheaders = {}
            headers = {}
        if not 'user-agent' in knownheaders:
            headers[
                'User-Agent'] = 'Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.5; en-US; rv:1.9.1b1) Gecko/20081007 Firefox/3.0'
        if not 'accept-charset' in knownheaders:
            headers['Accept-Charset'] = 'ISO-8859-1,utf-8;q=0.7,*;q=0.7'
        if not 'accept' in knownheaders:
            headers['Accept'] = '*/*'
        k = str([method, path, headers])
        if usecache:
            if k in self.cachedresponses.keys():
                self.log.debug('Using cached version of %s, %s' % (method, path))
                return self.cachedresponses[k]
            else:
                self.log.debug('%s not found in %s' % (k, self.cachedresponses.keys()))
        r = self._request(method, path, headers)
        if cacheresponse:
            self.cachedresponses[k] = r

        if r:
            response, responsebody = r
            if response.status in [301, 302, 307]:
                if followredirect:
                    if response.getheader('location'):
                        newloc = response.getheader('location')
                        self.log.info('Redirected to %s' % newloc)
                        pret = oururlparse(newloc)
                        if pret is not None:
                            (target, port, path, query, ssl) = pret
                            if not port: port = 80
                            if target == '':
                                target = self.target
                            if port is None:
                                port = self.port
                            if not path.startswith('/'):
                                path = '/' + path
                            if (target, port, ssl) == (self.target, self.port, ssl):
                                r = self.request(method, path, usecache, cacheresponse,
                                                 headers, comingfromredir=True)
                            else:
                                self.log.warn('Tried to redirect to a different server %s' % newloc)
                        else:
                            self.log.warn('%s is not a well formatted url' % response.getheader('location'))
        return r

    def _request(self, method, path, headers):
        original_socket = socket.socket
        try:
            conn_factory, connect_host, connect_port, query_path = \
                    self.proxy.prepare(self.target, self.port, path, self.ssl)

            params = dict()
            if sys.hexversion > 0x2060000:
                params['timeout'] = 4
            if (sys.hexversion >= 0x2070900) and self.ssl:
                import ssl as ssllib
                params['context'] = ssllib._create_unverified_context()
            h = conn_factory(connect_host, connect_port,**params)
            if self.debuglevel <= 10:
                if self.debuglevel > 1:
                    h.set_debuglevel(self.debuglevel)
            try:
                self.log.info('Sending %s %s' % (method, path))
                h.request(method, query_path, headers=headers)
            except socket.error:
                self.log.warn('Could not initialize connection to %s' % self.target)
                return
            self.requestnumber += 1

            response = h.getresponse()
            responsebody = response.read()
            h.close()
            r = response, responsebody
        except (socket.error, socket.timeout, httplib.BadStatusLine):
            self.log.warn('Hey.. they closed our connection!')
            r = None
        finally:
            self.proxy.terminate()

        return r

    def querycrawler(self, path=None, curdepth=0, maxdepth=1):
        self.log.debug('Crawler is visiting %s' % path)
        localcrawlpaths = list()
        if curdepth > maxdepth:
            self.log.info('maximum depth %s reached' % maxdepth)
            return
        r = self.request(path=path)
        if r is None:
            return
        response, responsebody = r
        try:
            soup = BeautifulSoup(responsebody)
        except:
            self.log.warn('could not parse the response body')
            return
        tags = soup('a')
        for tag in tags:
            try:
                href = tag["href"]
                if href is not None:
                    tmpu = urlparse(href)
                    if (tmpu[1] != '') and (self.target != tmpu[1]):
                        # not on the same domain name .. ignore
                        self.log.debug('Ignoring link because it is not on the same site %s' % href)
                        continue
                    if tmpu[0] not in ['http', 'https', '']:
                        self.log.debug('Ignoring link because it is not an http uri %s' % href)
                        continue
                    path = tmpu[2]
                    if not path.startswith('/'):
                        path = '/' + path
                    if len(tmpu[4]) > 0:
                        # found a query .. thats all we need
                        location = urlunparse(('', '', path, tmpu[3], tmpu[4], ''))
                        self.log.info('Found query %s' % location)
                        return href
                    if path not in self.crawlpaths:
                        href = unquote(path)
                        self.log.debug('adding %s for crawling' % href)
                        self.crawlpaths.append(href)
                        localcrawlpaths.append(href)
            except KeyError:
                pass
        for nextpath in localcrawlpaths:
            r = self.querycrawler(path=nextpath, curdepth=curdepth + 1, maxdepth=maxdepth)
            if r:
                return r

    def _parse_proxy(self, proxy):
        parts = urlparse(proxy)
        if not parts.scheme or not parts.netloc:
            raise Exception("Invalid proxy specified, scheme required")

        netloc = parts.netloc.split(":")
        if len(netloc) != 2:
            raise Exception("Proxy port unspecified")

        try:
            if parts.scheme == "socks5":
                if socks is None:
                    raise Exception("socks5 proxy requires PySocks")

                return Socks5Proxy(netloc[0], int(netloc[1]))
            elif parts.scheme == "http":
                return HttpProxy(netloc[0], int(netloc[1]))
            else:
                raise Exception("Unsupported proxy scheme")
        except ValueError:
            raise Exception("Invalid port number")

def scrambledheader(header):
    c = 'connection'
    if len(header) != len(c):
        return False
    if header == c:
        return False
    for character in c:
        if c.count(character) != header.count(character):
            return False
    return True

<|endoftext|>


<commit_before># coding: utf-8
from django.conf.urls.defaults import *
from regressiontests.templates import views

urlpatterns = patterns('',

    # Test urls for testing reverse lookups
    (r'^$', views.index),
    (r'^client/(\d+)/$', views.client),
    (r'^client/(\d+)/(?P<action>[^/]+)/$', views.client_action),
    url(r'^named-client/(\d+)/$', views.client, name="named.client"),

    # Unicode strings are permitted everywhere.
    url(ur'^Юникод/(\w+)/$', views.client, name=u"метка_оператора"),
)
<commit_msg>Fixed an invalid URL specification.<commit_after># coding: utf-8
from django.conf.urls.defaults import *
from regressiontests.templates import views

urlpatterns = patterns('',

    # Test urls for testing reverse lookups
    (r'^$', views.index),
    (r'^client/(\d+)/$', views.client),
    (r'^client/(?P<id>\d+)/(?P<action>[^/]+)/$', views.client_action),
    url(r'^named-client/(\d+)/$', views.client, name="named.client"),

    # Unicode strings are permitted everywhere.
    url(ur'^Юникод/(\w+)/$', views.client, name=u"метка_оператора"),
)
<|endoftext|>


<commit_before>import numpy as np
import pytest

from simulators.fake_simulator import fake_bhm_simulation, fake_iam_simulation
from simulators.fake_simulator import main as fake_main
from simulators.fake_simulator import parse_args

def test_fake_simulator_parser():
    args = ["HD30501", "01", "-p", "2300, 4.5, -3.0", "--params2", "2100, 3.5, 0.0", "-g", "10"]
    parsed = parse_args(args)
    assert parsed.star == "HD30501"
    assert parsed.sim_num == "01"
    assert parsed.params1 == "2300, 4.5, -3.0"
    assert parsed.params2 == "2100, 3.5, 0.0"
    assert parsed.gamma == 10
    assert parsed.rv is None
    assert isinstance(parsed.gamma, float)
    assert parsed.replace is False
    assert parsed.noplots is False
    assert parsed.test is False
    assert parsed.mode == "iam"
    assert parsed.noise is None

def test_fake_simulator_parser_toggle():
    args = ["HDTEST", "02", "-t", "-r", '-v', "10", "-m", "bhm", "-s", "100", "-n"]
    parsed = parse_args(args)
    assert parsed.star == "HDTEST"
    assert parsed.sim_num == "02"
    assert parsed.params1 is None
    assert parsed.params2 is None
    assert parsed.rv == 10
    assert parsed.gamma is None
    assert parsed.replace is True
    assert parsed.noplots is True
    assert parsed.test is True
    assert parsed.mode == "bhm"
    assert parsed.noise == "100"

def test_fake_sim_main_with_no_params1_returns_error():
    with pytest.raises(ValueError):
        fake_main("hdtest", 1, params1=None, params2=[5800, 4.0, -0.5], rv=7, gamma=5, mode="iam")

    with pytest.raises(ValueError):
        fake_main("hdtest2", 2, params1=None, gamma=7, mode="bmh")

@pytest.mark.xfail()
def test_fake_simulator_main():
    fake_main("teststar", 1, "4500, 4.0, -1.0", "2300, 4.5, 0.0")

    assert False

@pytest.mark.xfail()
@pytest.mark.parametrize("params", [(2500, 4.5, 0.0), (2800, 4.5, 0.5)])
@pytest.mark.parametrize("wav", [np.linspace(2130, 2145, 40), np.linspace(2147, 2160, 200)])
@pytest.mark.parametrize("gamma", [-5, 1, 7])
@pytest.mark.parametrize("limits", [[2070, 2180]])
def test_fake_iam_simulation_with_wav(params, wav, gamma, limits):
    fake_wav, fake_flux = fake_iam_simulation(wav, [5000, 4.5, 0.5], params2=params,
                                              rv=-rv, gamma=gamma, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert np.all(fake_wav == wav)
    assert fake_wav.shape == fake_flux.shape

@pytest.mark.xfail()
@pytest.mark.parametrize("params", [(2500, 4.5, 0.0), (2800, 4.5, 0.5)])
@pytest.mark.parametrize("limits", [[2030, 2180], [2100, 2140]])
@pytest.mark.parametrize("independant", [True, False])
def test_fake_iam_simulation_without_wav(params, limits, independant):
    fake_wav, fake_flux = fake_iam_simulation(None, [58000, 4.0, -0.5], params2=params, rv=7,
                                              gamma=5, limits=limits, independent=independant)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert fake_wav.shape == fake_flux.shape

# @pytest.mark.xfail()
@pytest.mark.parametrize("wav", [np.linspace(2130, 2145, 40), np.linspace(2147, 2160, 200)])
@pytest.mark.parametrize("gamma", [-5, 1, 7])
@pytest.mark.parametrize("limits", [[2070, 2180]])
def test_fake_bhm_simulation_with_wav(wav, gamma, limits):
    fake_wav, fake_flux = fake_bhm_simulation(wav, [5000, 4.5, 0.5], gamma=gamma, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert np.all(fake_wav == wav)
    assert fake_wav.shape == fake_flux.shape

@pytest.mark.xfail()
@pytest.mark.parametrize("limits", [[2030, 2180], [2100, 2140]])
def test_fake_bhm_simulation_without_wav(limits):
    fake_wav, fake_flux = fake_bhm_simulation(None, [4000, 4.0, -0.5], gamma=5, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert fake_wav.shape == fake_flux.shape
<commit_msg>Fix a fake simualtor test. missing parametrize value.<commit_after>import numpy as np
import pytest

from simulators.fake_simulator import fake_bhm_simulation, fake_iam_simulation
from simulators.fake_simulator import main as fake_main
from simulators.fake_simulator import parse_args

def test_fake_simulator_parser():
    args = ["HD30501", "01", "-p", "2300, 4.5, -3.0", "--params2", "2100, 3.5, 0.0", "-g", "10"]
    parsed = parse_args(args)
    assert parsed.star == "HD30501"
    assert parsed.sim_num == "01"
    assert parsed.params1 == "2300, 4.5, -3.0"
    assert parsed.params2 == "2100, 3.5, 0.0"
    assert parsed.gamma == 10
    assert parsed.rv is None
    assert isinstance(parsed.gamma, float)
    assert parsed.replace is False
    assert parsed.noplots is False
    assert parsed.test is False
    assert parsed.mode == "iam"
    assert parsed.noise is None

def test_fake_simulator_parser_toggle():
    args = ["HDTEST", "02", "-t", "-r", '-v', "10", "-m", "bhm", "-s", "100", "-n"]
    parsed = parse_args(args)
    assert parsed.star == "HDTEST"
    assert parsed.sim_num == "02"
    assert parsed.params1 is None
    assert parsed.params2 is None
    assert parsed.rv == 10
    assert parsed.gamma is None
    assert parsed.replace is True
    assert parsed.noplots is True
    assert parsed.test is True
    assert parsed.mode == "bhm"
    assert parsed.noise == "100"

def test_fake_sim_main_with_no_params1_returns_error():
    with pytest.raises(ValueError):
        fake_main("hdtest", 1, params1=None, params2=[5800, 4.0, -0.5], rv=7, gamma=5, mode="iam")

    with pytest.raises(ValueError):
        fake_main("hdtest2", 2, params1=None, gamma=7, mode="bmh")

@pytest.mark.xfail()
def test_fake_simulator_main():
    fake_main("teststar", 1, "4500, 4.0, -1.0", "2300, 4.5, 0.0")

    assert False

# @pytest.mark.xfail()
@pytest.mark.parametrize("params", [(2500, 4.5, 0.0), (2800, 4.5, 0.5)])
@pytest.mark.parametrize("wav", [np.linspace(2130, 2145, 40), np.linspace(2147, 2160, 200)])
@pytest.mark.parametrize("rv", [5, 2, -6])
@pytest.mark.parametrize("gamma", [-5, 1, 7])
@pytest.mark.parametrize("limits", [[2070, 2180]])
def test_fake_iam_simulation_with_wav(params, wav, rv, gamma, limits):
    fake_wav, fake_flux = fake_iam_simulation(wav, [5000, 4.5, 0.5], params2=params,
                                              rv=-rv, gamma=gamma, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert np.all(fake_wav == wav)
    assert fake_wav.shape == fake_flux.shape

@pytest.mark.xfail()
@pytest.mark.parametrize("params", [(2500, 4.5, 0.0), (2800, 4.5, 0.5)])
@pytest.mark.parametrize("limits", [[2030, 2180], [2100, 2140]])
@pytest.mark.parametrize("independant", [True, False])
def test_fake_iam_simulation_without_wav(params, limits, independant):
    fake_wav, fake_flux = fake_iam_simulation(None, [58000, 4.0, -0.5], params2=params, rv=7,
                                              gamma=5, limits=limits, independent=independant)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert fake_wav.shape == fake_flux.shape

# @pytest.mark.xfail()
@pytest.mark.parametrize("wav", [np.linspace(2130, 2145, 40), np.linspace(2147, 2160, 200)])
@pytest.mark.parametrize("gamma", [-5, 1, 7])
@pytest.mark.parametrize("limits", [[2070, 2180]])
def test_fake_bhm_simulation_with_wav(wav, gamma, limits):
    fake_wav, fake_flux = fake_bhm_simulation(wav, [5000, 4.5, 0.5], gamma=gamma, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert np.all(fake_wav == wav)
    assert fake_wav.shape == fake_flux.shape

# @pytest.mark.xfail()
@pytest.mark.parametrize("limits", [[2030, 2180], [2100, 2140]])
def test_fake_bhm_simulation_without_wav(limits):
    fake_wav, fake_flux = fake_bhm_simulation(None, [4000, 4.0, -0.5], gamma=5, limits=limits)
    assert np.all(fake_wav < limits[1]) and np.all(fake_wav > limits[0])
    assert fake_wav.shape == fake_flux.shape
<|endoftext|>


<commit_before>from .abstract import AbstractTensor

class PointerTensor(AbstractTensor):
    """PointerTensor objects point to tensors (as their name implies). They exist to mimic
    the entire API of a normal tensor, but instead of computing a tensor function locally
    (such as addition, subtraction, etc.) they forward the computation to a remote machine
    as specified by self.location. Specifically, every PointerTensor has a tensor
    located somewhere that it points to (they should never exist by themselves). Note that
    PointerTensor objects can point to both torch.Tensor objects AND to other PointerTensor
    objects. Furthermore, the objects being pointed to can be on the same machine or
    (more commonly) on a different one. Note further that a PointerTensor does not know
    the nature how it sends messages to the tensor it points to (whether over socket, http,
    or some other protocol) as that functionality is abstracted in the BaseWorker object in
    self.location.

    :Example:

     >>> import syft as sy
     >>> hook = sy.TorchHook()

     >>> bob = sy.VirtualWorker(id="bob")

     >>> x = sy.Tensor([1,2,3,4,5])
     >>> y = sy.Tensor([1,1,1,1,1])

     >>> x_ptr = x.send(bob) # returns a PointerTensor, sends tensor to Bob
     >>> y_ptr = y.send(bob) # returns a PointerTensor, sends tensor to Bob

     >>> # executes command on Bob's machine
     >>> z_ptr = x_ptr + y_ptr

    """

    def __init__(
        self,
        parent: AbstractTensor = None,
        location=None,
        id_at_location=None,
        register=False,
        owner=None,
        id=None,
    ):
        """Initializes a PointerTensor

        :Parameters:

         * **parent (**:class:`.abstract.AbstractTensor` **, optional)** this is
            typically just the wrapper around the PointerTensor which makes it
            so that you can pass this PointerTensor to all the other methods/
            functions that PyTorch likes to use, although it can also be other
            tensors which extend AbstractTensor, such as custom tensors for
            Secure Multi-Party Computation or Federated Learning

         * **location (**:class:`....workers.BaseWorker` **, optional)** the
            BaseWorker object which points to the worker on which this
            pointer's object can be found.

         * **id_at_location ((str or int), optional)** - the id of the tensor
            being pointed to.

         * **register (bool, optional)** this parameter determines whether
            to automatically register the new pointer that gets created.

         * **owner (**:class:`....workers.BaseWorker` **, optional)** while
            "location" specifies where the pointer points to, this parameter
            specifies the worker on which the pointer is located. It is also
            where the pointer is registered if register is set to True.

         * **id: ((str or int), optional)** - this is the id of the PointerTensor

        """
        # TODO: finish AbstractTensor which should handle register/parent/id stuff

        self.location = location
        self.id_at_location = id_at_location
        self.owner = owner
        self.id = id

    def __str__(self):
        """This returns a string version of this pointer, which is primarily for
        end users to quickly see things about the tensor. This tostring shouldn't
        be used for anything else though as it's likely to change. (aka, don't try
        to parse it to extract information. Read the attribute you need directly).
        Also, don't use this to-string as a serialized form of the pointer."""

        type_name = type(self).__name__
        return (
            f"["
            f"{type_name} - "
            f"id:{self.id} "
            f"owner:{self.owner.id} "
            f"loc:{self.location.id} "
            f"id@loc:{self.id_at_location}"
            f"]"
        )

    def __repr__(self):
        """This returns the to-string method when called using __repr__, most
        commonly seen when returned as cells in Jupyter noteoboks."""
        return self.__str__()

    def get(self, deregister_ptr: bool = True):
        """Since PointerTensor objects always point to a remote tensor (or chain
        of tensors, where a chain is simply a linked-list of tensors linked via
        their .child attributes), this method will request that the tensor/chain
        being pointed to be serialized and returned from this function. This will
        typically mean that the remote object will be removed/destroyed. If you
        merely wish to bring a copy back to the local worker, call .copy() before
        calling .get(). TODO: add param get_copy which doesn't destroy remote if
        true.

        :Parameters:

         * **deregister_ptr (bool, optional)** this determines whether to deregister
            this pointer from the pointer's owner during this method. This defaults
            to True because the main reason people use this method is to move the
            tensor from the remote machine to the local one, at which time the
            pointer has no use.

         * **out (**:class:`.abstract.AbstractTensor` **)** this is the tensor
            (or chain) which this object used to point to on a remote machine.

        """

        # if the pointer happens to be pointing to a local object,
        # just return that object (this is an edge case)
        if self.location == self.owner:
            tensor = self.owner.get_obj(self.id_at_location).child
        else:
            # get tensor from remote machine
            tensor = self.owner.request_obj(self.id_at_location, self.location)

        # Register the result
        assigned_id = self.id_at_location
        self.owner.register_obj(tensor, assigned_id)

        # Remove this pointer by default
        if deregister_ptr:
            self.owner.de_register_obj(self)

        return tensor
<commit_msg>updated comments<commit_after>from .abstract import AbstractTensor

class PointerTensor(AbstractTensor):
    """A pointer to another tensor. A PointerTensor forwards all API calls to the remote.

    PointerTensor objects point to tensors (as their name implies). They exist to mimic
    the entire API of a normal tensor, but instead of computing a tensor function locally
    (such as addition, subtraction, etc.) they forward the computation to a remote machine
    as specified by self.location. Specifically, every PointerTensor has a tensor
    located somewhere that it points to (they should never exist by themselves). Note that
    PointerTensor objects can point to both torch.Tensor objects AND to other PointerTensor
    objects. Furthermore, the objects being pointed to can be on the same machine or
    (more commonly) on a different one. Note further that a PointerTensor does not know
    the nature how it sends messages to the tensor it points to (whether over socket, http,
    or some other protocol) as that functionality is abstracted in the BaseWorker object in
    self.location.

    :Example:

     >>> import syft as sy
     >>> hook = sy.TorchHook()

     >>> bob = sy.VirtualWorker(id="bob")

     >>> x = sy.Tensor([1,2,3,4,5])
     >>> y = sy.Tensor([1,1,1,1,1])

     >>> x_ptr = x.send(bob) # returns a PointerTensor, sends tensor to Bob
     >>> y_ptr = y.send(bob) # returns a PointerTensor, sends tensor to Bob

     >>> # executes command on Bob's machine
     >>> z_ptr = x_ptr + y_ptr

    """

    def __init__(
        self,
        parent: AbstractTensor = None,
        location=None,
        id_at_location=None,
        register=False,
        owner=None,
        id=None,
    ):
        """Initializes a PointerTensor

        :Parameters:

         * **parent (**:class:`.abstract.AbstractTensor` **, optional)** this is
            typically just the wrapper around the PointerTensor which makes it
            so that you can pass this PointerTensor to all the other methods/
            functions that PyTorch likes to use, although it can also be other
            tensors which extend AbstractTensor, such as custom tensors for
            Secure Multi-Party Computation or Federated Learning

         * **location (**:class:`....workers.BaseWorker` **, optional)** the
            BaseWorker object which points to the worker on which this
            pointer's object can be found.

         * **id_at_location ((str or int), optional)** - the id of the tensor
            being pointed to.

         * **register (bool, optional)** this parameter determines whether
            to automatically register the new pointer that gets created.

         * **owner (**:class:`....workers.BaseWorker` **, optional)** while
            "location" specifies where the pointer points to, this parameter
            specifies the worker on which the pointer is located. It is also
            where the pointer is registered if register is set to True.

         * **id: ((str or int), optional)** - this is the id of the PointerTensor

        """
        # TODO: finish AbstractTensor which should handle register/parent/id stuff

        self.location = location
        self.id_at_location = id_at_location
        self.owner = owner
        self.id = id

    def __str__(self):
        """This returns a string version of this pointer, which is primarily for
        end users to quickly see things about the tensor. This tostring shouldn't
        be used for anything else though as it's likely to change. (aka, don't try
        to parse it to extract information. Read the attribute you need directly).
        Also, don't use this to-string as a serialized form of the pointer."""

        type_name = type(self).__name__
        return (
            f"["
            f"{type_name} - "
            f"id:{self.id} "
            f"owner:{self.owner.id} "
            f"loc:{self.location.id} "
            f"id@loc:{self.id_at_location}"
            f"]"
        )

    def __repr__(self):
        """This returns the to-string method when called using __repr__, most
        commonly seen when returned as cells in Jupyter noteoboks."""
        return self.__str__()

    def get(self, deregister_ptr: bool = True):
        """Since PointerTensor objects always point to a remote tensor (or chain
        of tensors, where a chain is simply a linked-list of tensors linked via
        their .child attributes), this method will request that the tensor/chain
        being pointed to be serialized and returned from this function. This will
        typically mean that the remote object will be removed/destroyed. If you
        merely wish to bring a copy back to the local worker, call .copy() before
        calling .get(). TODO: add param get_copy which doesn't destroy remote if
        true.

        :Parameters:

         * **deregister_ptr (bool, optional)** this determines whether to deregister
            this pointer from the pointer's owner during this method. This defaults
            to True because the main reason people use this method is to move the
            tensor from the remote machine to the local one, at which time the
            pointer has no use.

         * **out (**:class:`.abstract.AbstractTensor` **)** this is the tensor
            (or chain) which this object used to point to on a remote machine.

        """

        # if the pointer happens to be pointing to a local object,
        # just return that object (this is an edge case)
        if self.location == self.owner:
            tensor = self.owner.get_obj(self.id_at_location).child
        else:
            # get tensor from remote machine
            tensor = self.owner.request_obj(self.id_at_location, self.location)

        # Register the result
        assigned_id = self.id_at_location
        self.owner.register_obj(tensor, assigned_id)

        # Remove this pointer by default
        if deregister_ptr:
            self.owner.de_register_obj(self)

        return tensor
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2015 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import ujson as json

from twisted.internet import defer

import baserules
from push_rule_evaluator import PushRuleEvaluatorForEvent

from synapse.api.constants import EventTypes

logger = logging.getLogger(__name__)

def decode_rule_json(rule):
    rule['conditions'] = json.loads(rule['conditions'])
    rule['actions'] = json.loads(rule['actions'])
    return rule

@defer.inlineCallbacks
def _get_rules(room_id, user_ids, store):
    rules_by_user = yield store.bulk_get_push_rules(user_ids)
    rules_by_user = {
        uid: baserules.list_with_base_rules([
            decode_rule_json(rule_list)
            for rule_list in rules_by_user.get(uid, [])
        ])
        for uid in user_ids
    }
    defer.returnValue(rules_by_user)

@defer.inlineCallbacks
def evaluator_for_room_id(room_id, store):
    users = yield store.get_users_in_room(room_id)
    rules_by_user = yield _get_rules(room_id, users, store)

    defer.returnValue(BulkPushRuleEvaluator(
        room_id, rules_by_user, users, store
    ))

class BulkPushRuleEvaluator:
    """
    Runs push rules for all users in a room.
    This is faster than running PushRuleEvaluator for each user because it
    fetches all the rules for all the users in one (batched) db query
    rather than doing multiple queries per-user. It currently uses
    the same logic to run the actual rules, but could be optimised further
    (see https://matrix.org/jira/browse/SYN-562)
    """
    def __init__(self, room_id, rules_by_user, users_in_room, store):
        self.room_id = room_id
        self.rules_by_user = rules_by_user
        self.users_in_room = users_in_room
        self.store = store

    @defer.inlineCallbacks
    def action_for_event_by_user(self, event, handler):
        actions_by_user = {}

        users_dict = yield self.store.are_guests(self.rules_by_user.keys())

        filtered_by_user = yield handler._filter_events_for_clients(
            users_dict.items(), [event]
        )

        evaluator = PushRuleEvaluatorForEvent(event, len(self.users_in_room))

        condition_cache = {}

        member_state = yield self.store.get_state_for_event(
            event.event_id,
        )

        display_names = {}
        for ev in member_state.values():
            nm = ev.content.get("displayname", None)
            if nm and ev.type == EventTypes.Member:
                display_names[ev.state_key] = nm

        for uid, rules in self.rules_by_user.items():
            display_name = display_names.get(uid, None)

            filtered = filtered_by_user[uid]
            if len(filtered) == 0:
                continue

            for rule in rules:
                if 'enabled' in rule and not rule['enabled']:
                    continue

                matches = _condition_checker(
                    evaluator, rule['conditions'], uid, display_name, condition_cache
                )
                if matches:
                    actions = [x for x in rule['actions'] if x != 'dont_notify']
                    if actions:
                        actions_by_user[uid] = actions
                    break
        defer.returnValue(actions_by_user)

def _condition_checker(evaluator, conditions, uid, display_name, cache):
    for cond in conditions:
        _id = cond.get("_id", None)
        if _id:
            res = cache.get(_id, None)
            if res is False:
                break
            elif res is True:
                continue

        res = evaluator.matches(cond, uid, display_name, None)
        if _id:
            cache[_id] = res

        if res is False:
            return False

    return True
<commit_msg>Except truthy values<commit_after># -*- coding: utf-8 -*-
# Copyright 2015 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import ujson as json

from twisted.internet import defer

import baserules
from push_rule_evaluator import PushRuleEvaluatorForEvent

from synapse.api.constants import EventTypes

logger = logging.getLogger(__name__)

def decode_rule_json(rule):
    rule['conditions'] = json.loads(rule['conditions'])
    rule['actions'] = json.loads(rule['actions'])
    return rule

@defer.inlineCallbacks
def _get_rules(room_id, user_ids, store):
    rules_by_user = yield store.bulk_get_push_rules(user_ids)
    rules_by_user = {
        uid: baserules.list_with_base_rules([
            decode_rule_json(rule_list)
            for rule_list in rules_by_user.get(uid, [])
        ])
        for uid in user_ids
    }
    defer.returnValue(rules_by_user)

@defer.inlineCallbacks
def evaluator_for_room_id(room_id, store):
    users = yield store.get_users_in_room(room_id)
    rules_by_user = yield _get_rules(room_id, users, store)

    defer.returnValue(BulkPushRuleEvaluator(
        room_id, rules_by_user, users, store
    ))

class BulkPushRuleEvaluator:
    """
    Runs push rules for all users in a room.
    This is faster than running PushRuleEvaluator for each user because it
    fetches all the rules for all the users in one (batched) db query
    rather than doing multiple queries per-user. It currently uses
    the same logic to run the actual rules, but could be optimised further
    (see https://matrix.org/jira/browse/SYN-562)
    """
    def __init__(self, room_id, rules_by_user, users_in_room, store):
        self.room_id = room_id
        self.rules_by_user = rules_by_user
        self.users_in_room = users_in_room
        self.store = store

    @defer.inlineCallbacks
    def action_for_event_by_user(self, event, handler):
        actions_by_user = {}

        users_dict = yield self.store.are_guests(self.rules_by_user.keys())

        filtered_by_user = yield handler._filter_events_for_clients(
            users_dict.items(), [event]
        )

        evaluator = PushRuleEvaluatorForEvent(event, len(self.users_in_room))

        condition_cache = {}

        member_state = yield self.store.get_state_for_event(
            event.event_id,
        )

        display_names = {}
        for ev in member_state.values():
            nm = ev.content.get("displayname", None)
            if nm and ev.type == EventTypes.Member:
                display_names[ev.state_key] = nm

        for uid, rules in self.rules_by_user.items():
            display_name = display_names.get(uid, None)

            filtered = filtered_by_user[uid]
            if len(filtered) == 0:
                continue

            for rule in rules:
                if 'enabled' in rule and not rule['enabled']:
                    continue

                matches = _condition_checker(
                    evaluator, rule['conditions'], uid, display_name, condition_cache
                )
                if matches:
                    actions = [x for x in rule['actions'] if x != 'dont_notify']
                    if actions:
                        actions_by_user[uid] = actions
                    break
        defer.returnValue(actions_by_user)

def _condition_checker(evaluator, conditions, uid, display_name, cache):
    for cond in conditions:
        _id = cond.get("_id", None)
        if _id:
            res = cache.get(_id, None)
            if res is False:
                break
            elif res is True:
                continue

        res = evaluator.matches(cond, uid, display_name, None)
        if _id:
            cache[_id] = res

        if not res:
            return False

    return True
<|endoftext|>


<commit_before>import os
import abc
import asyncio

import boto.glacier.utils

class BaseStream(asyncio.StreamReader, metaclass=abc.ABCMeta):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.readers = {}
        self.writers = {}

    @abc.abstractproperty
    def size(self):
        pass

    def add_reader(self, name, reader):
        self.readers[name] = reader

    def remove_reader(self, name):
        del self.readers[name]

    def add_writer(self, name, writer):
        self.writers[name] = writer

    def remove_writer(self, name):
        del self.writers[name]

    def feed_eof(self):
        super().feed_eof()
        for reader in self.readers.values():
            reader.feed_eof()
        for writer in self.writers.values():
            if hasattr(writer, 'can_write_eof') and writer.can_write_eof():
                writer.write_eof()

    @asyncio.coroutine
    def read(self, size=-1):
        eof = self.at_eof()
        data = yield from self._read(size)
        if not eof:
            for reader in self.readers.values():
                reader.feed_data(data)
            for writer in self.writers.values():
                writer.write(data)
        return data

    @abc.abstractmethod
    @asyncio.coroutine
    def _read(self, size):
        pass

class ResponseStreamReader(BaseStream):

    def __init__(self, response):
        super().__init__()
        self.response = response
        self.content_type = self.response.headers.get('Content-Type', 'application/octet-stream')

    @property
    def size(self):
        return self.response.headers.get('Content-Length')

    @asyncio.coroutine
    def _read(self, size):
        return (yield from self.response.content.read(size))

class RequestStreamReader(BaseStream):

    def __init__(self, request):
        super().__init__()
        self.request = request

    @property
    def size(self):
        return self.request.headers.get('Content-Length')

    @asyncio.coroutine
    def _read(self, size):
        return (yield from asyncio.StreamReader.read(self, size))

class FileStreamReader(BaseStream):

    def __init__(self, file_pointer):
        super().__init__()
        self.file_gen = None
        self.file_pointer = file_pointer
        self.read_size = None

    @property
    def size(self):
        cursor = self.file_pointer.tell()
        self.file_pointer.seek(0, os.SEEK_END)
        ret = self.file_pointer.tell()
        self.file_pointer.seek(cursor)
        return ret

    def close(self):
        self.file_pointer.close()
        self.feed_eof()

    def read_as_gen(self):
        self.file_pointer.seek(0)
        while True:
            data = self.file_pointer.read(self.read_size)
            if not data:
                break
            yield data

    @asyncio.coroutine
    def _read(self, size):
        self.file_gen = self.file_gen or self.read_as_gen()
        # add sleep of 0 so read will yield and continue in next io loop iteration
        yield from asyncio.sleep(0)
        self.read_size = size
        try:
            return next(self.file_gen)
        except StopIteration:
            return b''

class HashStreamWriter:
    """Stream-like object that hashes and discards its input."""
    def __init__(self, hasher):
        self.hash = hasher()

    @property
    def hexdigest(self):
        return self.hash.hexdigest()

    def can_write_eof(self):
        return False

    def write(self, data):
        self.hash.update(data)

    def close(self):
        pass

class HashTreeStreamWriter:
    """Stream-like object that hashes and discards its input."""
    def __init__(self, hasher, chunk_size=1024*1024):
        self._hasher = hasher
        self._hashes = []
        self._chunks = b''
        self._chunk_size = chunk_size

    @property
    def hexdigest(self):
        return boto.glacier.utils.bytes_to_hex(boto.glacier.utils.tree_hash(self._hashes))

    def can_write_eof(self):
        return False

    def write(self, data):
        self._chunks += data
        if len(self._chunks) >= self._chunk_size:
            to_hash = self._chunks[:self._chunk_size]
            self._chunks = self._chunks[self._chunk_size:]
            self._hashes.append(self._hasher(to_hash).digest())

    def close(self):
        if self._chunks:
            self._hashes.append(self._hasher(self._chunks).digest())
<commit_msg>Fix flake error.<commit_after>import os
import abc
import asyncio

import boto.glacier.utils

class BaseStream(asyncio.StreamReader, metaclass=abc.ABCMeta):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.readers = {}
        self.writers = {}

    @abc.abstractproperty
    def size(self):
        pass

    def add_reader(self, name, reader):
        self.readers[name] = reader

    def remove_reader(self, name):
        del self.readers[name]

    def add_writer(self, name, writer):
        self.writers[name] = writer

    def remove_writer(self, name):
        del self.writers[name]

    def feed_eof(self):
        super().feed_eof()
        for reader in self.readers.values():
            reader.feed_eof()
        for writer in self.writers.values():
            if hasattr(writer, 'can_write_eof') and writer.can_write_eof():
                writer.write_eof()

    @asyncio.coroutine
    def read(self, size=-1):
        eof = self.at_eof()
        data = yield from self._read(size)
        if not eof:
            for reader in self.readers.values():
                reader.feed_data(data)
            for writer in self.writers.values():
                writer.write(data)
        return data

    @abc.abstractmethod
    @asyncio.coroutine
    def _read(self, size):
        pass

class ResponseStreamReader(BaseStream):

    def __init__(self, response):
        super().__init__()
        self.response = response
        self.content_type = self.response.headers.get('Content-Type', 'application/octet-stream')

    @property
    def size(self):
        return self.response.headers.get('Content-Length')

    @asyncio.coroutine
    def _read(self, size):
        return (yield from self.response.content.read(size))

class RequestStreamReader(BaseStream):

    def __init__(self, request):
        super().__init__()
        self.request = request

    @property
    def size(self):
        return self.request.headers.get('Content-Length')

    @asyncio.coroutine
    def _read(self, size):
        return (yield from asyncio.StreamReader.read(self, size))

class FileStreamReader(BaseStream):

    def __init__(self, file_pointer):
        super().__init__()
        self.file_gen = None
        self.file_pointer = file_pointer
        self.read_size = None

    @property
    def size(self):
        cursor = self.file_pointer.tell()
        self.file_pointer.seek(0, os.SEEK_END)
        ret = self.file_pointer.tell()
        self.file_pointer.seek(cursor)
        return ret

    def close(self):
        self.file_pointer.close()
        self.feed_eof()

    def read_as_gen(self):
        self.file_pointer.seek(0)
        while True:
            data = self.file_pointer.read(self.read_size)
            if not data:
                break
            yield data

    @asyncio.coroutine
    def _read(self, size):
        self.file_gen = self.file_gen or self.read_as_gen()
        # add sleep of 0 so read will yield and continue in next io loop iteration
        yield from asyncio.sleep(0)
        self.read_size = size
        try:
            return next(self.file_gen)
        except StopIteration:
            return b''

class HashStreamWriter:
    """Stream-like object that hashes and discards its input."""
    def __init__(self, hasher):
        self.hash = hasher()

    @property
    def hexdigest(self):
        return self.hash.hexdigest()

    def can_write_eof(self):
        return False

    def write(self, data):
        self.hash.update(data)

    def close(self):
        pass

class HashTreeStreamWriter:
    """Stream-like object that hashes and discards its input."""
    def __init__(self, hasher, chunk_size=1024 * 1024):
        self._hasher = hasher
        self._hashes = []
        self._chunks = b''
        self._chunk_size = chunk_size

    @property
    def hexdigest(self):
        return boto.glacier.utils.bytes_to_hex(boto.glacier.utils.tree_hash(self._hashes))

    def can_write_eof(self):
        return False

    def write(self, data):
        self._chunks += data
        if len(self._chunks) >= self._chunk_size:
            to_hash = self._chunks[:self._chunk_size]
            self._chunks = self._chunks[self._chunk_size:]
            self._hashes.append(self._hasher(to_hash).digest())

    def close(self):
        if self._chunks:
            self._hashes.append(self._hasher(self._chunks).digest())
<|endoftext|>


<commit_before>import json

from django.http import HttpResponse
from django.views.generic import TemplateView, View

from map.models import CityBorder
from map.utils import generateGeoJson

class MapView(TemplateView):
    template_name = 'map.html'

class FetchGridView(View):

    def get(self, request, *args, **kwargs):
        pk = self.request.GET.get('pk')
        size = int(self.request.GET.get('size'))
        city = CityBorder.objects.get(pk=pk)
        grid = city.generateGrid(size)
        geojson = generateGeoJson(list(grid))
        return HttpResponse(json.dumps(geojson))
<commit_msg>added crime types in context data in views<commit_after>import json

from django.http import HttpResponse
from django.views.generic import TemplateView, View

from crime.models import CriminalRecord
from map.models import CityBorder
from map.utils import generateGeoJson

class MapView(TemplateView):
    template_name = 'map.html'

    def get_context_data(self, **kwargs):
        context = {}
        context['crime_types'] = CriminalRecord.objects.all().order_by(
            'primary_type').distinct().values_list('primary_type', flat=True)
        return context

class FetchGridView(View):

    def get(self, request, *args, **kwargs):
        pk = self.request.GET.get('pk')
        size = int(self.request.GET.get('size'))
        city = CityBorder.objects.get(pk=pk)
        grid = city.generateGrid(size)
        geojson = generateGeoJson(list(grid))
        return HttpResponse(json.dumps(geojson))
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
#
# MIT License
#
# Copyright (c) 2017  Miha Purg <miha.purg@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#
#

import re
import logging
from collections import namedtuple

from Qpyl.common import raise_or_log

logger = logging.getLogger(__name__)

PosVector = namedtuple("PosVector", ["x", "y", "z"])

_PLACEHOLDER_RE = re.compile("\$\S+\.\S+\$")

def find_placeholders(inputstring):
    """Find atom placeholders ($514.C3).

    See also QStruct.convert_placeholders
    """
    return _PLACEHOLDER_RE.findall(inputstring)

class QStructError(Exception):
    pass

class QStruct(object):
    """
    Class for processing the structure (coordinates)

    Args:
      filename (path to the structure file)
      filetype (type of structure file:  'pdb' or 'mol2')
      ignore_errors (boolean):  Optional, default is False.
                                If set to True, some non-vital
                                exceptions are logged instead.

    In contrast to QLib and QPrm, the 'read' methods in this
    class should not be called since the object should
    contain data from only one structure file.

    The structure data is stored in three lists:
      atoms, residues, molecules
    which contain _StructAtom, _StructResidue and _StructMolecule
    objects.

    """

    def __init__(self, filename, filetype, ignore_errors=False):
        self.ignore_errors = ignore_errors
        FILE_TYPES = {'pdb': self._read_pdb,
                      'mol2': self._read_mol2}
        self.filetype = filetype.lower()
        if self.filetype not in FILE_TYPES.keys():
            raise QStructError("Filetype {} not supported. Use {}"
                               .format(filetype,
                                       " or ".join(FILE_TYPES.keys())))

        self.atoms = []
        self.residues = []
        self.molecules = []
        self.filename = filename
        # TODO: some sort of lookup hashes if needed

        # run the parser function
        FILE_TYPES[self.filetype](filename)

        # check if we actually got something
        for t in ["atoms", "residues", "molecules"]:
            if len(self.__dict__[t]) == 0:
                raise QStructError("No {} found, check file '{}' and"
                                   " filetype '{}'".format(t,
                                                           self.filename,
                                                           self.filetype))

    def _read_mol2(self, mol2_file):
        """
        Read and parse a mol2 file for coordinates.

        Args:
            mol2_file (string):  name/path of file
        """

        molecule = None
        residue = None
        aindex, old_aindex = None, None
        section = None
        for line in open(mol2_file, 'r').readlines():

            if line.startswith("@<TRIPOS>"):
                section = line.replace("@<TRIPOS>", "").strip()
                if section == "MOLECULE":
                    if molecule != None:
                        self.molecules.append(molecule)

                    molecule = _StructMolecule(self)
                continue

            if section == "ATOM":
                if aindex != None:
                    old_aindex = aindex

                lf = line.split()

                aindex, aname = int(lf[0]), lf[1]
                x, y, z = map(float, lf[2:5])
                rindex = int(lf[6])
                rname = lf[7][0:4].upper()

                if old_aindex != None and aindex - old_aindex != 1:
                    raise_or_log("Bad Mol2 format - atom "
                                 "index {} followed by {}"
                                 .format(old_aindex, aindex),
                                 QStructError, logger, self.ignore_errors)

                if not residue or residue.index != rindex:
                    if residue and rindex - residue.index != 1:
                        raise_or_log("Bad Mol2 format - residue "
                                     "index {} followed by {}"
                                     .format(residue.index, rindex),
                                     QStructError, logger, self.ignore_errors)
                    residue = _StructResidue(rindex, rname, molecule, self)
                    self.residues.append(residue)
                    molecule.add_residue(residue)

                atom = _StructAtom(aindex, aname, x, y, z, residue, self)
                self.atoms.append(atom)
                residue.add_atom(atom)

        # append last one after parsing
        if molecule != None and len(molecule.residues) > 0:
            self.molecules.append(molecule)

    def _read_pdb(self, pdb_file):
        """
        Read and parse a PDB file for coordinates.

        Args:
            pdb_file (string):  name/path of file

        """
        # make a new _StructMolecule object
        molecule = _StructMolecule(self)
        # parse the PDB file
        residue = None
        aindex, old_aindex = None, None
        for line in open(pdb_file, 'r').readlines():
            if line.startswith("ATOM") or line.startswith("HETATM"):
                if aindex != None:
                    old_aindex = aindex

                aindex = int(line[6:12])
                if old_aindex != None and aindex - old_aindex != 1:
                    raise_or_log("Bad PDB format - atom "
                                 "index {} followed by {}"
                                 .format(old_aindex, aindex),
                                 QStructError, logger, self.ignore_errors)

                aname = line[12:17].strip()
                rname = line[17:20].strip().upper()
                rindex = int(line[22:26])
                x, y, z = map(float, (line[30:38], line[38:46], line[46:54]))

                if not residue or residue.index != rindex:
                    if residue and rindex - residue.index != 1:
                        raise_or_log("Bad PDB format - residue "
                                     "index {} followed by {}"
                                     .format(residue.index, rindex),
                                     QStructError, logger, self.ignore_errors)

                    residue = _StructResidue(rindex, rname, molecule, self)
                    self.residues.append(residue)
                    molecule.add_residue(residue)

                atom = _StructAtom(aindex, aname, x, y, z, residue, self)
                self.atoms.append(atom)
                residue.add_atom(atom)

            elif line.startswith("TER") or line.startswith("GAP"):
                self.molecules.append(molecule)
                residue = None
                molecule = _StructMolecule(self)

        # append last one if it didn't gave a TER/GAP
        if molecule != None and len(molecule.residues) > 0:
            self.molecules.append(molecule)

    def convert_placeholders(self, inputstring):
        """Convert atom placeholders ($514.C3$) to indexes.

        Placeholders are a combination of the residue id and
          atom name, encapsulated in $$ - $RESID.ATOM_NAME$
        In addition,there are some special values:
          $LAST.ID$ - id of last atom in the system

        Arguments:
          inputstring (string):  string with placeholders (input file contents)

        Returns:
          outputstring (string):  converted string

        """

        id_map = {"{}.{}".format(a.residue.index, a.name): str(a.index)
                                                         for a in self.atoms}
        last_id = "{}.{}".format(self.atoms[-1].residue.index,
                                 self.atoms[-1].name)

        outputstring = ""
        for line in inputstring.split("\n"):
            comment = ""
            if "#" in line:
                i = line.index("#")
                line, comment = line[:i], line[i:]

            c = find_placeholders(line)
            for pid in c:
                pid = pid.strip("$")
                pid2 = pid.replace("LAST.ID", last_id)
                try:
                    padding = (len(pid2)+2  - len(id_map[pid2])) * " "
                except KeyError:
                    raise QStructError("Atom '${}$' does not exist in the pdb "
                                       "structure.".format(pid2))
                line = re.sub("\$" + pid + "\$", id_map[pid2] + padding, line)

            outputstring += line + comment + "\n"
        return outputstring

class _StructAtom(object):
    """Contains structural information for an atom.

    Arguments:
      index_struct (int):  index as written in pdb or mol2
      name (string):  atom name
      x,y,z (float):  coordinates
      residue (_StructResidue):  parent residue object
      structure (_QStruct):  parent structure object

    Property 'index' (int) is the actual 1-based index of the atom
    in the atom list (as opposed to index_struct which was read from
    the file). It should correspond to the index in the generated topology.
    """
    def __init__(self, index_struct, name, x, y, z, residue, structure):
        self.index_struct = int(index_struct)
        self.name = name
        self.coordinates = PosVector(float(x), float(y), float(z))
        self.residue = residue
        self.structure = structure

    @property
    def index(self):
        return self.structure.atoms.index(self) + 1

    def __repr__(self):
        res = self.residue
        mol = res.molecule
        return "_StructAtom: {}.{}.{}".format(mol.index,
                                              res.index,
                                              self.index)

class _StructResidue(object):
    """Contains structural information for a residue.

    Arguments:
      index_struct (int):  index as written in pdb or mol2
      name (string):  residue name
      molecule (_StructMolecule):  parent molecule object
      structure (_QStruct):  parent structure object

    Property 'index' (int) is the actual 1-based index of the residue
    in the residue list (as opposed to index_struct which was read from
    the file). It should correspond to the index in the generated topology.
    """
    def __init__(self, index_struct, name, molecule, structure):
        self.atoms = []
        self.index_struct = int(index_struct)
        self.name = name
        self.molecule = molecule
        self.structure = structure

    @property
    def index(self):
        return self.structure.residues.index(self) + 1

    def add_atom(self, atom):
        self.atoms.append(atom)

    def __repr__(self):
        mol = self.molecule
        return "_StructResidue: {}.{}{}".format(mol.index,
                                                self.name,
                                                self.index)

class _StructMolecule(object):
    """Contains structural information for a molecule.

    Arguments:
      structure (_QStruct):  parent structure object

    Special property is 'index' (int). It is the actual
    1-based index of the molecule in the residue list (as it was appended).
    This should corresponds to the index in the generated topology.
    """
    def __init__(self, structure):
        self.residues = []
        self.structure = structure

    @property
    def index(self):
        return self.structure.molecules.index(self) + 1

    def add_residue(self, residue):
        self.residues.append(residue)

    def __repr__(self):
        return "_StructMolecule: {}".format(self.index)

<commit_msg>Bugfix in Qpyl.core.qstructure. Current atom residue index during parsing was compared to previous atom's assigned residue index (1,2,3...) rather than the actual PDB index (32,33,34,...), resulting in a lot of warnings and every atom being treated as its own residue.<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
#
# MIT License
#
# Copyright (c) 2017  Miha Purg <miha.purg@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#
#

import re
import logging
from collections import namedtuple

from Qpyl.common import raise_or_log

logger = logging.getLogger(__name__)

PosVector = namedtuple("PosVector", ["x", "y", "z"])

_PLACEHOLDER_RE = re.compile("\$\S+\.\S+\$")

def find_placeholders(inputstring):
    """Find atom placeholders ($514.C3).

    See also QStruct.convert_placeholders
    """
    return _PLACEHOLDER_RE.findall(inputstring)

class QStructError(Exception):
    pass

class QStruct(object):
    """
    Class for processing the structure (coordinates)

    Args:
      filename (path to the structure file)
      filetype (type of structure file:  'pdb' or 'mol2')
      ignore_errors (boolean):  Optional, default is False.
                                If set to True, some non-vital
                                exceptions are logged instead.

    In contrast to QLib and QPrm, the 'read' methods in this
    class should not be called since the object should
    contain data from only one structure file.

    The structure data is stored in three lists:
      atoms, residues, molecules
    which contain _StructAtom, _StructResidue and _StructMolecule
    objects.

    """

    def __init__(self, filename, filetype, ignore_errors=False):
        self.ignore_errors = ignore_errors
        FILE_TYPES = {'pdb': self._read_pdb,
                      'mol2': self._read_mol2}
        self.filetype = filetype.lower()
        if self.filetype not in FILE_TYPES.keys():
            raise QStructError("Filetype {} not supported. Use {}"
                               .format(filetype,
                                       " or ".join(FILE_TYPES.keys())))

        self.atoms = []
        self.residues = []
        self.molecules = []
        self.filename = filename
        # TODO: some sort of lookup hashes if needed

        # run the parser function
        FILE_TYPES[self.filetype](filename)

        # check if we actually got something
        for t in ["atoms", "residues", "molecules"]:
            if len(self.__dict__[t]) == 0:
                raise QStructError("No {} found, check file '{}' and"
                                   " filetype '{}'".format(t,
                                                           self.filename,
                                                           self.filetype))

    def _read_mol2(self, mol2_file):
        """
        Read and parse a mol2 file for coordinates.

        Args:
            mol2_file (string):  name/path of file
        """

        molecule = None
        residue = None
        aindex, old_aindex = None, None
        section = None
        for line in open(mol2_file, 'r').readlines():

            if line.startswith("@<TRIPOS>"):
                section = line.replace("@<TRIPOS>", "").strip()
                if section == "MOLECULE":
                    if molecule != None:
                        self.molecules.append(molecule)

                    molecule = _StructMolecule(self)
                continue

            if section == "ATOM":
                if aindex != None:
                    old_aindex = aindex

                lf = line.split()

                aindex, aname = int(lf[0]), lf[1]
                x, y, z = map(float, lf[2:5])
                rindex = int(lf[6])
                rname = lf[7][0:4].upper()

                if old_aindex != None and aindex - old_aindex != 1:
                    raise_or_log("Bad Mol2 format - atom "
                                 "index {} followed by {}"
                                 .format(old_aindex, aindex),
                                 QStructError, logger, self.ignore_errors)

                if not residue or residue.index_struct != rindex:
                    if residue and rindex - residue.index_struct != 1:
                        raise_or_log("Bad Mol2 format - residue "
                                     "index {} followed by {}"
                                     .format(residue.index_struct, rindex),
                                     QStructError, logger, self.ignore_errors)
                    residue = _StructResidue(rindex, rname, molecule, self)
                    self.residues.append(residue)
                    molecule.add_residue(residue)

                atom = _StructAtom(aindex, aname, x, y, z, residue, self)
                self.atoms.append(atom)
                residue.add_atom(atom)

        # append last one after parsing
        if molecule != None and len(molecule.residues) > 0:
            self.molecules.append(molecule)

    def _read_pdb(self, pdb_file):
        """
        Read and parse a PDB file for coordinates.

        Args:
            pdb_file (string):  name/path of file

        """
        # make a new _StructMolecule object
        molecule = _StructMolecule(self)
        # parse the PDB file
        residue = None
        aindex, old_aindex = None, None
        for line in open(pdb_file, 'r').readlines():
            if line.startswith("ATOM") or line.startswith("HETATM"):
                if aindex != None:
                    old_aindex = aindex

                aindex = int(line[6:12])
                if old_aindex != None and aindex - old_aindex != 1:
                    raise_or_log("Bad PDB format - atom "
                                 "index {} followed by {}"
                                 .format(old_aindex, aindex),
                                 QStructError, logger, self.ignore_errors)

                aname = line[12:17].strip()
                rname = line[17:20].strip().upper()
                rindex = int(line[22:26])
                x, y, z = map(float, (line[30:38], line[38:46], line[46:54]))

                if not residue or residue.index_struct != rindex:
                    if residue and rindex - residue.index_struct != 1:
                        raise_or_log("Bad PDB format - residue "
                                     "index {} followed by {}"
                                     .format(residue.index_struct, rindex),
                                     QStructError, logger, self.ignore_errors)

                    residue = _StructResidue(rindex, rname, molecule, self)
                    self.residues.append(residue)
                    molecule.add_residue(residue)

                atom = _StructAtom(aindex, aname, x, y, z, residue, self)
                self.atoms.append(atom)
                residue.add_atom(atom)

            elif line.startswith("TER") or line.startswith("GAP"):
                self.molecules.append(molecule)
                residue = None
                molecule = _StructMolecule(self)

        # append last one if it didn't gave a TER/GAP
        if molecule != None and len(molecule.residues) > 0:
            self.molecules.append(molecule)

    def convert_placeholders(self, inputstring):
        """Convert atom placeholders ($514.C3$) to indexes.

        Placeholders are a combination of the residue id and
          atom name, encapsulated in $$ - $RESID.ATOM_NAME$
        In addition,there are some special values:
          $LAST.ID$ - id of last atom in the system

        Arguments:
          inputstring (string):  string with placeholders (input file contents)

        Returns:
          outputstring (string):  converted string

        """

        id_map = {"{}.{}".format(a.residue.index, a.name): str(a.index)
                                                         for a in self.atoms}
        last_id = "{}.{}".format(self.atoms[-1].residue.index,
                                 self.atoms[-1].name)

        outputstring = ""
        for line in inputstring.split("\n"):
            comment = ""
            if "#" in line:
                i = line.index("#")
                line, comment = line[:i], line[i:]

            c = find_placeholders(line)
            for pid in c:
                pid = pid.strip("$")
                pid2 = pid.replace("LAST.ID", last_id)
                try:
                    padding = (len(pid2)+2  - len(id_map[pid2])) * " "
                except KeyError:
                    raise QStructError("Atom '${}$' does not exist in the pdb "
                                       "structure.".format(pid2))
                line = re.sub("\$" + pid + "\$", id_map[pid2] + padding, line)

            outputstring += line + comment + "\n"
        return outputstring

class _StructAtom(object):
    """Contains structural information for an atom.

    Arguments:
      index_struct (int):  index as written in pdb or mol2
      name (string):  atom name
      x,y,z (float):  coordinates
      residue (_StructResidue):  parent residue object
      structure (_QStruct):  parent structure object

    Property 'index' (int) is the actual 1-based index of the atom
    in the atom list (as opposed to index_struct which was read from
    the file). It should correspond to the index in the generated topology.
    """
    def __init__(self, index_struct, name, x, y, z, residue, structure):
        self.index_struct = int(index_struct)
        self.name = name
        self.coordinates = PosVector(float(x), float(y), float(z))
        self.residue = residue
        self.structure = structure

    @property
    def index(self):
        return self.structure.atoms.index(self) + 1

    def __repr__(self):
        res = self.residue
        mol = res.molecule
        return "_StructAtom: {}.{}.{}".format(mol.index,
                                              res.index,
                                              self.index)

class _StructResidue(object):
    """Contains structural information for a residue.

    Arguments:
      index_struct (int):  index as written in pdb or mol2
      name (string):  residue name
      molecule (_StructMolecule):  parent molecule object
      structure (_QStruct):  parent structure object

    Property 'index' (int) is the actual 1-based index of the residue
    in the residue list (as opposed to index_struct which was read from
    the file). It should correspond to the index in the generated topology.
    """
    def __init__(self, index_struct, name, molecule, structure):
        self.atoms = []
        self.index_struct = int(index_struct)
        self.name = name
        self.molecule = molecule
        self.structure = structure

    @property
    def index(self):
        return self.structure.residues.index(self) + 1

    def add_atom(self, atom):
        self.atoms.append(atom)

    def __repr__(self):
        mol = self.molecule
        return "_StructResidue: {}.{}{}".format(mol.index,
                                                self.name,
                                                self.index)

class _StructMolecule(object):
    """Contains structural information for a molecule.

    Arguments:
      structure (_QStruct):  parent structure object

    Special property is 'index' (int). It is the actual
    1-based index of the molecule in the residue list (as it was appended).
    This should corresponds to the index in the generated topology.
    """
    def __init__(self, structure):
        self.residues = []
        self.structure = structure

    @property
    def index(self):
        return self.structure.molecules.index(self) + 1

    def add_residue(self, residue):
        self.residues.append(residue)

    def __repr__(self):
        return "_StructMolecule: {}".format(self.index)

<|endoftext|>