<commit_before>import time, math
import numpy as np

def resample(ts, values, num_samples):
    """Convert a list of times and a list of values to evenly spaced samples with linear interpolation"""
    ts = normalize(ts)
    return np.interp(np.linspace(0.0, 1.0, num_samples), ts, values)

def upsample(signal, factor):
    """Increase the sampling rate of a signal (by an integer factor), with linear interpolation"""
    assert type(factor) == int and factor > 1
    result = [None] * ((len(signal) - 1) * factor)
    for i, v in enumerate(signal):
        if i == len(signal) - 1:
            result[-1] = v
            break
        v_ = signal[i+1]
        delta = v_ - v
        for j in range(factor):
            f = (i * factor) + j
            result[f] = v + ((delta / factor) * j)
    return result

def downsample(signal, factor):
    """Decrease the sampling rate of a signal (by an integer factor), with averaging"""
    signal = np.array(signal)
    xs = signal.shape[0]
    signal = signal[:xs - (xs % int(factor))]
    result = np.mean(np.concatenate([[signal[i::factor] for i in range(factor)]]), axis=0)
    return result

def normalize(signal, minimum=None, maximum=None):
    """Normalize a signal to the range 0, 1"""
    signal = np.array(signal).astype('float')
    if minimum is None:
        minimum = np.min(signal)
    signal -= minimum
    if maximum is None:
        maximum = np.max(signal)
    signal /= maximum
    signal = np.clip(signal, 0.0, 1.0)
    return signal

def magnitude(signal):
    """Absolute value of a signal"""
    ## why am I seeing some negative values?
    signal = np.absolute(signal)
    signal = threshold(signal, 0)
    return signal

def threshold(signal, value):
    """Drop all values in a signal to 0 if below the given threshold"""
    signal = np.array(signal)
    return (signal > value) * signal

def limit(signal, value):
    """Limit all values in a signal to the given value"""
    return np.clip(signal, 0, value)

def remove_shots(signal, threshold=None, devs=2, positive_only=False):
    """Replace values in a signal that are above a threshold or vary by a given number of deviations with the average of the surrounding samples"""
    average = np.average(signal)
    if threshold is not None:
        shot_indexes = [i for (i, sample) in enumerate(signal) if sample > threshold]
    else:
        deviation = np.std(signal)
        shot_indexes = [i for (i, sample) in enumerate(signal) if (sample - average if positive_only else abs(sample - average)) > deviation * devs]
    for i in shot_indexes:
        neighbors = []
        j = i + 1
        k = i - 1
        while j in shot_indexes:
            j += 1
        if j < len(signal):
            neighbors.append(signal[j])
        while k in shot_indexes:
            k -= 1
        if k >= 0:
            neighbors.append(signal[k])
        signal[i] = sum(neighbors) / float(len(neighbors))
    return signal

def compress(signal, value=2.0, normalize=False):
    """Compress the signal by an exponential value (will expand if value<0)"""
    signal = np.array(signal)
    signal = np.power(signal, 1.0 / value)
    return normalize(signal) if normalize else signal

def smooth(signal, size=10, window='blackman'):
    """Apply averaging / low-pass filter to a signal with the given window shape and size"""
    types = ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']
    signal = np.array(signal)
    if size < 3:
        return signal
    s = np.r_[2 * signal[0] - signal[size:1:-1], signal, 2 * signal[-1] - signal[-1:-size:-1]]
    if window == 'flat': # running average
        w = np.ones(size,'d')
    else:
        w = getattr(np, window)(size)
    y = np.convolve(w / w.sum(), s, mode='same')
    return y[size - 1:-size + 1]

## todo: def hipass
## smooth and then subtract result

def detect_peaks(signal, lookahead=300, delta=0):   ## probably a better scipy module...
    """ Detect the local maximas and minimas in a signal
        lookahead -- samples to look ahead from a potential peak to see if a bigger one is coming
        delta -- minimum difference between a peak and surrounding points to be considered a peak (no hills) and makes things faster
        Note: careful if you have flat regions, may affect lookahead
    """
    signal = np.array(signal)
    peaks = []
    valleys = []
    min_value, max_value = np.Inf, -np.Inf
    for index, value in enumerate(signal[:-lookahead]):
        if value > max_value:
            max_value = value
            max_pos = index
        if value < min_value:
            min_value = value
            min_pos = index
        if value < max_value - delta and max_value != np.Inf:
            if signal[index:index + lookahead].max() < max_value:
                peaks.append([max_pos, max_value])
                drop_first_peak = True
                max_value = np.Inf
                min_value = np.Inf
                if index + lookahead >= signal.size:
                    break
                continue
        if value > min_value + delta and min_value != -np.Inf:
            if signal[index:index + lookahead].min() > min_value:
                valleys.append([min_pos, min_value])
                drop_first_valley = True
                min_value = -np.Inf
                max_value = -np.Inf
                if index + lookahead >= signal.size:
                    break
    return peaks, valleys

def autocorrelate(signal):
    """Get the auto-correlation function of a signal"""
    x = np.hstack((signal, np.zeros(len(signal))))
    sp = np.fft.rfft(x)
    tmp = np.empty_like(sp)
    tmp = np.conj(sp, tmp)
    tmp = np.multiply(tmp, sp, tmp)
    ac = np.fft.irfft(tmp)
    ac = np.divide(ac, signal.size, ac)[:signal.size/ 2]
    tmp = signal.size / (signal.size - np.arange(signal.size / 2, dtype=np.float64))
    ac = np.multiply(ac, tmp, ac)
    ac = np.concatenate([ac, np.zeros(signal.size - ac.size)])
    return normalize(ac)

def derivative(signal):
    """Return a signal that is the derivative function of a given signal"""
    def f(x):
        x = int(x)
        return signal[x]
    def df(x, h=0.1e-5):
        return (f(x + h * 0.5) - f(x - h * 0.5)) / h
    return [df(x) for x in range(len(signal))]

def integral(signal):
    """Return a signal that is the integral function of a given signal"""
    result = []
    v = 0.0
    for i in xrange(len(signal)):
        v += signal[i]
        result.append(v)
    return result

def trendline(signal):
    """Returns a  line (slope, intersect) that is the regression line given a series of values."""
    signal = list(signal)
    n = len(signal) - 1
    sum_x = 0
    sum_y = 0
    sum_xx = 0
    sum_xy = 0
    for i in range(1, n + 1):
        x = i
        y = signal[i]
        sum_x = sum_x + x
        sum_y = sum_y + y
        xx = math.pow(x, 2)
        sum_xx = sum_xx + xx
        xy = x*y
        sum_xy = sum_xy + xy
    try:
        a = (-sum_x * sum_xy + sum_xx * sum_y) / (n * sum_xx - sum_x * sum_x)
        b = (-sum_x * sum_y + n * sum_xy) / (n * sum_xx - sum_x * sum_x)
    except ZeroDivisionError:
        a, b = 0, 0
    return (b, a) # (slope, intersect)

def bandpass_filter(signal, sampling_rate, lowcut, highcut, order=6):
    """In hz"""
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    low = lowcut / nyquist
    high = highcut / nyquist
    b, a = butter(order, [low, high], btype='band')
    signal = lfilter(b, a, signal)
    return signal

def lowpass_filter(signal, sampling_rate, cutoff, order=6):
    """smooth appears to be much faster in certain situations"""
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    signal = lfilter(b, a, signal)
    return signal

def highpass_filter(signal, sampling_rate, cutoff, order=6):
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='high', analog=False)
    signal = lfilter(b, a, signal)
    return signal    <commit_msg>added rescaling<commit_after>import time, math
import numpy as np

def resample(ts, values, num_samples=None):
    """Convert a list of times and a list of values to evenly spaced samples with linear interpolation"""
    if num_samples is None:
        num_samples = len(ts)
    ts = normalize(ts)
    return np.interp(np.linspace(0.0, 1.0, num_samples), ts, values)

def upsample(signal, factor):
    """Increase the sampling rate of a signal (by an integer factor), with linear interpolation"""
    assert type(factor) == int and factor > 1
    result = [None] * ((len(signal) - 1) * factor)
    for i, v in enumerate(signal):
        if i == len(signal) - 1:
            result[-1] = v
            break
        v_ = signal[i+1]
        delta = v_ - v
        for j in range(factor):
            f = (i * factor) + j
            result[f] = v + ((delta / factor) * j)
    return result

def downsample(signal, factor):
    """Decrease the sampling rate of a signal (by an integer factor), with averaging"""
    signal = np.array(signal)
    xs = signal.shape[0]
    signal = signal[:xs - (xs % int(factor))]
    result = np.mean(np.concatenate([[signal[i::factor] for i in range(factor)]]), axis=0)
    return result

def normalize(signal, minimum=None, maximum=None):
    """Normalize a signal to the range 0, 1"""
    signal = np.array(signal).astype('float')
    if minimum is None:
        minimum = np.min(signal)
    signal -= minimum
    if maximum is None:
        maximum = np.max(signal)
    signal /= maximum
    signal = np.clip(signal, 0.0, 1.0)
    return signal

def rescale(signal, low, high):
    """Rescale a signal (normalize it first) to a given range"""
    signal *= high - low
    signal += low
    return signal

def make_audio(signal):
    signal = normalize(signal)
    signal = rescale(signal, -32768, 32767)
    signal = signal.astype(np.int16)
    return signal

def magnitude(signal):
    """Absolute value of a signal"""
    ## why am I seeing some negative values?
    signal = np.absolute(signal)
    signal = threshold(signal, 0)
    return signal

def threshold(signal, value):
    """Drop all values in a signal to 0 if below the given threshold"""
    signal = np.array(signal)
    return (signal > value) * signal

def limit(signal, value):
    """Limit all values in a signal to the given value"""
    return np.clip(signal, 0, value)

def remove_shots(signal, threshold=None, devs=2, positive_only=False):
    """Replace values in a signal that are above a threshold or vary by a given number of deviations with the average of the surrounding samples"""
    average = np.average(signal)
    if threshold is not None:
        shot_indexes = [i for (i, sample) in enumerate(signal) if sample > threshold]
    else:
        deviation = np.std(signal)
        shot_indexes = [i for (i, sample) in enumerate(signal) if (sample - average if positive_only else abs(sample - average)) > deviation * devs]
    for i in shot_indexes:
        neighbors = []
        j = i + 1
        k = i - 1
        while j in shot_indexes:
            j += 1
        if j < len(signal):
            neighbors.append(signal[j])
        while k in shot_indexes:
            k -= 1
        if k >= 0:
            neighbors.append(signal[k])
        signal[i] = sum(neighbors) / float(len(neighbors))
    return signal

def compress(signal, value=2.0, normalize=False):
    """Compress the signal by an exponential value (will expand if value<0)"""
    signal = np.array(signal)
    signal = np.power(signal, 1.0 / value)
    return normalize(signal) if normalize else signal

def smooth(signal, size=10, window='blackman'):
    """Apply weighted moving average (aka low-pass filter) via convolution function to a signal with the given window shape and size"""
    """This is going to be faster than highpass_filter"""
    types = ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']
    signal = np.array(signal)
    if size < 3:
        return signal
    s = np.r_[2 * signal[0] - signal[size:1:-1], signal, 2 * signal[-1] - signal[-1:-size:-1]]
    if window == 'flat': # running average
        w = np.ones(size,'d')
    else:
        w = getattr(np, window)(size) # get a series of weights that matches the window and is the correct size
    y = np.convolve(w / w.sum(), s, mode='same') # convolve the signals
    return y[size - 1:-size + 1]

## todo: detrend
## smooth and then subtract result

def detect_peaks(signal, lookahead=300, delta=0):   ## probably a better scipy module...
    """ Detect the local maximas and minimas in a signal
        lookahead -- samples to look ahead from a potential peak to see if a bigger one is coming
        delta -- minimum difference between a peak and surrounding points to be considered a peak (no hills) and makes things faster
        Note: careful if you have flat regions, may affect lookahead
    """
    signal = np.array(signal)
    peaks = []
    valleys = []
    min_value, max_value = np.Inf, -np.Inf
    for index, value in enumerate(signal[:-lookahead]):
        if value > max_value:
            max_value = value
            max_pos = index
        if value < min_value:
            min_value = value
            min_pos = index
        if value < max_value - delta and max_value != np.Inf:
            if signal[index:index + lookahead].max() < max_value:
                peaks.append([max_pos, max_value])
                drop_first_peak = True
                max_value = np.Inf
                min_value = np.Inf
                if index + lookahead >= signal.size:
                    break
                continue
        if value > min_value + delta and min_value != -np.Inf:
            if signal[index:index + lookahead].min() > min_value:
                valleys.append([min_pos, min_value])
                drop_first_valley = True
                min_value = -np.Inf
                max_value = -np.Inf
                if index + lookahead >= signal.size:
                    break
    return peaks, valleys

def autocorrelate(signal):
    """Get the auto-correlation function of a signal"""
    x = np.hstack((signal, np.zeros(len(signal))))
    sp = np.fft.rfft(x)
    tmp = np.empty_like(sp)
    tmp = np.conj(sp, tmp)
    tmp = np.multiply(tmp, sp, tmp)
    ac = np.fft.irfft(tmp)
    ac = np.divide(ac, signal.size, ac)[:signal.size/ 2]
    tmp = signal.size / (signal.size - np.arange(signal.size / 2, dtype=np.float64))
    ac = np.multiply(ac, tmp, ac)
    ac = np.concatenate([ac, np.zeros(signal.size - ac.size)])
    return normalize(ac)

def derivative(signal):
    """Return a signal that is the derivative function of a given signal"""
    def f(x):
        x = int(x)
        return signal[x]
    def df(x, h=0.1e-5):
        return (f(x + h * 0.5) - f(x - h * 0.5)) / h
    return [df(x) for x in range(len(signal))]

def integral(signal):
    """Return a signal that is the integral function of a given signal"""
    result = []
    v = 0.0
    for i in xrange(len(signal)):
        v += signal[i]
        result.append(v)
    return result

def trendline(signal):
    """Returns a  line (slope, intersect) that is the regression line given a series of values."""
    signal = list(signal)
    n = len(signal) - 1
    sum_x = 0
    sum_y = 0
    sum_xx = 0
    sum_xy = 0
    for i in range(1, n + 1):
        x = i
        y = signal[i]
        sum_x = sum_x + x
        sum_y = sum_y + y
        xx = math.pow(x, 2)
        sum_xx = sum_xx + xx
        xy = x*y
        sum_xy = sum_xy + xy
    try:
        a = (-sum_x * sum_xy + sum_xx * sum_y) / (n * sum_xx - sum_x * sum_x)
        b = (-sum_x * sum_y + n * sum_xy) / (n * sum_xx - sum_x * sum_x)
    except ZeroDivisionError:
        a, b = 0, 0
    return (b, a) # (slope, intersect)

def bandpass_filter(signal, sampling_rate, lowcut, highcut, order=6):
    """In hz"""
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    low = lowcut / nyquist
    high = highcut / nyquist
    b, a = butter(order, [low, high], btype='band')
    signal = lfilter(b, a, signal)
    return signal

def lowpass_filter(signal, sampling_rate, cutoff, order=6):
    """fyi, convolution-based smooth is probably faster"""
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    signal = lfilter(b, a, signal)
    return signal

def highpass_filter(signal, sampling_rate, cutoff, order=6):
    from scipy.signal import butter, lfilter
    nyquist = 0.5 * sampling_rate
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='high', analog=False)
    signal = lfilter(b, a, signal)
    return signal    <|endoftext|>


<commit_before>"""distutils

The main package for the Python Module Distribution Utilities.  Normally
used from a setup script as

   from distutils.core import setup

   setup (...)
"""

# Distutils version
#
# Updated automatically by the Python release process.
#
#--start constants--
__version__ = "3.3.0b2"
#--end constants--
<commit_msg>Bump to 3.3.0rc1.<commit_after>"""distutils

The main package for the Python Module Distribution Utilities.  Normally
used from a setup script as

   from distutils.core import setup

   setup (...)
"""

# Distutils version
#
# Updated automatically by the Python release process.
#
#--start constants--
__version__ = "3.3.0rc1"
#--end constants--
<|endoftext|>


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

from datetime import datetime
from functools import partial
import logging
import mimetypes
import os
from string import Template
import sys
import tempfile
from threading import Thread
import time

import asyncirc
from imgurpython import ImgurClient
from imgurpython.helpers.error import ImgurClientError
from twx import botapi

import config
from models.image import ImageInfo, ImageDatabase

CONFIG_FILE = "config.yaml"
IMAGE_EXTENSIONS = (".jpg", ".png", ".gif")

l = logging.getLogger(__name__)

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

class TelegramImageBot(botapi.TelegramBot):

    def __init__(self, conf, on_image=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._offset = None
        self.conf = conf
        self.on_image = on_image

    @staticmethod
    def build_name(user):
        return user.username or ' '.join(filter([user.first_name, user.last_name]))

    @property
    def offset(self):
        return self._offset

    @offset.setter
    def offset(self, offset):
        l.info("new offset: {}", offset)
        self._offset = offset

    def handle_updates(self, updates):
        if not updates:
            return

        for update in updates:
            upd_id, message = update.update_id, update.message

            l.debug("handling update: {}", update)

            # Out data storage object
            img = ImageInfo(f_id=None,
                            time=message.date,
                            username=message.sender.username,
                            c_id=message.chat.id, m_id=message.message_id,
                            caption=message.caption, ext='.jpg',
                            remote_path=None, local_path=None, url=None, finished=False)

            if message.document:
                # Check for image mime types
                mime_type = message.document.mime_type
                l.info("received document from {0.sender.username}: {0.document}", message)
                if mime_type:
                    ext = mimetypes.guess_extension(mime_type)
                    l.debug("guessed extension '{}' from MIME-type '{}'", ext, mime_type)
                    if ext in IMAGE_EXTENSIONS:
                        # Download document (image file)
                        img = img._replace(ext=ext, f_id=message.document.file_id)
                        self.on_image(img)
                    else:
                        self.send_message(message.chat.id, "I do not know how to handle that")

            elif message.photo:
                l.info("received photo from {1} ({0.sender.id}): {0.photo}",
                       message, self.build_name(message.sender))
                sorted_photo = sorted(message.photo, key=lambda p: p.file_size)
                if sorted_photo != message.photo:
                    l.critical("PhotoSizes were not sorted by size; {}", message)

                # Download the file (always jpg)
                img = img._replace(f_id=sorted_photo[-1].file_id)
                self.on_image(img)

            elif message.text:
                l.info("received text from {1} ({0.sender.id}): {0.text}",
                       message, self.build_name(message.sender))
                self.send_message(message.chat.id, "Just send me photos or images")

            else:
                l.warn("didn't handle update: {}", update)
                self.send_message(message.chat.id, "I do not know how to handle that")

            if not self.offset or upd_id >= self.offset:
                self.offset = upd_id + 1

    def handle_error(self, error):
        l.error("failed to fetch data; {}", error)
        # Delay next poll if there was an error
        time.sleep(self.conf.telegram.timeout or 5)

    def poll_loop(self):
        timeout = self.conf.telegram.timeout or 5
        l.info("poll loop initiated with timeout {}", timeout)

        i = 0
        while True:
            i += 1
            l.debug("poll #{}", i)

            # Long polling
            self.get_updates(
                timeout=timeout,
                offset=self.offset,
                on_success=self.handle_updates,
                on_error=self.handle_error
            ).wait()

class MyIRCClient(asyncirc.IRCClient):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._connected = False

    def _process_data(self, line):
        try:
            code = int(line.split()[1])
        except:
            pass
        else:
            if code == 376:  # End of /MOTD command
                self._connected = True
                l.info("IRC client connected as {}", self.nick)
            elif code == 433:  # Nickname is already in use
                self.nick += "_"
                self.send_raw("NICK {nick}".format(nick=self.nick))

        super()._process_data(line)

    def wait_connected(self, timeout=7):
        start = time.time()
        l.debug("Waiting for IRC client to connect")
        while time.time() < start + timeout:
            if self._connected:
                return True
            time.sleep(0.1)
        else:
            return False

class ImageReceivedThread(Thread):
    def __init__(self, conf, irc_bot, tg_bot, img, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conf = conf
        self.irc_bot = irc_bot
        self.tg_bot = tg_bot
        self.img = img

    def reply(self, msg):
        self.tg_bot.send_message(
            self.img.c_id,
            msg,
            disable_web_page_preview=True,
            reply_to_message_id=self.img.m_id,
            on_success=partial(l.info, "sent message to {0.chat.username} ({0.chat.id}): {0.text}")
        )

    def run(self):
        # Show that we're doing something
        self.tg_bot.send_chat_action(self.img.c_id, botapi.ChatAction.PHOTO)

        # Must be created in thread because multi-threading is now allowed
        db = ImageDatabase(self.conf.storage.database) if self.conf.storage.database else None

        try:
            l.debug("Running ImageReceivedThread with {}", self.img)
            # Check if we recieved the file already and see how far we got
            if db:
                db_img = db.find_image(self.img)
                if db_img:
                    self.img = db_img

            # Download file if necessary
            if not self.img.local_path or not os.path.exists(self.img.local_path):
                if not self.download_file():
                    return
            else:
                l.warn("File exists already, skipping download: {}", self.img.local_path)

            # Upload file if necessary
            if not self.img.url:
                self.upload_file()
            else:
                l.warn("File already uploaded: {}", self.img.url)

            # Post to IRC
            self.post_to_irc()

            # Report success
            self.reply("Image delivered. Uploaded to: " + self.img.url)
            self.img = self.img._replace(finished=True)

            # Cleanup
            if self.conf.storage.delete_images:
                os.remove(self.img.local_path)
                self.img = self.img._replace(local_path=None)

        except Exception as e:
            self.reply("Oops, there was an error. Contact @fichtefoll and run in circles.\n"
                       "Error: " + str(e))
            l.error("Uncaught error in ImageReceivedThread: {}", e)
            raise

        finally:
            if db:
                if not db_img:
                    db.insert_image(self.img)
                elif self.img != db_img:
                    db.update_image(self.img)
                db.close()

    def download_file(self):
        # Get file info
        file_info = self.tg_bot.get_file(self.img.f_id).wait()
        if isinstance(file_info, botapi.Error):
            msg = "Error getting file info: {}".format(file_info)
            l.error(msg)
            self.reply(msg)
            return False

        l.info("file info: {}", file_info)

        # Build file path
        directory = (Template(self.conf.storage.directory or "$temp/telegram")
                     .substitute(temp=tempfile.gettempdir()))
        directory = os.path.abspath(directory)
        basename = file_info.file_path.replace("/", "_")
        out_file = os.path.join(directory, basename)
        self.img = self.img._replace(remote_path=file_info.file_path, local_path=out_file)

        # Do download
        os.makedirs(directory, exist_ok=True)
        result = self.tg_bot.download_file(self.img.remote_path,
                                           out_file=self.img.local_path).wait()
        if isinstance(result, Exception):
            msg = "Error downloading file: {}".format(result)
            l.error(msg)
            self.reply(msg)
            return False
        else:
            l.info("Downloaded file to: {}", self.img.local_path)
            return True

    def upload_file(self):
        timestamp = datetime.fromtimestamp(self.img.time).strftime(
            self.conf.imgur.timestamp_format or "%Y-%m-%dT%H:%M:%S"
        )
        config = dict(
            album=self.conf.imgur.album,
            name="{}_{}".format(timestamp, self.img.username).replace(":", "-"),
            title="{} (by {}; {})".format(self.img.caption or "No caption",
                                          self.img.username, timestamp)
        )

        try:
            client = ImgurClient(self.conf.imgur.client_id, self.conf.imgur.client_secret,
                                 refresh_token=self.conf.imgur.refresh_token)
            data = client.upload_from_path(self.img.local_path, config=config, anon=False)
        except ImgurClientError as e:
            msg = "Error uploading to imgur: {0.status_code} {0.error_message}".format(e)
            l.error(msg)
            self.reply(msg)
            raise

        l.info("uploaded image: {}", data)
        l.debug("X-RateLimit-ClientRemaining: {}", client.credits['ClientRemaining'])

        self.img = self.img._replace(url=data['link'])
        return True

    def post_to_irc(self):
        pre_msg = ("<{{0.username}}>: {}{{0.url}}"
                   .format("{0.caption} " if self.img.caption else ""))
        msg = pre_msg.format(self.img)
        self.irc_bot.msg(self.conf.irc.channel, msg)

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

def verify_config(conf):
    if not conf.telegram.token:
        l.critical("no telegram token found")

    elif not conf.imgur.client_id or not conf.imgur.client_secret:
        l.critical("no imgur client info found")

    elif not conf.imgur.refresh_token:
        l.critical("no imgur refresh_token found. Create one with authenticate_imgur.py")

    elif not conf.irc.host or not conf.irc.channel:
        l.critical("no sufficient irc configuration found")

    else:
        return True

def init_logging(conf, console_level):
    class NewStyleLogRecord(logging.LogRecord):
        def getMessage(self):
            msg = self.msg
            if not isinstance(self.msg, str):
                msg = str(self.msg)
            if not isinstance(self.args, tuple):
                self.args = (self.args,)
            return msg.rstrip().format(*self.args)
    logging.setLogRecordFactory(NewStyleLogRecord)

    fmt = logging.Formatter("| {levelname:^8} | {message} (from {name}; {threadName})",
                            style='{')
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(fmt)
    # Filter out requests logging
    handler.addFilter(
        lambda r: (not r.name.startswith("requests")) or r.levelno > 20
    )

    handlers = [handler]

    if conf.logging.active:
        fmt = logging.Formatter(
            "| {asctime} | {levelname:^8} | {message} (from {name}; {threadName})",
            style='{'
        )
        f = open(conf.logging.path or "errors.log", "a")
        handler = logging.StreamHandler(f)
        handler.setFormatter(fmt)
        conf_level = getattr(logging, (conf.logging.level or "WARN").upper())
        handler.addFilter(lambda r: r.levelno >= conf_level)

        f.write("-- started application; logging level: {}\n".format(conf_level))
        handlers.append(handler)

    logging.basicConfig(level=console_level, handlers=handlers)
    print("-- console logging level: {}".format(l.getEffectiveLevel()))

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

def main():
    # Read config, init logging
    conf = config.read_file(CONFIG_FILE)
    init_logging(conf=conf, console_level=logging.INFO)
    l.info("config: {!s}", conf)

    # Verify other config
    if not verify_config(conf):
        return 2

    # Start IRC bot
    irc_bot = MyIRCClient(
        host=conf.irc.host,
        port=conf.irc.port or 6667,
        nick=conf.irc.nick or "TelegramBot",
        realname=conf.irc.nick,
        # use_ssl=conf.irc.ssl or False
    )
    irc_bot.start()
    if not irc_bot.wait_connected(conf.irc.timeout or 7):
        l.critical("Couldn't connect to IRC")
        return 3
    # Don't need to join channel because chanmode 'n' is not set
    irc_bot.join(conf.irc.channel)

    # Start Telegram bot
    tg_bot = TelegramImageBot(conf, token=conf.telegram.token)
    l.info("Me: {}", tg_bot.update_bot_info().wait())

    # Register image callback as a closure
    def on_image(img):
        nonlocal conf, irc_bot, tg_bot
        thread = ImageReceivedThread(
            conf=conf,
            irc_bot=irc_bot,
            tg_bot=tg_bot,
            img=img
        )
        thread.start()
        return thread

    tg_bot.on_image = on_image

    # Go through backlog and reschedule failed image uploads
    if conf.storage.database:
        with ImageDatabase(conf.storage.database) as db:
            backlog = db.get_unfinished_images()
        if backlog:
            l.info("Going through backlog, size: {}", len(backlog))
            for img in backlog:
                on_image(img).join()
            l.info("Finished backlog")

    # Main loop
    tg_bot.poll_loop()

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Also filter requests logging in file handler<commit_after>#!/usr/bin/env python3

from datetime import datetime
from functools import partial
import logging
import mimetypes
import os
from string import Template
import sys
import tempfile
from threading import Thread
import time

import asyncirc
from imgurpython import ImgurClient
from imgurpython.helpers.error import ImgurClientError
from twx import botapi

import config
from models.image import ImageInfo, ImageDatabase

CONFIG_FILE = "config.yaml"
IMAGE_EXTENSIONS = (".jpg", ".png", ".gif")

l = logging.getLogger(__name__)

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

class TelegramImageBot(botapi.TelegramBot):

    def __init__(self, conf, on_image=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._offset = None
        self.conf = conf
        self.on_image = on_image

    @staticmethod
    def build_name(user):
        return user.username or ' '.join(filter([user.first_name, user.last_name]))

    @property
    def offset(self):
        return self._offset

    @offset.setter
    def offset(self, offset):
        l.info("new offset: {}", offset)
        self._offset = offset

    def handle_updates(self, updates):
        if not updates:
            return

        for update in updates:
            upd_id, message = update.update_id, update.message

            l.debug("handling update: {}", update)

            # Out data storage object
            img = ImageInfo(f_id=None,
                            time=message.date,
                            username=message.sender.username,
                            c_id=message.chat.id, m_id=message.message_id,
                            caption=message.caption, ext='.jpg',
                            remote_path=None, local_path=None, url=None, finished=False)

            if message.document:
                # Check for image mime types
                mime_type = message.document.mime_type
                l.info("received document from {0.sender.username}: {0.document}", message)
                if mime_type:
                    ext = mimetypes.guess_extension(mime_type)
                    l.debug("guessed extension '{}' from MIME-type '{}'", ext, mime_type)
                    if ext in IMAGE_EXTENSIONS:
                        # Download document (image file)
                        img = img._replace(ext=ext, f_id=message.document.file_id)
                        self.on_image(img)
                    else:
                        self.send_message(message.chat.id, "I do not know how to handle that")

            elif message.photo:
                l.info("received photo from {1} ({0.sender.id}): {0.photo}",
                       message, self.build_name(message.sender))
                sorted_photo = sorted(message.photo, key=lambda p: p.file_size)
                if sorted_photo != message.photo:
                    l.critical("PhotoSizes were not sorted by size; {}", message)

                # Download the file (always jpg)
                img = img._replace(f_id=sorted_photo[-1].file_id)
                self.on_image(img)

            elif message.text:
                l.info("received text from {1} ({0.sender.id}): {0.text}",
                       message, self.build_name(message.sender))
                self.send_message(message.chat.id, "Just send me photos or images")

            else:
                l.warn("didn't handle update: {}", update)
                self.send_message(message.chat.id, "I do not know how to handle that")

            if not self.offset or upd_id >= self.offset:
                self.offset = upd_id + 1

    def handle_error(self, error):
        l.error("failed to fetch data; {}", error)
        # Delay next poll if there was an error
        time.sleep(self.conf.telegram.timeout or 5)

    def poll_loop(self):
        timeout = self.conf.telegram.timeout or 5
        l.info("poll loop initiated with timeout {}", timeout)

        i = 0
        while True:
            i += 1
            l.debug("poll #{}", i)

            # Long polling
            self.get_updates(
                timeout=timeout,
                offset=self.offset,
                on_success=self.handle_updates,
                on_error=self.handle_error
            ).wait()

class MyIRCClient(asyncirc.IRCClient):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._connected = False

    def _process_data(self, line):
        try:
            code = int(line.split()[1])
        except:
            pass
        else:
            if code == 376:  # End of /MOTD command
                self._connected = True
                l.info("IRC client connected as {}", self.nick)
            elif code == 433:  # Nickname is already in use
                self.nick += "_"
                self.send_raw("NICK {nick}".format(nick=self.nick))

        super()._process_data(line)

    def wait_connected(self, timeout=7):
        start = time.time()
        l.debug("Waiting for IRC client to connect")
        while time.time() < start + timeout:
            if self._connected:
                return True
            time.sleep(0.1)
        else:
            return False

class ImageReceivedThread(Thread):
    def __init__(self, conf, irc_bot, tg_bot, img, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conf = conf
        self.irc_bot = irc_bot
        self.tg_bot = tg_bot
        self.img = img

    def reply(self, msg):
        self.tg_bot.send_message(
            self.img.c_id,
            msg,
            disable_web_page_preview=True,
            reply_to_message_id=self.img.m_id,
            on_success=partial(l.info, "sent message to {0.chat.username} ({0.chat.id}): {0.text}")
        )

    def run(self):
        # Show that we're doing something
        self.tg_bot.send_chat_action(self.img.c_id, botapi.ChatAction.PHOTO)

        # Must be created in thread because multi-threading is now allowed
        db = ImageDatabase(self.conf.storage.database) if self.conf.storage.database else None

        try:
            l.debug("Running ImageReceivedThread with {}", self.img)
            # Check if we recieved the file already and see how far we got
            if db:
                db_img = db.find_image(self.img)
                if db_img:
                    self.img = db_img

            # Download file if necessary
            if not self.img.local_path or not os.path.exists(self.img.local_path):
                if not self.download_file():
                    return
            else:
                l.warn("File exists already, skipping download: {}", self.img.local_path)

            # Upload file if necessary
            if not self.img.url:
                self.upload_file()
            else:
                l.warn("File already uploaded: {}", self.img.url)

            # Post to IRC
            self.post_to_irc()

            # Report success
            self.reply("Image delivered. Uploaded to: " + self.img.url)
            self.img = self.img._replace(finished=True)

            # Cleanup
            if self.conf.storage.delete_images:
                os.remove(self.img.local_path)
                self.img = self.img._replace(local_path=None)

        except Exception as e:
            self.reply("Oops, there was an error. Contact @fichtefoll and run in circles.\n"
                       "Error: " + str(e))
            l.error("Uncaught error in ImageReceivedThread: {}", e)
            raise

        finally:
            if db:
                if not db_img:
                    db.insert_image(self.img)
                elif self.img != db_img:
                    db.update_image(self.img)
                db.close()

    def download_file(self):
        # Get file info
        file_info = self.tg_bot.get_file(self.img.f_id).wait()
        if isinstance(file_info, botapi.Error):
            msg = "Error getting file info: {}".format(file_info)
            l.error(msg)
            self.reply(msg)
            return False

        l.info("file info: {}", file_info)

        # Build file path
        directory = (Template(self.conf.storage.directory or "$temp/telegram")
                     .substitute(temp=tempfile.gettempdir()))
        directory = os.path.abspath(directory)
        basename = file_info.file_path.replace("/", "_")
        out_file = os.path.join(directory, basename)
        self.img = self.img._replace(remote_path=file_info.file_path, local_path=out_file)

        # Do download
        os.makedirs(directory, exist_ok=True)
        result = self.tg_bot.download_file(self.img.remote_path,
                                           out_file=self.img.local_path).wait()
        if isinstance(result, Exception):
            msg = "Error downloading file: {}".format(result)
            l.error(msg)
            self.reply(msg)
            return False
        else:
            l.info("Downloaded file to: {}", self.img.local_path)
            return True

    def upload_file(self):
        timestamp = datetime.fromtimestamp(self.img.time).strftime(
            self.conf.imgur.timestamp_format or "%Y-%m-%dT%H:%M:%S"
        )
        config = dict(
            album=self.conf.imgur.album,
            name="{}_{}".format(timestamp, self.img.username).replace(":", "-"),
            title="{} (by {}; {})".format(self.img.caption or "No caption",
                                          self.img.username, timestamp)
        )

        try:
            client = ImgurClient(self.conf.imgur.client_id, self.conf.imgur.client_secret,
                                 refresh_token=self.conf.imgur.refresh_token)
            data = client.upload_from_path(self.img.local_path, config=config, anon=False)
        except ImgurClientError as e:
            msg = "Error uploading to imgur: {0.status_code} {0.error_message}".format(e)
            l.error(msg)
            self.reply(msg)
            raise

        l.info("uploaded image: {}", data)
        l.debug("X-RateLimit-ClientRemaining: {}", client.credits['ClientRemaining'])

        self.img = self.img._replace(url=data['link'])
        return True

    def post_to_irc(self):
        pre_msg = ("<{{0.username}}>: {}{{0.url}}"
                   .format("{0.caption} " if self.img.caption else ""))
        msg = pre_msg.format(self.img)
        self.irc_bot.msg(self.conf.irc.channel, msg)

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

def verify_config(conf):
    if not conf.telegram.token:
        l.critical("no telegram token found")

    elif not conf.imgur.client_id or not conf.imgur.client_secret:
        l.critical("no imgur client info found")

    elif not conf.imgur.refresh_token:
        l.critical("no imgur refresh_token found. Create one with authenticate_imgur.py")

    elif not conf.irc.host or not conf.irc.channel:
        l.critical("no sufficient irc configuration found")

    else:
        return True

def init_logging(conf, console_level):
    class NewStyleLogRecord(logging.LogRecord):
        def getMessage(self):
            msg = self.msg
            if not isinstance(self.msg, str):
                msg = str(self.msg)
            if not isinstance(self.args, tuple):
                self.args = (self.args,)
            return msg.rstrip().format(*self.args)
    logging.setLogRecordFactory(NewStyleLogRecord)

    fmt = logging.Formatter("| {levelname:^8} | {message} (from {name}; {threadName})",
                            style='{')
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(fmt)
    # Filter out requests logging
    handler.addFilter(
        lambda r: (not r.name.startswith("requests")) or r.levelno > 20
    )

    handlers = [handler]

    if conf.logging.active:
        fmt = logging.Formatter(
            "| {asctime} | {levelname:^8} | {message} (from {name}; {threadName})",
            style='{'
        )
        f = open(conf.logging.path or "errors.log", "a")
        handler = logging.StreamHandler(f)
        handler.setFormatter(fmt)
        conf_level = getattr(logging, (conf.logging.level or "WARN").upper())
        handler.addFilter(lambda r: r.levelno >= conf_level)
        # Filter out requests logging
        handler.addFilter(
            lambda r: (not r.name.startswith("requests")) or r.levelno > 20
        )

        f.write("-- started application; logging level: {}\n".format(conf_level))
        handlers.append(handler)

    logging.basicConfig(level=console_level, handlers=handlers)
    print("-- console logging level: {}".format(l.getEffectiveLevel()))

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

def main():
    # Read config, init logging
    conf = config.read_file(CONFIG_FILE)
    init_logging(conf=conf, console_level=logging.INFO)
    l.info("config: {!s}", conf)

    # Verify other config
    if not verify_config(conf):
        return 2

    # Start IRC bot
    irc_bot = MyIRCClient(
        host=conf.irc.host,
        port=conf.irc.port or 6667,
        nick=conf.irc.nick or "TelegramBot",
        realname=conf.irc.nick,
        # use_ssl=conf.irc.ssl or False
    )
    irc_bot.start()
    if not irc_bot.wait_connected(conf.irc.timeout or 7):
        l.critical("Couldn't connect to IRC")
        return 3
    # Don't need to join channel because chanmode 'n' is not set
    irc_bot.join(conf.irc.channel)

    # Start Telegram bot
    tg_bot = TelegramImageBot(conf, token=conf.telegram.token)
    l.info("Me: {}", tg_bot.update_bot_info().wait())

    # Register image callback as a closure
    def on_image(img):
        nonlocal conf, irc_bot, tg_bot
        thread = ImageReceivedThread(
            conf=conf,
            irc_bot=irc_bot,
            tg_bot=tg_bot,
            img=img
        )
        thread.start()
        return thread

    tg_bot.on_image = on_image

    # Go through backlog and reschedule failed image uploads
    if conf.storage.database:
        with ImageDatabase(conf.storage.database) as db:
            backlog = db.get_unfinished_images()
        if backlog:
            l.info("Going through backlog, size: {}", len(backlog))
            for img in backlog:
                on_image(img).join()
            l.info("Finished backlog")

    # Main loop
    tg_bot.poll_loop()

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


<commit_before>###############################################################################
# Caleydo - Visualization for Molecular Biology - http://caleydo.org
# Copyright (c) The Caleydo Team. All rights reserved.
# Licensed under the new BSD license, available at http://caleydo.org/license
###############################################################################

import os
import sys
sys.path.insert(0, os.path.abspath('/phovea/phovea_server/phovea_server'))

if __name__ == '__main__':
  import launch
  launch.run()
else:
  import launch
  application = launch.create_embedded()
  print('run as embedded version: %s', application)
<commit_msg>removed hard coded file path<commit_after>###############################################################################
# Caleydo - Visualization for Molecular Biology - http://caleydo.org
# Copyright (c) The Caleydo Team. All rights reserved.
# Licensed under the new BSD license, available at http://caleydo.org/license
###############################################################################

import os
import sys
path = os.path.dirname(os.path.realpath(__file__))
# required to import files in python 3.7
sys.path.insert(0, os.path.abspath(path + '/phovea_server'))

if __name__ == '__main__':
  import launch
  launch.run()
else:
  import launch
  application = launch.create_embedded()
  print('run as embedded version: %s', application)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

from domaintools.exceptions import NotFoundException
from domaintools.exceptions import NotAuthorizedException
from domaintools.exceptions import ServiceUnavailableException

from domaintools import API

from cortexutils.analyzer import Analyzer

class DomainToolsAnalyzer(Analyzer):

    def __init__(self):
        Analyzer.__init__(self)
        self.service = self.get_param(
            'config.service', None, 'Service parameter is missing')

    def domaintools(self, data):
        """

        :param service:
        :return:
        """
        if (self.service == 'reverse-ip' and self.data_type == 'ip'):
            self.service = 'host-domains'

        api = API(self.get_param('config.username'), self.get_param('config.key'))

        if self.service == 'reverse-ip' and self.data_type in ['domain', 'ip', 'fqdn']:
            response = api.reverse_ip(data).response()

        elif self.service == 'host-domains' and self.data_type == 'ip':
            response = api.host_domains(data).response()

        elif self.service == 'name-server-domains' and self.data_type == 'domain':
            response = api.reverse_name_server(data).response()

        elif self.service == 'whois/history' and self.data_type == 'domain':
            response = api.whois_history(data).response()

        elif self.service == 'whois/parsed' and self.data_type == 'domain':
            response = api.parsed_whois(data).response()

        elif self.service == 'risk_evidence' and self.data_type in ['domain', 'fqdn']:
            response = api.risk_evidence(data).response()

        elif self.service == 'reputation' and self.data_type in ['domain', 'fqdn']:
            response = api.reputation(data, include_reasons=True).response()

        elif self.service == 'reverse-whois':
            response = api.reverse_whois(data, mode='purchase').response()

        elif self.service == 'whois' and self.data_type == 'ip':
            response = api.whois(data).response()

        return response

    def summary(self, raw):
        r = {
            "service": self.service,
            "dataType": self.data_type
        }

        if "ip_addresses" in raw:
            if type(raw["ip_addresses"]) == dict:
                r["ip"] = {
                    "address": raw["ip_addresses"]["ip_address"],
                    "domain_count": raw["ip_addresses"]["domain_count"]
                }
            elif type(raw["ip_addresses"]) == list:
                r["ip"] = {
                    "address": "{} IP addresses".format(len(r)),
                    "domain_count": sum(d["domain_count"] for d in raw["ip_addresses"])
                }

        if "domain_count" in raw:
            r["domain_count"] = {
                "current": raw["domain_count"]["current"],
                "historic": raw["domain_count"]["historic"]
            }

        if "registrant" in raw:
            r["registrant"] = raw["registrant"]
        elif "response" in raw and "registrant" in raw["response"]:
            r["registrant"] = raw["response"]["registrant"]

        if "parsed_whois" in raw:
            r["registrar"] = raw["parsed_whois"]["registrar"]["name"]
            #

        if "name_server" in raw:
            r["name_server"] = raw["name_server"]["hostname"]
            r["domain_count"] = raw["name_server"]["total"]

        if "risk_score" in raw:
            r["risk_score"] = raw["risk_score"]
            if "reasons" in raw:
                r["reputation"] = True

        taxonomies = []

        # Prepare predicate and value for each service
        if r["service"] in ["reverse-ip", "host-domains"]:
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_IP",
                                                  "{}, {} domains".format(r["ip"]["address"],
                                                                              r["ip"]["domain_count"])))

        if r["service"] == "name-server-domains":
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_Name_Server",
                                                  "{}, {} domains".format(r["name_server"], r["domain_count"])))

        if r["service"] == "reverse-whois":
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_Whois",
                                                  "curr:{} / hist:{} domains".format(r["domain_count"]["current"],
                                                                                r["domain_count"][
                                                                                "historic"])))

        if r["service"] == "whois/history":
            taxonomies.append(self.build_taxonomy("info", "DT", "Whois_History",
                                                  "{}, {} domains ".format(r["name_server"], r["domain_count"])))

        if r["service"] == "whois/parsed" or r['service'] == "whois":
            if r["registrar"]:
                taxonomies.append(self.build_taxonomy("info", "DT", "Whois", "REGISTRAR:{}".format(r["registrar"])))
            if r["registrant"]:
                taxonomies.append(
                    self.build_taxonomy("info", "DT", "Whois", "REGISTRANT:{}".format(r["registrant"])))

        if "risk_score" in r:
            risk_service = "Risk"
            if "reputation" in r:
                risk_service = "Reputation"
            if r["risk_score"] == 0:
                level = "safe"
            elif 0 < r["risk_score"] <= 50:
                level = "suspicious"
            elif r["risk_score"] > 50:
                level = "malicious"
            taxonomies.append(
                self.build_taxonomy(level, "DT", risk_service, "{}".format(r["risk_score"])))

        result = {'taxonomies': taxonomies}
        return result

    def run(self):
        data = self.get_data()

        try:
            r = self.domaintools(data)

            if 'response' in r:
                self.report(r.get('response'))
            elif 'error' in r and 'message' in r['error']:
                self.error(r['error']['message'])
            else:
                self.report(r)

        except NotFoundException:
            self.error(self.data_type.capitalize() + " not found")
        except NotAuthorizedException:
            self.error("An authorization error occurred")
        except ServiceUnavailableException:
            self.error("DomainTools Service is currenlty unavailable")
        except Exception as e:
            self.unexpectedError(e)

        else:
            self.error('Unknown DomainTools service or invalid data type')

if __name__ == '__main__':
    DomainToolsAnalyzer().run()
<commit_msg>#339 fix short report in summary()<commit_after>#!/usr/bin/env python
# encoding: utf-8

from domaintools.exceptions import NotFoundException
from domaintools.exceptions import NotAuthorizedException
from domaintools.exceptions import ServiceUnavailableException

from domaintools import API

from cortexutils.analyzer import Analyzer

class DomainToolsAnalyzer(Analyzer):

    def __init__(self):
        Analyzer.__init__(self)
        self.service = self.get_param(
            'config.service', None, 'Service parameter is missing')

    def domaintools(self, data):
        """

        :param service:
        :return:
        """
        if (self.service == 'reverse-ip' and self.data_type == 'ip'):
            self.service = 'host-domains'

        api = API(self.get_param('config.username'), self.get_param('config.key'))

        if self.service == 'reverse-ip' and self.data_type in ['domain', 'ip', 'fqdn']:
            response = api.reverse_ip(data).response()

        elif self.service == 'host-domains' and self.data_type == 'ip':
            response = api.host_domains(data).response()

        elif self.service == 'name-server-domains' and self.data_type == 'domain':
            response = api.reverse_name_server(data).response()

        elif self.service == 'whois/history' and self.data_type == 'domain':
            response = api.whois_history(data).response()

        elif self.service == 'whois/parsed' and self.data_type == 'domain':
            response = api.parsed_whois(data).response()

        elif self.service == 'risk_evidence' and self.data_type in ['domain', 'fqdn']:
            response = api.risk_evidence(data).response()

        elif self.service == 'reputation' and self.data_type in ['domain', 'fqdn']:
            response = api.reputation(data, include_reasons=True).response()

        elif self.service == 'reverse-whois':
            response = api.reverse_whois(data, mode='purchase').response()

        elif self.service == 'whois' and self.data_type == 'ip':
            response = api.whois(data).response()

        return response

    def summary(self, raw):
        r = {
            "service": self.service,
            "dataType": self.data_type
        }

        if "ip_addresses" in raw:
            if type(raw["ip_addresses"]) == dict:
                r["ip"] = {
                    "address": raw["ip_addresses"]["ip_address"],
                    "domain_count": raw["ip_addresses"]["domain_count"]
                }
            elif type(raw["ip_addresses"]) == list:
                r["ip"] = {
                    "address": "{} IP addresses".format(len(r)),
                    "domain_count": sum(d["domain_count"] for d in raw["ip_addresses"])
                }

        if "record_count" in raw:
            r["record_count"] = raw.get('record_count')

        if "domain_count" in raw:
            r["domain_count"] = {
                "current": raw["domain_count"]["current"],
                "historic": raw["domain_count"]["historic"]
            }

        if "registrant" in raw:
            r["registrant"] = raw["registrant"]
        elif "response" in raw and "registrant" in raw["response"]:
            r["registrant"] = raw["response"]["registrant"]

        if "parsed_whois" in raw:
            r["registrar"] = raw["parsed_whois"]["registrar"]["name"]
            #

        if "name_server" in raw:
            r["name_server"] = raw["name_server"]["hostname"]
            r["domain_count"] = raw["name_server"]["total"]

        if "risk_score" in raw:
            r["risk_score"] = raw["risk_score"]
            if "reasons" in raw:
                r["reputation"] = True

        taxonomies = []

        # Prepare predicate and value for each service
        if r["service"] in ["reverse-ip", "host-domains"]:
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_IP",
                                                  "{}, {} domains".format(r["ip"]["address"],
                                                                              r["ip"]["domain_count"])))

        if r["service"] == "name-server-domains":
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_Name_Server",
                                                  "{}, {} domains".format(r["name_server"], r["domain_count"])))

        if r["service"] == "reverse-whois":
            taxonomies.append(self.build_taxonomy("info", "DT", "Reverse_Whois",
                                                  "curr:{} / hist:{} domains".format(r["domain_count"]["current"],
                                                                                r["domain_count"][
                                                                                "historic"])))

        if r["service"] == "whois/history":
            taxonomies.append(self.build_taxonomy("info", "DT", "Whois_History",
                                                  "{} {}".format(r["record_count"], "records" if r["record_count"] > 1 else "record")))

        if r["service"] == "whois/parsed" or r['service'] == "whois":
            if r["registrar"]:
                taxonomies.append(self.build_taxonomy("info", "DT", "Whois", "REGISTRAR:{}".format(r["registrar"])))
            if r["registrant"]:
                taxonomies.append(
                    self.build_taxonomy("info", "DT", "Whois", "REGISTRANT:{}".format(r["registrant"])))

        if "risk_score" in r:
            risk_service = "Risk"
            if "reputation" in r:
                risk_service = "Reputation"
            if r["risk_score"] == 0:
                level = "safe"
            elif 0 < r["risk_score"] <= 50:
                level = "suspicious"
            elif r["risk_score"] > 50:
                level = "malicious"
            taxonomies.append(
                self.build_taxonomy(level, "DT", risk_service, "{}".format(r["risk_score"])))

        result = {'taxonomies': taxonomies}
        return result

    def run(self):
        data = self.get_data()

        try:
            r = self.domaintools(data)

            if 'response' in r:
                self.report(r.get('response'))
            elif 'error' in r and 'message' in r['error']:
                self.error(r['error']['message'])
            else:
                self.report(r)

        except NotFoundException:
            self.error(self.data_type.capitalize() + " not found")
        except NotAuthorizedException:
            self.error("An authorization error occurred")
        except ServiceUnavailableException:
            self.error("DomainTools Service is currenlty unavailable")
        except Exception as e:
            self.unexpectedError(e)

        else:
            self.error('Unknown DomainTools service or invalid data type')

if __name__ == '__main__':
    DomainToolsAnalyzer().run()
<|endoftext|>


<commit_before>from Cookie import Cookie
from pazzo.exceptions import PazzoException
from pazzo.sessions.stores.store_registry import StoreRegistry
from pazzo.sessions.wsgi_session import Session
from werkzeug.http import cookie_date
import time

def me_want_cookie(key, value='', max_age=None, expires=None, path='/',
                   domain=None, secure=False, httponly=False):
        """
        Sets a cookie.

        ``expires`` can be:
        - a string in the correct format

        """

        cookie = Cookie()
        cookie[key] = value

        if expires is not None:
            cookie[key]['expires'] = expires
        if max_age is not None:
            cookie[key]['max-age'] = max_age
            # IE requires expires, so set it if hasn't been already.
            if not expires:
                cookie[key]['expires'] = cookie_date(time.time() +
                                                           max_age)
        if path is not None:
            cookie[key]['path'] = path
        if domain is not None:
            cookie[key]['domain'] = domain
        if secure:
            cookie[key]['secure'] = True
        if httponly:
            cookie[key]['httponly'] = True

        return cookie

class SessionMiddleware(object):
    """ Pazzo session middleware for WSGI apps
    Handles loading, saving, expiry so app doesn't
    have to
    """
    def __init__(self, app):
        store_registry = StoreRegistry()
        if not store_registry.initialized:
            raise PazzoException("Session store registry not initialized")
        self.settings = store_registry.settings
        self.app = app
        self.registry = store_registry

    def __call__(self, environ, start_response):
        # just make sure it exists and throw it in the environ
        session = Session(environ)
        if session.session_key is None:
            session.create()
        environ[self.settings.SESSION_WSGI_ENVIRON_NAME] = session

        def session_start_response(status, headers, exc_info=None):
            session = Session(environ)
            if session.modified:
                if session.get_expire_at_browser_close():
                    max_age = None
                    expires = None
                else:
                    max_age = session.get_expiry_age()
                    expires_time = time.time() + max_age
                    expires = cookie_date(expires_time)
                cookie = me_want_cookie('sessionid',
                            value=session.session_key,
                            max_age=max_age,
                            expires=expires,
                            path=self.settings.SESSION_COOKIE_PATH,
                            domain=self.settings.SESSION_COOKIE_DOMAIN,
                            secure=self.settings.SESSION_COOKIE_SECURE,
                            httponly=self.settings.SESSION_COOKIE_HTTPONLY)
                headers.append(("Set-Cookie", cookie.output(header='')))
            return start_response(status, headers, exc_info)

        return self.app(environ, session_start_response)
<commit_msg>Clarifies wsgi middleware docstring<commit_after>from Cookie import Cookie
from pazzo.exceptions import PazzoException
from pazzo.sessions.stores.store_registry import StoreRegistry
from pazzo.sessions.wsgi_session import Session
from werkzeug.http import cookie_date
import time

def me_want_cookie(key, value='', max_age=None, expires=None, path='/',
                   domain=None, secure=False, httponly=False):
        """
        Sets a cookie.

        ``expires`` can be:
        - a string in the correct format

        """

        cookie = Cookie()
        cookie[key] = value

        if expires is not None:
            cookie[key]['expires'] = expires
        if max_age is not None:
            cookie[key]['max-age'] = max_age
            # IE requires expires, so set it if hasn't been already.
            if not expires:
                cookie[key]['expires'] = cookie_date(time.time() +
                                                           max_age)
        if path is not None:
            cookie[key]['path'] = path
        if domain is not None:
            cookie[key]['domain'] = domain
        if secure:
            cookie[key]['secure'] = True
        if httponly:
            cookie[key]['httponly'] = True

        return cookie

class SessionMiddleware(object):
    """ Pazzo session middleware for WSGI apps
    Handles loading, saving, cookie expiry, etc.
    so app doesn't have to deal with these things
    """
    def __init__(self, app):
        store_registry = StoreRegistry()
        if not store_registry.initialized:
            raise PazzoException("Session store registry not initialized")
        self.settings = store_registry.settings
        self.app = app
        self.registry = store_registry

    def __call__(self, environ, start_response):
        # just make sure it exists and throw it in the environ
        session = Session(environ)
        if session.session_key is None:
            session.create()
        environ[self.settings.SESSION_WSGI_ENVIRON_NAME] = session

        def session_start_response(status, headers, exc_info=None):
            session = Session(environ)
            if session.modified:
                if session.get_expire_at_browser_close():
                    max_age = None
                    expires = None
                else:
                    max_age = session.get_expiry_age()
                    expires_time = time.time() + max_age
                    expires = cookie_date(expires_time)
                cookie = me_want_cookie('sessionid',
                            value=session.session_key,
                            max_age=max_age,
                            expires=expires,
                            path=self.settings.SESSION_COOKIE_PATH,
                            domain=self.settings.SESSION_COOKIE_DOMAIN,
                            secure=self.settings.SESSION_COOKIE_SECURE,
                            httponly=self.settings.SESSION_COOKIE_HTTPONLY)
                headers.append(("Set-Cookie", cookie.output(header='')))
            return start_response(status, headers, exc_info)

        return self.app(environ, session_start_response)
<|endoftext|>


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

"""Runs Aerospike (http://www.aerospike.com).

Aerospike is an opensource NoSQL solution. This benchmark runs a read/update
load test with varying numbers of client threads against an Aerospike server.

This test can be run in a variety of configurations including memory only,
remote/persistent ssd, and local ssd. The Aerospike configuration is controlled
by the "aerospike_storage_type", "scratch_disk_type", and "use_local_disk"
flags.
"""

import re
import time

import gflags as flags

from perfkitbenchmarker import data
from perfkitbenchmarker import vm_util

FLAGS = flags.FLAGS

MEMORY = 'memory'
DISK = 'disk'
flags.DEFINE_enum('aerospike_storage_type', MEMORY, [MEMORY, DISK],
                  'The type of storage to use for Aerospike data. The type of '
                  'disk is controlled by a combination of the '
                  '"scratch_disk_type" and "use_local_disk" flags.')

BENCHMARK_INFO = {'name': 'aerospike',
                  'description': 'Runs Aerospike',
                  'num_machines': 2}
PACKAGES = ('build-essential git autoconf libtool libssl-dev lua5.1 '
            'liblua5.1-dev')
AEROSPIKE_SERVER = 'https://github.com/aerospike/aerospike-server.git'
AEROSPIKE_CLIENT = 'https://github.com/aerospike/aerospike-client-c.git'
CLIENT_DIR = 'aerospike-client-c'
CLIENT_VERSION = '3.0.84'
SERVER_DIR = 'aerospike-server'
SERVER_VERSION = '3.3.19'
READ_PERCENT = 90
MAX_THREADS = 129
MIN_THREADS = 8
THREAD_STEP = 8

def GetInfo():
  if FLAGS.aerospike_storage_type == DISK and not FLAGS.use_local_disk:
    BENCHMARK_INFO['scratch_disk'] = True
  else:
    BENCHMARK_INFO['scratch_disk'] = False
  return BENCHMARK_INFO

def _PrepareClient(client):
  """Prepare the Aerospike C client on a VM."""
  client.InstallPackage(PACKAGES)
  clone_command = 'git clone %s'
  client.RemoteCommand(clone_command % AEROSPIKE_CLIENT)
  build_command = ('cd %s && git checkout %s && git submodule update --init '
                   '&& make')
  client.RemoteCommand(build_command % (CLIENT_DIR, CLIENT_VERSION))

  # Apply a patch to the client benchmark so we have access to average latency
  # of requests. Switching over to YCSB should obviate this.
  client.PushDataFile('aerospike.patch')
  benchmark_dir = '%s/benchmarks/src/main' % CLIENT_DIR
  client.RemoteCommand('cp aerospike.patch %s' % benchmark_dir)
  client.RemoteCommand('cd %s && patch -p1 -f  < aerospike.patch'
                       % benchmark_dir)
  client.RemoteCommand('sed -i -e "s/lpthread/lpthread -lz/" '
                       '%s/benchmarks/Makefile' % CLIENT_DIR)
  client.RemoteCommand('cd %s/benchmarks && make' % CLIENT_DIR)

def _PrepareServer(server):
  """Prepare the Aerospike server on a VM."""
  server.InstallPackage(PACKAGES)
  clone_command = 'git clone %s'
  server.RemoteCommand(clone_command % AEROSPIKE_SERVER)
  build_command = ('cd %s && git checkout %s && git submodule update --init '
                   '&& make')
  server.RemoteCommand(build_command % (SERVER_DIR, SERVER_VERSION))

  if FLAGS.aerospike_storage_type == DISK:
    if FLAGS.use_local_disk:
      devices = server.GetLocalDrives()
    else:
      devices = [disk.GetDevicePath() for disk in server.scratch_disks]

    server.RenderTemplate(data.ResourcePath('aerospike.conf.j2'),
                          'aerospike-server/as/etc/aerospike_dev.conf',
                          {'devices': devices})

  for disk in server.scratch_disks:
    server.RemoteCommand('sudo umount %s' % disk.mount_point)

  server.RemoteCommand('cd %s && make init' % SERVER_DIR)
  server.RemoteCommand(
      'cd %s; nohup sudo make start &> /dev/null &' % SERVER_DIR)
  time.sleep(5)  # Wait for server to come up

def Prepare(benchmark_spec):
  """Install Aerospike server on one VM and Aerospike C client on the other.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  """
  server, client = benchmark_spec.vms

  def _Prepare(vm):
    if vm == client:
      _PrepareClient(vm)
    else:
      _PrepareServer(vm)

  vm_util.RunThreaded(_Prepare, benchmark_spec.vms)

def Run(benchmark_spec):
  """Runs a read/update load test on Aerospike.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  server, client = benchmark_spec.vms
  samples = []

  def ParseOutput(output):
    """Parses Aerospike output.

    Args:
      output: The stdout from running the benchmark.

    Returns:
      A tuple of average TPS and average latency.
    """
    write_latency, read_latency = re.findall(
        r'Overall Average Latency \(ms\) ([0-9]+\.[0-9]+)', output)[-2:]
    average_latency = ((READ_PERCENT / 100.0) * float(read_latency) +
                       ((100 - READ_PERCENT) / 100.0) * float(write_latency))
    tps = map(int, re.findall(r'total\(tps=([0-9]+)', output)[:-1])
    return float(sum(tps) / len(tps)), average_latency

  load_command = ('./%s/benchmarks/target/benchmarks -z 32 -n test -w I '
                  '-o B:1000 -k 1000000 -h %s' %
                  (CLIENT_DIR, server.internal_ip))
  client.RemoteCommand(load_command, should_log=True)

  for threads in range(MIN_THREADS, MAX_THREADS, THREAD_STEP):
    load_command = ('timeout 15 ./%s/benchmarks/target/benchmarks '
                    '-z %s -n test -w RU,%s -o B:1000 -k 1000000 '
                    '--latency 5,1 -h %s;:' %
                    (CLIENT_DIR, threads, READ_PERCENT,
                     server.internal_ip))
    stdout, _ = client.RemoteCommand(load_command, should_log=True)
    tps, latency = ParseOutput(stdout)

    metadata = {
        'Average Transactions Per Second': tps,
        'Client Threads': threads,
        'Storage Type': FLAGS.aerospike_storage_type,
    }
    if FLAGS.aerospike_storage_type == DISK:
      metadata['Disk Type'] = 'Local' if FLAGS.use_local_disk else 'Remote'
    samples.append(('Average Latency', latency, 'ms', metadata))

  return samples

def Cleanup(benchmark_spec):
  """Cleanup Aerospike.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  server, client = benchmark_spec.vms

  client.RemoteCommand('sudo rm -rf aerospike*')
  server.RemoteCommand('cd %s && nohup sudo make stop' % SERVER_DIR)
  server.RemoteCommand('sudo rm -rf aerospike*')
<commit_msg>Responded to more review comments<commit_after># Copyright 2014 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs Aerospike (http://www.aerospike.com).

Aerospike is an opensource NoSQL solution. This benchmark runs a read/update
load test with varying numbers of client threads against an Aerospike server.

This test can be run in a variety of configurations including memory only,
remote/persistent ssd, and local ssd. The Aerospike configuration is controlled
by the "aerospike_storage_type", "scratch_disk_type", and "use_local_disk"
flags.
"""

import re
import time

import gflags as flags

from perfkitbenchmarker import data
from perfkitbenchmarker import vm_util

FLAGS = flags.FLAGS

MEMORY = 'memory'
DISK = 'disk'
flags.DEFINE_enum('aerospike_storage_type', MEMORY, [MEMORY, DISK],
                  'The type of storage to use for Aerospike data. The type of '
                  'disk is controlled by a combination of the '
                  '"scratch_disk_type" and "use_local_disk" flags.')

BENCHMARK_INFO = {'name': 'aerospike',
                  'description': 'Runs Aerospike',
                  'num_machines': 2}
PACKAGES = ('build-essential git autoconf libtool libssl-dev lua5.1 '
            'liblua5.1-dev')
AEROSPIKE_SERVER = 'https://github.com/aerospike/aerospike-server.git'
AEROSPIKE_CLIENT = 'https://github.com/aerospike/aerospike-client-c.git'
CLIENT_DIR = 'aerospike-client-c'
CLIENT_VERSION = '3.0.84'
SERVER_DIR = 'aerospike-server'
SERVER_VERSION = '3.3.19'
READ_PERCENT = 90
MAX_THREADS = 128
MIN_THREADS = 8
THREAD_STEP = 8

def GetInfo():
  if FLAGS.aerospike_storage_type == DISK and not FLAGS.use_local_disk:
    BENCHMARK_INFO['scratch_disk'] = True
  else:
    BENCHMARK_INFO['scratch_disk'] = False
  return BENCHMARK_INFO

def _PrepareClient(client):
  """Prepare the Aerospike C client on a VM."""
  client.InstallPackage(PACKAGES)
  clone_command = 'git clone %s'
  client.RemoteCommand(clone_command % AEROSPIKE_CLIENT)
  build_command = ('cd %s && git checkout %s && git submodule update --init '
                   '&& make')
  client.RemoteCommand(build_command % (CLIENT_DIR, CLIENT_VERSION))

  # Apply a patch to the client benchmark so we have access to average latency
  # of requests. Switching over to YCSB should obviate this.
  client.PushDataFile('aerospike.patch')
  benchmark_dir = '%s/benchmarks/src/main' % CLIENT_DIR
  client.RemoteCommand('cp aerospike.patch %s' % benchmark_dir)
  client.RemoteCommand('cd %s && patch -p1 -f  < aerospike.patch'
                       % benchmark_dir)
  client.RemoteCommand('sed -i -e "s/lpthread/lpthread -lz/" '
                       '%s/benchmarks/Makefile' % CLIENT_DIR)
  client.RemoteCommand('cd %s/benchmarks && make' % CLIENT_DIR)

def _PrepareServer(server):
  """Prepare the Aerospike server on a VM."""
  server.InstallPackage(PACKAGES)
  clone_command = 'git clone %s'
  server.RemoteCommand(clone_command % AEROSPIKE_SERVER)
  build_command = ('cd %s && git checkout %s && git submodule update --init '
                   '&& make')
  server.RemoteCommand(build_command % (SERVER_DIR, SERVER_VERSION))

  if FLAGS.aerospike_storage_type == DISK:
    if FLAGS.use_local_disk:
      devices = server.GetLocalDrives()
    else:
      devices = [disk.GetDevicePath() for disk in server.scratch_disks]

    server.RenderTemplate(data.ResourcePath('aerospike.conf.j2'),
                          'aerospike-server/as/etc/aerospike_dev.conf',
                          {'devices': devices})

  for disk in server.scratch_disks:
    server.RemoteCommand('sudo umount %s' % disk.mount_point)

  server.RemoteCommand('cd %s && make init' % SERVER_DIR)
  server.RemoteCommand(
      'cd %s; nohup sudo make start &> /dev/null &' % SERVER_DIR)
  time.sleep(5)  # Wait for server to come up

def Prepare(benchmark_spec):
  """Install Aerospike server on one VM and Aerospike C client on the other.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  """
  server, client = benchmark_spec.vms

  def _Prepare(vm):
    if vm == client:
      _PrepareClient(vm)
    else:
      _PrepareServer(vm)

  vm_util.RunThreaded(_Prepare, benchmark_spec.vms)

def Run(benchmark_spec):
  """Runs a read/update load test on Aerospike.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  server, client = benchmark_spec.vms
  samples = []

  def ParseOutput(output):
    """Parses Aerospike output.

    Args:
      output: The stdout from running the benchmark.

    Returns:
      A tuple of average TPS and average latency.
    """
    write_latency, read_latency = re.findall(
        r'Overall Average Latency \(ms\) ([0-9]+\.[0-9]+)', output)[-2:]
    average_latency = ((READ_PERCENT / 100.0) * float(read_latency) +
                       ((100 - READ_PERCENT) / 100.0) * float(write_latency))
    tps = map(int, re.findall(r'total\(tps=([0-9]+)', output)[:-1])
    return float(sum(tps)) / len(tps), average_latency

  load_command = ('./%s/benchmarks/target/benchmarks -z 32 -n test -w I '
                  '-o B:1000 -k 1000000 -h %s' %
                  (CLIENT_DIR, server.internal_ip))
  client.RemoteCommand(load_command, should_log=True)

  for threads in range(MIN_THREADS, MAX_THREADS + 1, THREAD_STEP):
    load_command = ('timeout 15 ./%s/benchmarks/target/benchmarks '
                    '-z %s -n test -w RU,%s -o B:1000 -k 1000000 '
                    '--latency 5,1 -h %s;:' %
                    (CLIENT_DIR, threads, READ_PERCENT,
                     server.internal_ip))
    stdout, _ = client.RemoteCommand(load_command, should_log=True)
    tps, latency = ParseOutput(stdout)

    metadata = {
        'Average Transactions Per Second': tps,
        'Client Threads': threads,
        'Storage Type': FLAGS.aerospike_storage_type,
    }
    if FLAGS.aerospike_storage_type == DISK:
      metadata['Disk Type'] = 'Local' if FLAGS.use_local_disk else 'Remote'
    samples.append(('Average Latency', latency, 'ms', metadata))

  return samples

def Cleanup(benchmark_spec):
  """Cleanup Aerospike.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  server, client = benchmark_spec.vms

  client.RemoteCommand('sudo rm -rf aerospike*')
  server.RemoteCommand('cd %s && nohup sudo make stop' % SERVER_DIR)
  server.RemoteCommand('sudo rm -rf aerospike*')
<|endoftext|>


<commit_before>'''pylrparsing.Choice Class.

Choose one of the given elements. When most_greedy is True we will choose the 'longest'
element when multiple elements are valid. If most_greedy is False we will return the first
match.

:copyright: 2015, Jeroen van der Heijden (Transceptor Technology)
:license: need to decide
'''

from .noderesult import NodeResult
from .elements import NamedElement

class Choice(NamedElement):

    __slots__ = ('_elements', '_get_node_result')

    def __init__(self, *elements, most_greedy=True):
        self._elements = self._validate_elements(elements)
        self._get_node_result = self._most_greedy_result if most_greedy else self._stop_at_first_match

    def _most_greedy_result(self, root, tree, rule, s, node):
        most_greedy = NodeResult(False, node.start)

        for elem in self._elements:
            children = []
            node_res = root._walk(elem, node.start, children, rule, True)

            if node_res.is_valid and node_res.pos > most_greedy.pos:
                node.children = children
                most_greedy = node_res

        if most_greedy.is_valid:
            root._append_tree(tree, node, most_greedy.pos)

        return most_greedy

    def _stop_at_first_match(self, root, tree, rule, s, node):
        for elem in self._elements:
            children = []
            node_res = root._walk(elem, node.start, children, rule, True)

            if node_res.is_valid:
                node.children = children
                root._append_tree(tree, node, node_res.pos)
                break

        return node_res
<commit_msg>Rename module<commit_after>'''pyleri.Choice Class.

Choose one of the given elements. When most_greedy is True we will choose the 'longest'
element when multiple elements are valid. If most_greedy is False we will return the first
match.

:copyright: 2015, Jeroen van der Heijden (Transceptor Technology)
:license: need to decide
'''

from .noderesult import NodeResult
from .elements import NamedElement

class Choice(NamedElement):

    __slots__ = ('_elements', '_get_node_result')

    def __init__(self, *elements, most_greedy=True):
        self._elements = self._validate_elements(elements)
        self._get_node_result = self._most_greedy_result if most_greedy else self._stop_at_first_match

    def _most_greedy_result(self, root, tree, rule, s, node):
        most_greedy = NodeResult(False, node.start)

        for elem in self._elements:
            children = []
            node_res = root._walk(elem, node.start, children, rule, True)

            if node_res.is_valid and node_res.pos > most_greedy.pos:
                node.children = children
                most_greedy = node_res

        if most_greedy.is_valid:
            root._append_tree(tree, node, most_greedy.pos)

        return most_greedy

    def _stop_at_first_match(self, root, tree, rule, s, node):
        for elem in self._elements:
            children = []
            node_res = root._walk(elem, node.start, children, rule, True)

            if node_res.is_valid:
                node.children = children
                root._append_tree(tree, node, node_res.pos)
                break

        return node_res
<|endoftext|>