<commit_before>__author__ = 'richard'

import numpy as np
import pandas as pd

import scripts.plot_windtunnel as pwt
from scripts.i_o import get_directory
from scipy.interpolate import griddata

class Windtunnel():
    def __init__(self, experimental_condition):
        """
        experimental_condition

        """
        self.walls = Walls()
        self.boundary = self.walls.boundary
        self.experimental_condition = experimental_condition

        self.heater_l = Heater("Left", self.experimental_condition)
        self.heater_r = Heater("Right", self.experimental_condition)

    def show(self):
        ax = pwt.plot_windtunnel(self)
        return ax

class Walls():
    def __init__(self):
        # these are real dims of our wind tunnel
        self.left = -0.127
        self.right = 0.127
        self.upwind = 1.0
        self.downwind = 0.0
        self.ceiling = 0.254
        self.floor = 0.
        self.boundary = [self.downwind, self.upwind, self.left, self.right, self.floor, self.ceiling]

    def in_bounds(self, position):
        xpos, ypos, zpos = position
        inside = True
        past_wall = []

        if xpos > self.upwind:  # beyond upwind(upwind) wall (end)
            inside = False
            past_wall.append('upwind')
        if xpos < self.downwind:  # too far behind
            inside = False
            past_wall.append('downwind')
        if ypos < self.left:  # too left
            inside = False
            past_wall.append('left')
        if ypos > self.right:  # too far right
            inside = False
            past_wall.append('right')
        if zpos > self.ceiling:  # too far above
            inside = False
            past_wall.append('ceiling')
        if zpos < self.floor:  # too far below
            inside = False
            past_wall.append('floor')

        return inside, past_wall

class Heater():
    def __init__(self, side, experimental_condition):
        ''' given {left, right, none, custom coords} place heater in the windtunnel

        Args:
        location

        returns [x,y, zmin, zmax, diam]
        '''
        self.side = side
        self.experimental_condition = experimental_condition

        if side == experimental_condition:
            self.is_on = True
        else:
            self.is_on = False

        colors = {False: 'black', True: 'red'}
        self.color = colors[self.is_on]

        self.zmin = 0.03800
        self.zmax = 0.11340
        self.diam = 0.00635
        (self.x_position, self.y_position) = self._set_xy_coords()

    def _set_xy_coords(self):
        x_coord = 0.864

        if self.side in "leftLeftLEFT":
            y_coord = -0.0507
        elif self.side in "rightRightRIGHT":
            y_coord = 0.0507
        elif self.side in 'controlControlCONTROL':
            x_coord, y_coord = None, None
        else:
            raise Exception('invalid location type specified')

        return (x_coord, y_coord)

class Plume(object):
    ''' The Plume superclass

    Args:
    condition: {left|right|None}
    '''

    def __init__(self, experiment):
        # useful aliases
        self.experiment = experiment
        self.condition = experiment.condition
        self.walls = experiment.windtunnel.walls

class Boolean_Plume(Plume):
    """Are you in the plume Y/N"""
    def __init__(self, experiment):
        super(self.__class__, self).__init__(experiment)

        self.data = self._load_plume_data()

        try:
            self.resolution = abs(self.data.x_position.diff()[1])
        except AttributeError:  # if no plume, can't take diff() of no data
            self.resolution = None

    def in_plume(self, position):
        in_bounds, _ = self.walls.in_bounds(position)
        x, y, z = position

        if self.condition in 'controlControlCONTROL':
            inPlume = False
        elif np.abs(self.data['x_position'] - x).min() > self.resolution:
            # too far from the plume in the upwind/downwind direction
            inPlume = False
        elif in_bounds is False:
            print("WARNING: sniffing outside of windtunnel bounds")
            inPlume = False
        else:
            plume_plane = self._get_nearest_plume_plane(x)
            minor_axis = plume_plane.small_radius
            minor_ax_major_ax_ratio = 3
            major_axis = minor_axis * minor_ax_major_ax_ratio

            # implementation of http://math.stackexchange.com/a/76463/291217

            value = (((y - plume_plane.y_position) ** 2) / minor_axis ** 2) + \
                    (((z - plume_plane.z_position) ** 2) / major_axis ** 2)

            if value <= 1:
                inPlume = True
            else:
                inPlume = False

        return inPlume

    def show(self):
        fig, ax = pwt.plot_windtunnel(self.experiment.windtunnel)
        ax.axis('off')
        pwt.draw_plume(self, ax=ax)

    def _get_nearest_plume_plane(self, x_position):
        """given x position, find nearest plan"""
        closest_plume_index = (np.abs(self.data['x_position'] - x_position)).argmin()
        plume_plane = self.data.loc[closest_plume_index]

        return plume_plane

    def _load_plume_data(self):
        col_names = ['x_position', 'z_position', 'small_radius']

        if self.condition in 'controlControlCONTROL':
            return None
        elif self.condition in 'lLleftLeft':
            plume_dir = get_directory('BOOL_LEFT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
            df['y_position'] = self.experiment.windtunnel.heater_l.y_position
        elif self.condition in 'rightRight':
            plume_dir = get_directory('BOOL_RIGHT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
            df['y_position'] = self.experiment.windtunnel.heater_r.y_position
        else:
            raise Exception('problem with loading plume data {}'.format(self.condition))

        return df

class Timeavg_Plume(Plume):
    """time-averaged temperature readings taken inside the windtunnel"""
    def __init__(self, experiment):
        super(self.__class__, self).__init__(experiment)

        # initialize vals
        self.interpolated_data = pd.DataFrame()
        self._grid_x, self._grid_y, self.grid_z, self._interpolated_temps = None, None, None, None
        self._gradient_x, self._gradient_y, self._gradient_z = None, None, None

        # number of x, y, z positions to interpolate the data
        Nx, Ny, Nz = 100, 25, 25

        self.data = self._load_plume_data()
        self._interpolate_data((Nx, Ny, Nz))
        self._calc_gradient()

    def in_plume(self, position):
        pass

    def show_gradient(self):
        from mpl_toolkits.mplot3d import axes3d
        import matplotlib.pyplot as plt

        fig = plt.figure()
        ax = fig.gca(projection='3d')

        ax.quiver(self._grid_x, self._grid_y, self._grid_z, self._gradient_x, self._gradient_y, self._gradient_z, length=0.01)

        plt.show()

    def _load_plume_data(self):

        col_names = ['x', 'y', 'z', 'temperature']

        if self.condition in 'controlControlCONTROL':
            return None
        elif self.condition in 'lLleftLeft':
            plume_dir = get_directory('THERMOCOUPLE_TIMEAVG_LEFT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
        elif self.condition in 'rightRight':
            plume_dir = get_directory('THERMOCOUPLE_TIMEAVG_RIGHT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
        else:
            raise Exception('problem with loading plume data {}'.format(self.condition))

        return df

    def _interpolate_data(self, resolution):
        self._grid_x, self._grid_y, self._grid_z = np.mgrid[0.:1.:resolution[0], -0.127:0.127:resolution[1], 0:0.254:resolution[2]]
        print self._grid_x
        # grid_x, grid_y, grid_z = np.mgrid[0.:1.:100j, -0.127:0.127:25j, 0:0.254:25j]
        points = self.data[['x', 'y', 'z']].values  # (1382, 3)
        temps = self.data.temperature.values  # len 1382

        self._interpolated_temps = griddata(points,
                                      temps,
                                      (self._grid_x, self._grid_y, self._grid_z),
                                      method='nearest')
        print self._interpolated_temps.ravel().shape

        self.interpolated_data['x'] = self._grid_x.ravel()
        self.interpolated_data['y'] = self._grid_y.ravel()
        self.interpolated_data['z'] = self._grid_z.ravel()
        self.interpolated_data['avg_temp'] = self._interpolated_temps.ravel()

    def _calc_gradient(self):
        # Solve for the spatial gradient
        self._gradient_x, self._gradient_y, self._gradient_z = np.gradient(self._interpolated_temps,
                                                             self._grid_x,
                                                             self._grid_y,
                                                             self._grid_z)

        self.interpolated_data['gradient_x'] = self._gradient_x.ravel()
        self.interpolated_data['gradient_y'] = self._gradient_y.ravel()
        self.interpolated_data['gradient_z'] = self._gradient_z.ravel()

if __name__ == '__main__':
    windtunnel = Windtunnel('left')<commit_msg>gradient magnitude not quite working<commit_after>__author__ = 'richard'

import numpy as np
import pandas as pd

import scripts.plot_windtunnel as pwt
from scripts.i_o import get_directory
from scipy.interpolate import griddata

class Windtunnel():
    def __init__(self, experimental_condition):
        """
        experimental_condition

        """
        self.walls = Walls()
        self.boundary = self.walls.boundary
        self.experimental_condition = experimental_condition

        self.heater_l = Heater("Left", self.experimental_condition)
        self.heater_r = Heater("Right", self.experimental_condition)

    def show(self):
        ax = pwt.plot_windtunnel(self)
        return ax

class Walls():
    def __init__(self):
        # these are real dims of our wind tunnel
        self.left = -0.127
        self.right = 0.127
        self.upwind = 1.0
        self.downwind = 0.0
        self.ceiling = 0.254
        self.floor = 0.
        self.boundary = [self.downwind, self.upwind, self.left, self.right, self.floor, self.ceiling]

    def in_bounds(self, position):
        xpos, ypos, zpos = position
        inside = True
        past_wall = []

        if xpos > self.upwind:  # beyond upwind(upwind) wall (end)
            inside = False
            past_wall.append('upwind')
        if xpos < self.downwind:  # too far behind
            inside = False
            past_wall.append('downwind')
        if ypos < self.left:  # too left
            inside = False
            past_wall.append('left')
        if ypos > self.right:  # too far right
            inside = False
            past_wall.append('right')
        if zpos > self.ceiling:  # too far above
            inside = False
            past_wall.append('ceiling')
        if zpos < self.floor:  # too far below
            inside = False
            past_wall.append('floor')

        return inside, past_wall

class Heater():
    def __init__(self, side, experimental_condition):
        ''' given {left, right, none, custom coords} place heater in the windtunnel

        Args:
        location

        returns [x,y, zmin, zmax, diam]
        '''
        self.side = side
        self.experimental_condition = experimental_condition

        if side == experimental_condition:
            self.is_on = True
        else:
            self.is_on = False

        colors = {False: 'black', True: 'red'}
        self.color = colors[self.is_on]

        self.zmin = 0.03800
        self.zmax = 0.11340
        self.diam = 0.00635
        (self.x_position, self.y_position) = self._set_xy_coords()

    def _set_xy_coords(self):
        x_coord = 0.864

        if self.side in "leftLeftLEFT":
            y_coord = -0.0507
        elif self.side in "rightRightRIGHT":
            y_coord = 0.0507
        elif self.side in 'controlControlCONTROL':
            x_coord, y_coord = None, None
        else:
            raise Exception('invalid location type specified')

        return (x_coord, y_coord)

class Plume(object):
    ''' The Plume superclass

    Args:
    condition: {left|right|None}
    '''

    def __init__(self, experiment):
        # useful aliases
        self.experiment = experiment
        self.condition = experiment.condition
        self.walls = experiment.windtunnel.walls

class Boolean_Plume(Plume):
    """Are you in the plume Y/N"""
    def __init__(self, experiment):
        super(self.__class__, self).__init__(experiment)

        self.data = self._load_plume_data()

        try:
            self.resolution = abs(self.data.x_position.diff()[1])
        except AttributeError:  # if no plume, can't take diff() of no data
            self.resolution = None

    def in_plume(self, position):
        in_bounds, _ = self.walls.in_bounds(position)
        x, y, z = position

        if self.condition in 'controlControlCONTROL':
            inPlume = False
        elif np.abs(self.data['x_position'] - x).min() > self.resolution:
            # too far from the plume in the upwind/downwind direction
            inPlume = False
        elif in_bounds is False:
            print("WARNING: sniffing outside of windtunnel bounds")
            inPlume = False
        else:
            plume_plane = self._get_nearest_plume_plane(x)
            minor_axis = plume_plane.small_radius
            minor_ax_major_ax_ratio = 3
            major_axis = minor_axis * minor_ax_major_ax_ratio

            # implementation of http://math.stackexchange.com/a/76463/291217

            value = (((y - plume_plane.y_position) ** 2) / minor_axis ** 2) + \
                    (((z - plume_plane.z_position) ** 2) / major_axis ** 2)

            if value <= 1:
                inPlume = True
            else:
                inPlume = False

        return inPlume

    def show(self):
        fig, ax = pwt.plot_windtunnel(self.experiment.windtunnel)
        ax.axis('off')
        pwt.draw_plume(self, ax=ax)

    def _get_nearest_plume_plane(self, x_position):
        """given x position, find nearest plan"""
        closest_plume_index = (np.abs(self.data['x_position'] - x_position)).argmin()
        plume_plane = self.data.loc[closest_plume_index]

        return plume_plane

    def _load_plume_data(self):
        col_names = ['x_position', 'z_position', 'small_radius']

        if self.condition in 'controlControlCONTROL':
            return None
        elif self.condition in 'lLleftLeft':
            plume_dir = get_directory('BOOL_LEFT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
            df['y_position'] = self.experiment.windtunnel.heater_l.y_position
        elif self.condition in 'rightRight':
            plume_dir = get_directory('BOOL_RIGHT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
            df['y_position'] = self.experiment.windtunnel.heater_r.y_position
        else:
            raise Exception('problem with loading plume data {}'.format(self.condition))

        return df

class Timeavg_Plume(Plume):
    """time-averaged temperature readings taken inside the windtunnel"""
    def __init__(self, experiment):
        super(self.__class__, self).__init__(experiment)

        # initialize vals
        self.interpolated_data = pd.DataFrame()
        self._grid_x, self._grid_y, self.grid_z, self._interpolated_temps = None, None, None, None
        self._gradient_x, self._gradient_y, self._gradient_z = None, None, None

        # number of x, y, z positions to interpolate the data
        resolution = (100j, 25j, 25j) # stored as complex numbers for mgrid to work properly

        self.data = self._load_plume_data()
        self._interpolate_data(resolution)
        self._calc_gradient()

    def in_plume(self, position):
        pass

    def show_gradient(self):
        from mpl_toolkits.mplot3d import axes3d
        import matplotlib.pyplot as plt

        fig = plt.figure()
        ax = fig.gca(projection='3d')

        ax.quiver(self._grid_x, self._grid_y, self._grid_z, self._gradient_x, self._gradient_y, self._gradient_z, length=0.01)

        plt.show()

    def _load_plume_data(self):

        col_names = ['x', 'y', 'z', 'temperature']

        if self.condition in 'controlControlCONTROL':
            return None
        elif self.condition in 'lLleftLeft':
            plume_dir = get_directory('THERMOCOUPLE_TIMEAVG_LEFT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
        elif self.condition in 'rightRight':
            plume_dir = get_directory('THERMOCOUPLE_TIMEAVG_RIGHT_CSV')
            df = pd.read_csv(plume_dir, names=col_names)
        else:
            raise Exception('problem with loading plume data {}'.format(self.condition))

        return df

    def _interpolate_data(self, resolution):
        self._grid_x, self._grid_y, self._grid_z = np.mgrid[0.:1.:resolution[0], -0.127:0.127:resolution[1], 0:0.254:resolution[2]]
        print self._grid_x.shape
        # grid_x, grid_y, grid_z = np.mgrid[0.:1.:100j, -0.127:0.127:25j, 0:0.254:25j]
        points = self.data[['x', 'y', 'z']].values  # (1382, 3)
        temps = self.data.temperature.values  # len 1382

        self._interpolated_temps = griddata(points,
                                      temps,
                                      (self._grid_x, self._grid_y, self._grid_z),
                                      method='nearest')
        print self._interpolated_temps.ravel().shape

        self.interpolated_data['x'] = self._grid_x.ravel()
        self.interpolated_data['y'] = self._grid_y.ravel()
        self.interpolated_data['z'] = self._grid_z.ravel()
        self.interpolated_data['avg_temp'] = self._interpolated_temps.ravel()

    def _calc_gradient(self):
        # Solve for the spatial gradient
        self._gradient_x, self._gradient_y, self._gradient_z = np.gradient(self._interpolated_temps,
                                                             self._grid_x,
                                                             self._grid_y,
                                                             self._grid_z)

        self.interpolated_data['gradient_x'] = self._gradient_x.ravel()
        self.interpolated_data['gradient_y'] = self._gradient_y.ravel()
        self.interpolated_data['gradient_z'] = self._gradient_z.ravel()
        self.interpolated_data['gradient_mag'] = np.linalg.norm(self.interpolated_data[['gradient_x', 'gradient_y', 'gradient_z']])

if __name__ == '__main__':
    windtunnel = Windtunnel('left')<|endoftext|>


<commit_before>import logging
import re
import threading
import time
from asyncio import run_coroutine_threadsafe, get_event_loop_policy, \
                    new_event_loop, sleep

from aiohttp import ClientError
from discord import Client, Status, HTTPException, GatewayNotFound
from discord.utils import get
from websockets import InvalidHandshake, WebSocketProtocolError

from . import BaseBridge
from ..event import Event, Target

class DiscordBridge(BaseBridge):
    def __init__(self, config):
        BaseBridge.__init__(self, config)
        self.users = {}
        self.user_map = {}
        self.user_lock = threading.Lock()

        self.leaving_users = {}
        self.lv_thread = threading.Thread(target=self.leave_loop, daemon=True)
        self.lv_thread.start()

        self.loop = loop = new_event_loop()
        self.bridge_bot = DiscordBridgeBot(self.config, self, id(self), loop)

    def on_register(self):
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

        for name in self.config['channels']:
            self.send_event(self, Target.Manager, 'channel_join', name)

    def on_channel_add(self, channel):
        run_coroutine_threadsafe(self.bridge_bot.add_channel(channel),
                                 self.bridge_bot.loop).result()

    def on_channel_remove(self, channel):
        run_coroutine_threadsafe(self.bridge_bot.remove_channel(channel),
                                 self.bridge_bot.loop).result()

    def decode_mentions(self, content, channel_id):
        # TODO: Avoid using private member
        users = self.channels[channel_id]._users

        def replace(match):
            user_id = int(match.group(1))
            if user_id in self.users:
                return '<@{}>'.format(self.users[user_id].discord_id)
            elif user_id in users:
                return '@{}'.format(users[user_id].name)
            else:
                return match.group()

        return re.sub(r'<\[@([0-9]+)\]>', replace, content)

    def ev_message(self, event, content):
        if event.source_id in self.users:
            return

        name = self.get_user(event.source_id).name
        content = '<{}> {}'.format(name, content)

        if event.target_id in self.channels:
            content = self.decode_mentions(content, event.target_id)
            self.bridge_bot.message(self.channels[event.target_id].name,
                                    content)

    def ev_action(self, event, content):
        if event.source_id in self.users:
            return

        name = self.get_user(event.source_id).name
        content = '* {} {}'.format(name, content)

        if event.target_id in self.channels:
            content = self.decode_mentions(content, event.target_id)
            self.bridge_bot.action(self.channels[event.target_id].name,
                                   content)

    def ev_shutdown(self, event):
        if self.loop.is_running() and self.bridge_bot._is_ready.is_set():
            run_coroutine_threadsafe(self.bridge_bot.close(),
                                     self.loop).result()

        self.detach()

    def on_terminate(self):
        if self.loop.is_running() and self.bridge_bot._is_ready.is_set():
            run_coroutine_threadsafe(self.bridge_bot.close(),
                                     self.loop).result()

    def run(self):
        policy = get_event_loop_policy()
        policy.set_event_loop(self.loop)

        while True:
            try:
                task = self.bridge_bot.login(self.config['user'],
                                             self.config['password'])
                self.loop.run_until_complete(task)

            except (HTTPException, ClientError):
                logging.exception("Failed to log in to Discord")
                self.loop.run_until_complete(sleep(10))

            except BaseException as e:
                self.send_event(self, Target.Manager, 'exception', e)
                return

            else:
                break

        while not self.bridge_bot.is_closed:
            try:
                self.loop.run_until_complete(self.bridge_bot.sane_connect())

            except (HTTPException, ClientError, GatewayNotFound,
                    InvalidHandshake, WebSocketProtocolError):
                logging.exception("Lost connection with Discord")
                self.loop.run_until_complete(sleep(10))

            except BaseException as e:
                self.send_event(self, Target.Manager, 'exception', e)
                return

    def leave_loop(self):
        while True:
            time.sleep(15)
            with self.user_lock:
                self.check_user_timeouts()

    def check_user_timeouts(self):
        leaving = []
        for (discord_id, channel), timestamp in self.leaving_users.items():
            if timestamp+self.config['timeout'] < time.time():
                leaving.append((discord_id, channel))

        for discord_id, channel in leaving:
            del self.leaving_users[(discord_id, channel)]
            self.user_timeout(discord_id, channel)

    def user_timeout(self, discord_id, channel):
        user_id = self.user_map[discord_id]
        self.send_event(self, Target.Manager, 'user_leave',
                        channel.id, user_id)

        user_channels = self.users[user_id].channels
        user_channels.remove(channel)

        if not user_channels:
            del self.users[user_id]
            del self.user_map[discord_id]

    def discord_user_join(self, channel, discord_id, name):
        with self.user_lock:
            # Discard possible pending leave for the user
            self.leaving_users.pop((discord_id, channel), None)

            if discord_id not in self.user_map:
                user = DiscordUser(discord_id, {channel})
                self.send_event(self, Target.Manager, 'user_join',
                                channel.id, id(user), name)

                self.users[id(user)] = user
                self.user_map[discord_id] = id(user)

            else:
                user = self.users[self.user_map[discord_id]]
                if channel not in user.channels:
                    self.send_event(self, Target.Manager, 'user_join',
                                    channel.id, id(user), name)

                    user.channels.add(channel)

    def discord_name_change(self, channel, discord_id, new_name):
        with self.user_lock:
            if discord_id in self.user_map:
                self.send_event(self, Target.Manager, 'user_change',
                                channel.id, self.user_map[discord_id],
                                new_name)

    def discord_user_leave(self, channel, discord_id):
        with self.user_lock:
            if discord_id in self.user_map:
                if (discord_id, channel) not in self.leaving_users:
                    self.leaving_users[(discord_id, channel)] = time.time()

    def translate_mentions(self, content, mentions):
        def replace(match):
            discord_id = match.group(1)
            if discord_id in self.user_map:
                return '<[@{}]>'.format(self.user_map[discord_id])
            else:
                user = get(mentions, id=discord_id)
                if user:
                    return '@{}'.format(user.name)
                else:
                    return match.group()

        return re.sub(r'<@([0-9]+)>', replace, content)

    def discord_channel_message(self, channel, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                user_id = self.user_map[discord_id]
                content = self.translate_mentions(content, mentions)
                self.send_event(user_id, channel.id, 'message', content)

    def discord_channel_action(self, channel, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                user_id = self.user_map[discord_id]
                content = self.translate_mentions(content[1:-1], mentions)
                self.send_event(user_id, channel.id, 'action', content)

    def discord_private_message(self, user_id, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                content = self.translate_mentions(content, mentions)
                self.send_event(self.user_map[discord_id], user_id, 'message',
                                content)

    def discord_private_action(self, user_id, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                content = self.translate_mentions(content[1:-1], mentions)
                self.send_event(self.user_map[discord_id], user_id, 'message',
                                content)

class DiscordUser:
    def __init__(self, discord_id, channels):
        self._discord_id = discord_id
        self.channels = channels

    def __eq__(self, other):
        if isinstance(other, DiscordUser):
            return self._discord_id == other._discord_id
        elif isinstance(other, str):
            return self._discord_id == other
        else:
            return NotImplemented

    def __str__(self):
        return 'DiscordUser({!r})'.format(self._discord_id)

    @property
    def discord_id(self):
        return self._discord_id

class DiscordBot(Client):
    def __init__(self, config, bridge, user_id, loop):
        Client.__init__(self, loop=loop)

        self.config = config
        self.bridge = bridge
        self.user_id = user_id
        self.joined_channels = {}

    async def sane_connect(self):
        self.gateway = await self._get_gateway()
        await self._make_websocket()

        while not self.is_closed:
            msg = await self.ws.recv()
            if msg is None:
                if self.ws.close_code == 1012:
                    await self.redirect_websocket(self.gateway)
                    continue
                else:
                    # Connection was dropped, break out
                    break

            await self.received_message(msg)

    def action(self, target_id, content):
        target_id = self.config['channels'][target_id]
        self.loop.call_soon_threadsafe(self.do_msg, target_id, content)

    def message(self, target_id, content):
        target_id = self.config['channels'][target_id]
        self.loop.call_soon_threadsafe(self.do_msg, target_id, content)

    def do_msg(self, target_id, content):
        channel = self.get_channel(target_id)
        if channel is not None:
            task = self.loop.create_task(self.send_message(channel, content))
            task.add_done_callback(lambda task: self.msg_done(task, content))

        else:
            print("unkown target", target_id)

    def msg_done(self, task, content):
        try:
            task.result()
        except (HTTPException, ClientError):
            logging.exception('Error sending "{}"'.format(content))
        except BaseException as e:
            self.bridge.send_event(self, Target.Manager, 'exception', e)

    @staticmethod
    def is_action(message):
        content = message.content
        return (len(content) > 2 and content[0] == content[-1] == '_'
                and content[-2] != '_' and content[1] != '_')

    async def on_message(self, message):
        if message.channel.is_private:
            args = (self.user_id, message.author.id,
                    message.content, message.mentions)
            if self.is_action(message):
                self.bridge.discord_private_action(*args)
            else:
                self.bridge.discord_private_message(*args)

    async def add_channel(self, channel):
        channel_id = self.config['channels'][channel.name]
        self.joined_channels[channel_id] = channel

    async def remove_channel(self, channel):
        channel_id = self.config['channels'][channel.name]
        del self.joined_channels[channel_id]

class DiscordBridgeBot(DiscordBot):
    def _sync_member(self, channel, member, discord_channel):
        if member != self.user:
            if (discord_channel.permissions_for(member).read_messages
                    and member.status != Status.offline):
                self.bridge.discord_user_join(channel, member.id, member.name)
            else:
                self.bridge.discord_user_leave(channel, member.id)

    def _sync_channel_members(self, channel, discord_channel):
        for member in discord_channel.server.members:
            self._sync_member(channel, member, discord_channel)

    async def on_member_join(self, member):
        for channel_id, channel in self.joined_channels.items():
            discord_channel = self.get_channel(channel_id)
            if discord_channel is not None:
                self._sync_member(channel, member, discord_channel)

    async def on_member_remove(self, member):
        for channel_id, channel in self.joined_channels.items():
            self.bridge.discord_user_leave(channel, member.id)

    async def on_member_update(self, before, after):
        if after.name != before.name:
            for channel in self.joined_channels.values():
                self.bridge.discord_name_change(channel, after.id, after.name)

        for discord_channel in after.server.channels:
            if discord_channel.id in self.joined_channels:
                channel = self.joined_channels[discord_channel.id]
                self._sync_member(channel, after, discord_channel)

    async def on_ready(self):
        for channel_id, channel in self.joined_channels.items():
            discord_channel = self.get_channel(channel_id)
            if discord_channel is not None:
                self._sync_channel_members(channel, discord_channel)

    async def on_message(self, message):
        await DiscordBot.on_message(self, message)

        if message.channel.id in self.joined_channels:
            channel = self.joined_channels[message.channel.id]
            args = (channel, message.author.id)
            content, mentions = message.content, message.mentions
            if message.attachments:
                urls = ', '.join((a['url'] for a in message.attachments))
                self.bridge.discord_channel_message(*args, urls, mentions)
            elif self.is_action(message):
                self.bridge.discord_channel_action(*args, content, mentions)
            else:
                self.bridge.discord_channel_message(*args, content, mentions)

    async def add_channel(self, channel):
        await DiscordBot.add_channel(self, channel)
        discord_channel_id = self.config['channels'][channel.name]

        discord_channel = self.get_channel(discord_channel_id)
        if discord_channel is not None:
            self._sync_channel_members(channel, discord_channel)

    async def remove_channel(self, channel):
        await DiscordBot.remove_channel(self, channel)
        discord_channel_id = self.config['channels'][channel.name]

        discord_channel = self.get_channel(discord_channel_id)
        if discord_channel is not None:
            for member in discord_channel.server.members:
                self.bridge.discord_user_leave(channel, member.id)
<commit_msg>Handle Discord attachments with captions<commit_after>import logging
import re
import threading
import time
from asyncio import run_coroutine_threadsafe, get_event_loop_policy, \
                    new_event_loop, sleep

from aiohttp import ClientError
from discord import Client, Status, HTTPException, GatewayNotFound
from discord.utils import get
from websockets import InvalidHandshake, WebSocketProtocolError

from . import BaseBridge
from ..event import Event, Target

class DiscordBridge(BaseBridge):
    def __init__(self, config):
        BaseBridge.__init__(self, config)
        self.users = {}
        self.user_map = {}
        self.user_lock = threading.Lock()

        self.leaving_users = {}
        self.lv_thread = threading.Thread(target=self.leave_loop, daemon=True)
        self.lv_thread.start()

        self.loop = loop = new_event_loop()
        self.bridge_bot = DiscordBridgeBot(self.config, self, id(self), loop)

    def on_register(self):
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

        for name in self.config['channels']:
            self.send_event(self, Target.Manager, 'channel_join', name)

    def on_channel_add(self, channel):
        run_coroutine_threadsafe(self.bridge_bot.add_channel(channel),
                                 self.bridge_bot.loop).result()

    def on_channel_remove(self, channel):
        run_coroutine_threadsafe(self.bridge_bot.remove_channel(channel),
                                 self.bridge_bot.loop).result()

    def decode_mentions(self, content, channel_id):
        # TODO: Avoid using private member
        users = self.channels[channel_id]._users

        def replace(match):
            user_id = int(match.group(1))
            if user_id in self.users:
                return '<@{}>'.format(self.users[user_id].discord_id)
            elif user_id in users:
                return '@{}'.format(users[user_id].name)
            else:
                return match.group()

        return re.sub(r'<\[@([0-9]+)\]>', replace, content)

    def ev_message(self, event, content):
        if event.source_id in self.users:
            return

        name = self.get_user(event.source_id).name
        content = '<{}> {}'.format(name, content)

        if event.target_id in self.channels:
            content = self.decode_mentions(content, event.target_id)
            self.bridge_bot.message(self.channels[event.target_id].name,
                                    content)

    def ev_action(self, event, content):
        if event.source_id in self.users:
            return

        name = self.get_user(event.source_id).name
        content = '* {} {}'.format(name, content)

        if event.target_id in self.channels:
            content = self.decode_mentions(content, event.target_id)
            self.bridge_bot.action(self.channels[event.target_id].name,
                                   content)

    def ev_shutdown(self, event):
        if self.loop.is_running() and self.bridge_bot._is_ready.is_set():
            run_coroutine_threadsafe(self.bridge_bot.close(),
                                     self.loop).result()

        self.detach()

    def on_terminate(self):
        if self.loop.is_running() and self.bridge_bot._is_ready.is_set():
            run_coroutine_threadsafe(self.bridge_bot.close(),
                                     self.loop).result()

    def run(self):
        policy = get_event_loop_policy()
        policy.set_event_loop(self.loop)

        while True:
            try:
                task = self.bridge_bot.login(self.config['user'],
                                             self.config['password'])
                self.loop.run_until_complete(task)

            except (HTTPException, ClientError):
                logging.exception("Failed to log in to Discord")
                self.loop.run_until_complete(sleep(10))

            except BaseException as e:
                self.send_event(self, Target.Manager, 'exception', e)
                return

            else:
                break

        while not self.bridge_bot.is_closed:
            try:
                self.loop.run_until_complete(self.bridge_bot.sane_connect())

            except (HTTPException, ClientError, GatewayNotFound,
                    InvalidHandshake, WebSocketProtocolError):
                logging.exception("Lost connection with Discord")
                self.loop.run_until_complete(sleep(10))

            except BaseException as e:
                self.send_event(self, Target.Manager, 'exception', e)
                return

    def leave_loop(self):
        while True:
            time.sleep(15)
            with self.user_lock:
                self.check_user_timeouts()

    def check_user_timeouts(self):
        leaving = []
        for (discord_id, channel), timestamp in self.leaving_users.items():
            if timestamp+self.config['timeout'] < time.time():
                leaving.append((discord_id, channel))

        for discord_id, channel in leaving:
            del self.leaving_users[(discord_id, channel)]
            self.user_timeout(discord_id, channel)

    def user_timeout(self, discord_id, channel):
        user_id = self.user_map[discord_id]
        self.send_event(self, Target.Manager, 'user_leave',
                        channel.id, user_id)

        user_channels = self.users[user_id].channels
        user_channels.remove(channel)

        if not user_channels:
            del self.users[user_id]
            del self.user_map[discord_id]

    def discord_user_join(self, channel, discord_id, name):
        with self.user_lock:
            # Discard possible pending leave for the user
            self.leaving_users.pop((discord_id, channel), None)

            if discord_id not in self.user_map:
                user = DiscordUser(discord_id, {channel})
                self.send_event(self, Target.Manager, 'user_join',
                                channel.id, id(user), name)

                self.users[id(user)] = user
                self.user_map[discord_id] = id(user)

            else:
                user = self.users[self.user_map[discord_id]]
                if channel not in user.channels:
                    self.send_event(self, Target.Manager, 'user_join',
                                    channel.id, id(user), name)

                    user.channels.add(channel)

    def discord_name_change(self, channel, discord_id, new_name):
        with self.user_lock:
            if discord_id in self.user_map:
                self.send_event(self, Target.Manager, 'user_change',
                                channel.id, self.user_map[discord_id],
                                new_name)

    def discord_user_leave(self, channel, discord_id):
        with self.user_lock:
            if discord_id in self.user_map:
                if (discord_id, channel) not in self.leaving_users:
                    self.leaving_users[(discord_id, channel)] = time.time()

    def translate_mentions(self, content, mentions):
        def replace(match):
            discord_id = match.group(1)
            if discord_id in self.user_map:
                return '<[@{}]>'.format(self.user_map[discord_id])
            else:
                user = get(mentions, id=discord_id)
                if user:
                    return '@{}'.format(user.name)
                else:
                    return match.group()

        return re.sub(r'<@([0-9]+)>', replace, content)

    def discord_channel_message(self, channel, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                user_id = self.user_map[discord_id]
                content = self.translate_mentions(content, mentions)
                self.send_event(user_id, channel.id, 'message', content)

    def discord_channel_action(self, channel, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                user_id = self.user_map[discord_id]
                content = self.translate_mentions(content[1:-1], mentions)
                self.send_event(user_id, channel.id, 'action', content)

    def discord_private_message(self, user_id, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                content = self.translate_mentions(content, mentions)
                self.send_event(self.user_map[discord_id], user_id, 'message',
                                content)

    def discord_private_action(self, user_id, discord_id, content, mentions):
        with self.user_lock:
            if discord_id in self.user_map:
                content = self.translate_mentions(content[1:-1], mentions)
                self.send_event(self.user_map[discord_id], user_id, 'message',
                                content)

class DiscordUser:
    def __init__(self, discord_id, channels):
        self._discord_id = discord_id
        self.channels = channels

    def __eq__(self, other):
        if isinstance(other, DiscordUser):
            return self._discord_id == other._discord_id
        elif isinstance(other, str):
            return self._discord_id == other
        else:
            return NotImplemented

    def __str__(self):
        return 'DiscordUser({!r})'.format(self._discord_id)

    @property
    def discord_id(self):
        return self._discord_id

class DiscordBot(Client):
    def __init__(self, config, bridge, user_id, loop):
        Client.__init__(self, loop=loop)

        self.config = config
        self.bridge = bridge
        self.user_id = user_id
        self.joined_channels = {}

    async def sane_connect(self):
        self.gateway = await self._get_gateway()
        await self._make_websocket()

        while not self.is_closed:
            msg = await self.ws.recv()
            if msg is None:
                if self.ws.close_code == 1012:
                    await self.redirect_websocket(self.gateway)
                    continue
                else:
                    # Connection was dropped, break out
                    break

            await self.received_message(msg)

    def action(self, target_id, content):
        target_id = self.config['channels'][target_id]
        self.loop.call_soon_threadsafe(self.do_msg, target_id, content)

    def message(self, target_id, content):
        target_id = self.config['channels'][target_id]
        self.loop.call_soon_threadsafe(self.do_msg, target_id, content)

    def do_msg(self, target_id, content):
        channel = self.get_channel(target_id)
        if channel is not None:
            task = self.loop.create_task(self.send_message(channel, content))
            task.add_done_callback(lambda task: self.msg_done(task, content))

        else:
            print("unkown target", target_id)

    def msg_done(self, task, content):
        try:
            task.result()
        except (HTTPException, ClientError):
            logging.exception('Error sending "{}"'.format(content))
        except BaseException as e:
            self.bridge.send_event(self, Target.Manager, 'exception', e)

    @staticmethod
    def is_action(message):
        content = message.content
        return (len(content) > 2 and content[0] == content[-1] == '_'
                and content[-2] != '_' and content[1] != '_')

    async def on_message(self, message):
        if message.channel.is_private:
            args = (self.user_id, message.author.id,
                    message.content, message.mentions)
            if self.is_action(message):
                self.bridge.discord_private_action(*args)
            else:
                self.bridge.discord_private_message(*args)

    async def add_channel(self, channel):
        channel_id = self.config['channels'][channel.name]
        self.joined_channels[channel_id] = channel

    async def remove_channel(self, channel):
        channel_id = self.config['channels'][channel.name]
        del self.joined_channels[channel_id]

class DiscordBridgeBot(DiscordBot):
    def _sync_member(self, channel, member, discord_channel):
        if member != self.user:
            if (discord_channel.permissions_for(member).read_messages
                    and member.status != Status.offline):
                self.bridge.discord_user_join(channel, member.id, member.name)
            else:
                self.bridge.discord_user_leave(channel, member.id)

    def _sync_channel_members(self, channel, discord_channel):
        for member in discord_channel.server.members:
            self._sync_member(channel, member, discord_channel)

    async def on_member_join(self, member):
        for channel_id, channel in self.joined_channels.items():
            discord_channel = self.get_channel(channel_id)
            if discord_channel is not None:
                self._sync_member(channel, member, discord_channel)

    async def on_member_remove(self, member):
        for channel_id, channel in self.joined_channels.items():
            self.bridge.discord_user_leave(channel, member.id)

    async def on_member_update(self, before, after):
        if after.name != before.name:
            for channel in self.joined_channels.values():
                self.bridge.discord_name_change(channel, after.id, after.name)

        for discord_channel in after.server.channels:
            if discord_channel.id in self.joined_channels:
                channel = self.joined_channels[discord_channel.id]
                self._sync_member(channel, after, discord_channel)

    async def on_ready(self):
        for channel_id, channel in self.joined_channels.items():
            discord_channel = self.get_channel(channel_id)
            if discord_channel is not None:
                self._sync_channel_members(channel, discord_channel)

    async def on_message(self, message):
        await DiscordBot.on_message(self, message)

        if message.channel.id in self.joined_channels:
            channel = self.joined_channels[message.channel.id]
            args = (channel, message.author.id)
            content, mentions = message.content, message.mentions

            if self.is_action(message):
                self.bridge.discord_channel_action(*args, content, mentions)
            elif message.content:
                self.bridge.discord_channel_message(*args, content, mentions)

            if message.attachments:
                urls = ', '.join((a['url'] for a in message.attachments))
                self.bridge.discord_channel_message(*args, urls, mentions)

    async def add_channel(self, channel):
        await DiscordBot.add_channel(self, channel)
        discord_channel_id = self.config['channels'][channel.name]

        discord_channel = self.get_channel(discord_channel_id)
        if discord_channel is not None:
            self._sync_channel_members(channel, discord_channel)

    async def remove_channel(self, channel):
        await DiscordBot.remove_channel(self, channel)
        discord_channel_id = self.config['channels'][channel.name]

        discord_channel = self.get_channel(discord_channel_id)
        if discord_channel is not None:
            for member in discord_channel.server.members:
                self.bridge.discord_user_leave(channel, member.id)
<|endoftext|>


<commit_before># This file is part of beets.
# Copyright 2013, Thomas Scholtes
#
# 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.

import os
import tempfile
import shutil
from glob import glob

import _common
from _common import unittest

import beets
from beets import ui
from beets.library import Item, Album
from beets.mediafile import MediaFile

class ReplayGainGstCliTest(unittest.TestCase):

    def setUp(self):
        self.setupBeets()
        self.config['replaygain']['backend'] = u'gstreamer'
        self.config['plugins'] = ['replaygain']
        self.setupLibrary(2)

    def tearDown(self):
        del os.environ['BEETSDIR']
        shutil.rmtree(self.temp_dir)

    def setupBeets(self):
        self.temp_dir = tempfile.mkdtemp()
        os.environ['BEETSDIR'] = self.temp_dir

        self.config = beets.config
        self.config.clear()
        self.config.read(user=False)

        self.config['verbose'] = True
        self.config['color'] = False
        self.config['threaded'] = False
        self.config['import']['copy'] = False

        self.libdir = os.path.join(self.temp_dir, 'libdir')
        os.mkdir(self.libdir)
        self.config['directory'] = self.libdir

        self.libpath = os.path.join(self.temp_dir, 'lib')
        self.config['library'] = self.libpath

        self.lib = beets.library.Library(self.config['library'].as_filename(),
                                         self.libdir)

    def setupLibrary(self, file_count):
        """Add an album to the library with ``file_count`` items.
        """
        album = Album(id=1)
        album.add(self.lib)

        fixture_glob = os.path.join(_common.RSRC, '*.mp3')
        for src in glob(fixture_glob)[0:file_count]:
            dst = os.path.join(self.libdir, os.path.basename(src))
            shutil.copy(src, dst)
            item = Item.from_path(dst)
            item.album_id = 1
            item.add(self.lib)
            self._reset_replaygain(item)

    def _reset_replaygain(self, item):
        item['rg_track_peak'] = 0
        item['rg_track_gain'] = 0
        item['rg_album_gain'] = 0
        item['rg_album_gain'] = 0
        item.write()
        item.store()

    def test_cli_saves_track_gain(self):
        for item in self.lib.items():
            self.assertEqual(item.rg_track_peak, 0.0)
            self.assertEqual(item.rg_track_gain, 0.0)
            mediafile = MediaFile(item.path)
            self.assertEqual(mediafile.rg_track_peak, 0.0)
            self.assertEqual(mediafile.rg_track_gain, 0.0)

        ui._raw_main(['replaygain'])
        for item in self.lib.items():
            self.assertNotEqual(item.rg_track_peak, 0.0)
            self.assertNotEqual(item.rg_track_gain, 0.0)
            mediafile = MediaFile(item.path)
            self.assertAlmostEqual(
                mediafile.rg_track_peak, item.rg_track_peak, places=6)
            self.assertAlmostEqual(
                mediafile.rg_track_gain, item.rg_track_gain, places=6)

    def test_cli_saves_album_gain_to_file(self):
        for item in self.lib.items():
            mediafile = MediaFile(item.path)
            self.assertEqual(mediafile.rg_album_peak, 0.0)
            self.assertEqual(mediafile.rg_album_gain, 0.0)

        ui._raw_main(['replaygain', '-a'])

        peaks = []
        gains = []
        for item in self.lib.items():
            mediafile = MediaFile(item.path)
            peaks.append(mediafile.rg_album_peak)
            gains.append(mediafile.rg_album_gain)

        # Make sure they are all the same
        self.assertEqual(max(peaks), min(peaks))
        self.assertEqual(max(gains), min(gains))

        self.assertNotEqual(max(gains), 0.0)
        self.assertNotEqual(max(peaks), 0.0)

def suite():
    return unittest.TestLoader().loadTestsFromName(__name__)

if __name__ == '__main__':
    unittest.main(defaultTest='suite')
<commit_msg>Reset config and plugins after tests<commit_after># This file is part of beets.
# Copyright 2013, Thomas Scholtes
#
# 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.

import os
import tempfile
import shutil
from glob import glob

import _common
from _common import unittest

import beets
from beets import ui
from beets import plugins
from beets.library import Item, Album
from beets.mediafile import MediaFile

class ReplayGainGstCliTest(unittest.TestCase):

    def setUp(self):
        self.setupBeets()
        self.config['replaygain']['backend'] = u'gstreamer'
        self.config['plugins'] = ['replaygain']
        self.setupLibrary(2)

    def tearDown(self):
        del os.environ['BEETSDIR']
        shutil.rmtree(self.temp_dir)
        self.config.clear()
        self.config.read(user=False)
        plugins._classes = set()
        plugins._instances = {}

    def setupBeets(self):
        self.temp_dir = tempfile.mkdtemp()
        os.environ['BEETSDIR'] = self.temp_dir

        self.config = beets.config
        self.config.clear()
        self.config.read(user=False)

        self.config['verbose'] = True
        self.config['color'] = False
        self.config['threaded'] = False
        self.config['import']['copy'] = False

        self.libdir = os.path.join(self.temp_dir, 'libdir')
        os.mkdir(self.libdir)
        self.config['directory'] = self.libdir

        self.libpath = os.path.join(self.temp_dir, 'lib')
        self.config['library'] = self.libpath

        self.lib = beets.library.Library(self.config['library'].as_filename(),
                                         self.libdir)

    def setupLibrary(self, file_count):
        """Add an album to the library with ``file_count`` items.
        """
        album = Album(id=1)
        album.add(self.lib)

        fixture_glob = os.path.join(_common.RSRC, '*.mp3')
        for src in glob(fixture_glob)[0:file_count]:
            dst = os.path.join(self.libdir, os.path.basename(src))
            shutil.copy(src, dst)
            item = Item.from_path(dst)
            item.album_id = 1
            item.add(self.lib)
            self._reset_replaygain(item)

    def _reset_replaygain(self, item):
        item['rg_track_peak'] = 0
        item['rg_track_gain'] = 0
        item['rg_album_gain'] = 0
        item['rg_album_gain'] = 0
        item.write()
        item.store()

    def test_cli_saves_track_gain(self):
        for item in self.lib.items():
            self.assertEqual(item.rg_track_peak, 0.0)
            self.assertEqual(item.rg_track_gain, 0.0)
            mediafile = MediaFile(item.path)
            self.assertEqual(mediafile.rg_track_peak, 0.0)
            self.assertEqual(mediafile.rg_track_gain, 0.0)

        ui._raw_main(['replaygain'])
        for item in self.lib.items():
            self.assertNotEqual(item.rg_track_peak, 0.0)
            self.assertNotEqual(item.rg_track_gain, 0.0)
            mediafile = MediaFile(item.path)
            self.assertAlmostEqual(
                mediafile.rg_track_peak, item.rg_track_peak, places=6)
            self.assertAlmostEqual(
                mediafile.rg_track_gain, item.rg_track_gain, places=6)

    def test_cli_saves_album_gain_to_file(self):
        for item in self.lib.items():
            mediafile = MediaFile(item.path)
            self.assertEqual(mediafile.rg_album_peak, 0.0)
            self.assertEqual(mediafile.rg_album_gain, 0.0)

        ui._raw_main(['replaygain', '-a'])

        peaks = []
        gains = []
        for item in self.lib.items():
            mediafile = MediaFile(item.path)
            peaks.append(mediafile.rg_album_peak)
            gains.append(mediafile.rg_album_gain)

        # Make sure they are all the same
        self.assertEqual(max(peaks), min(peaks))
        self.assertEqual(max(gains), min(gains))

        self.assertNotEqual(max(gains), 0.0)
        self.assertNotEqual(max(peaks), 0.0)

def suite():
    return unittest.TestLoader().loadTestsFromName(__name__)

if __name__ == '__main__':
    unittest.main(defaultTest='suite')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import division

import os
import random
from random import randint
import socket
import threading
import time

from shinken.objects.module import Module
from shinken_modules import TestConfig
from shinken_test import time_hacker, unittest
from test_livestatus import LiveStatusTest

from livestatus.livestatus_wait_query import LiveStatusWaitQuery
from livestatus.livestatus_query import LiveStatusQuery

class Test_WaitQuery(LiveStatusTest):

    def test_wait_query(self):
        request = b'''GET hosts
WaitObject: gstarck_test
WaitCondition: last_check >= 1419007690
WaitTimeout: 10000
WaitTrigger: check
Columns: last_check state plugin_output
Filter: host_name = gstarck_test
Localtime: 1419007690
OutputFormat: python
KeepAlive: on
ResponseHeader: fixed16
ColumnHeaders: off

'''

        response, keepalive = self.livestatus_broker.livestatus.handle_request(request)
        self.assertEqual(2, len(response), 'should contains Wait + Query')
        self.assertIsInstance(response[0], LiveStatusWaitQuery)
        self.assertIsInstance(response[1], LiveStatusQuery)

class TestFull_WaitQuery(TestConfig):
    ''' "Full" test : that is with connection to livestatus socket
    And so we have to start a thread to execute the actual main livestatus function (manage_lql_thread),
    which will handle new connections as it would do in real life.
    '''

    def tearDown(self):
        # stop thread
        self.livestatus_broker.interrupted = True
        self.lql_thread.join()
        super(TestFull_WaitQuery, self).tearDown()

    def setUp(self):
        super(TestFull_WaitQuery, self).setUp()
        time_hacker.set_real_time()
        self.testid = str(os.getpid() + random.randint(1, 1000))
        self.modconf = Module({'module_name': 'LiveStatus',
            'module_type': 'livestatus',
            'port': str(random.randint(50000, 65534)),
            'pnp_path': 'tmp/pnp4nagios_test' + self.testid,
            'host': '127.0.0.1',
            'name': 'test',
            'modules': ''
        })
        self.init_livestatus(self.modconf)

    def init_livestatus(self, conf):
        super(TestFull_WaitQuery, self).init_livestatus(conf)
        self.sched.brokers['Default-Broker'] = {'broks' : {}, 'has_full_broks' : False}
        self.sched.fill_initial_broks('Default-Broker')
        self.update_broker()
        # execute the livestatus by starting a dedicated thread to run the manage_lql_thread function:
        self.lql_thread = threading.Thread(target=self.livestatus_broker.manage_lql_thread, name='lqlthread')
        self.lql_thread.start()
        # wait for thread to init
        time.sleep(3)

    def query_livestatus(self, data, ip=None, port=None, timeout=60):
        if ip is None:
            ip = self.modconf.host
        if port is None:
            port = self.modconf.port
        port = int(port)
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            s.settimeout(timeout)
            s.connect((ip, port))
            s.send(data)
            s.shutdown(socket.SHUT_WR)
            ret = []
            while True:
                data = s.recv(1024)
                #print("received: %r" % data)
                if not data:
                    return b''.join(ret)
                ret.append(data)
        finally:
            s.close()

    def test_wait_query_1(self):

        ## NB NB NB:
        # if there was a scheduler and a poller running and connected to `self.livestatus_broker.from_q´
        # then the prepended command would be processed and the livestatus **could** return a different response,
        # **depending** on :
        # the wait_timeout_sec we use (which then should be greater than the randint(1,3) actually used)
        # and on how fast the command would be processed by the poller
        # and on how fast the host status would be updated after while within the livestatus process/thread.

        wait_timeout_sec = randint(1, 3)
        now = int( time.time() ) # current livestatus wants an INTEGER value for WaitTimeout ..

        # following request is nearly exactly what check_mk send (modulo the timestamps/hostname)
        # when you reschedule an immediate host check :
        request = b'''COMMAND [{now}] SCHEDULE_FORCED_HOST_CHECK;{host};{now}
GET hosts
WaitObject: {host}
WaitCondition: last_check >= {last_check}
WaitTimeout: {wait_timeout}
WaitTrigger: check
Columns: last_check state plugin_output
Filter: host_name = {host}
Localtime: {localtime}
OutputFormat: python
KeepAlive: on
ResponseHeader: fixed16
ColumnHeaders: true

'''.format(last_check=now, localtime=now, host='test_host_0', now=now,
           wait_timeout=wait_timeout_sec * 1000, # WaitTimeout header field is in millisecs
        )

        t0 = time.time()
        response = self.query_livestatus(request)
        t1 = time.time()
        self.assertLess(wait_timeout_sec, t1 - t0,
                        '(actually and in this very specific case) livestatus should take at least the requested WaitTimeout (%s sec) to complete ; response=%s' %
                        (wait_timeout_sec, response))
        goodresponse = "200          13\n[[0, 0, '']]\n"
        self.assertEqual(goodresponse, response)

if __name__ == '__main__':
    unittest.main()
<commit_msg>WaitQuery test: Also make sure the command has been queued into livestatus `from_q´.<commit_after># -*- coding: utf-8 -*-

from __future__ import division

import os
import random
from random import randint
import socket
import threading
import time

from shinken.objects.module import Module
from shinken.external_command import ExternalCommand

from shinken_modules import TestConfig
from shinken_test import time_hacker, unittest
from test_livestatus import LiveStatusTest

from livestatus.livestatus_wait_query import LiveStatusWaitQuery
from livestatus.livestatus_query import LiveStatusQuery

class Test_WaitQuery(LiveStatusTest):

    def test_wait_query(self):
        request = b'''GET hosts
WaitObject: gstarck_test
WaitCondition: last_check >= 1419007690
WaitTimeout: 10000
WaitTrigger: check
Columns: last_check state plugin_output
Filter: host_name = gstarck_test
Localtime: 1419007690
OutputFormat: python
KeepAlive: on
ResponseHeader: fixed16
ColumnHeaders: off

'''

        response, keepalive = self.livestatus_broker.livestatus.handle_request(request)
        self.assertEqual(2, len(response), 'should contains Wait + Query')
        self.assertIsInstance(response[0], LiveStatusWaitQuery)
        self.assertIsInstance(response[1], LiveStatusQuery)

class TestFull_WaitQuery(TestConfig):
    ''' "Full" test : that is with connection to livestatus socket
    And so we have to start a thread to execute the actual main livestatus function (manage_lql_thread),
    which will handle new connections as it would do in real life.
    '''

    def tearDown(self):
        # stop thread
        self.livestatus_broker.interrupted = True
        self.lql_thread.join()
        super(TestFull_WaitQuery, self).tearDown()

    def setUp(self):
        super(TestFull_WaitQuery, self).setUp()
        time_hacker.set_real_time()
        self.testid = str(os.getpid() + random.randint(1, 1000))
        self.modconf = Module({'module_name': 'LiveStatus',
            'module_type': 'livestatus',
            'port': str(random.randint(50000, 65534)),
            'pnp_path': 'tmp/pnp4nagios_test' + self.testid,
            'host': '127.0.0.1',
            'name': 'test',
            'modules': ''
        })
        self.init_livestatus(self.modconf)

    def init_livestatus(self, conf):
        super(TestFull_WaitQuery, self).init_livestatus(conf)
        self.sched.brokers['Default-Broker'] = {'broks' : {}, 'has_full_broks' : False}
        self.sched.fill_initial_broks('Default-Broker')
        self.update_broker()
        # execute the livestatus by starting a dedicated thread to run the manage_lql_thread function:
        self.lql_thread = threading.Thread(target=self.livestatus_broker.manage_lql_thread, name='lqlthread')
        self.lql_thread.start()
        # wait for thread to init
        time.sleep(3)

    def query_livestatus(self, data, ip=None, port=None, timeout=60):
        if ip is None:
            ip = self.modconf.host
        if port is None:
            port = self.modconf.port
        port = int(port)
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            s.settimeout(timeout)
            s.connect((ip, port))
            s.send(data)
            s.shutdown(socket.SHUT_WR)
            ret = []
            while True:
                data = s.recv(1024)
                #print("received: %r" % data)
                if not data:
                    return b''.join(ret)
                ret.append(data)
        finally:
            s.close()

    def test_wait_query_1(self):

        ## NB NB NB:
        # if there was a scheduler and a poller running and connected to `self.livestatus_broker.from_q´
        # then the prepended command would be processed and the livestatus **could** return a different response,
        # **depending** on :
        # the wait_timeout_sec we use (which then should be greater than the randint(1,3) actually used)
        # and on how fast the command would be processed by the poller
        # and on how fast the host status would be updated after while within the livestatus process/thread.

        now = int( time.time() ) # current livestatus wants an INTEGER value for WaitTimeout ..
        host = 'test_host_0'
        # following request is nearly exactly what check_mk send (modulo the timestamps/hostname)
        # when you reschedule an immediate host check :
        command = b'[{now}] SCHEDULE_FORCED_HOST_CHECK;{host};{now}'.format(
                    now=now, host=host)
        wait_timeout_sec = randint(1, 3)
        request = b'''COMMAND {cmd}
GET hosts
WaitObject: {host}
WaitCondition: last_check >= {last_check}
WaitTimeout: {wait_timeout}
WaitTrigger: check
Columns: last_check state plugin_output
Filter: host_name = {host}
Localtime: {localtime}
OutputFormat: python
KeepAlive: on
ResponseHeader: fixed16
ColumnHeaders: true

'''.format(last_check=now, localtime=now, host=host, now=now, cmd=command,
           wait_timeout=wait_timeout_sec * 1000, # WaitTimeout header field is in millisecs
        )

        t0 = time.time()
        response = self.query_livestatus(request)
        t1 = time.time()
        self.assertLess(wait_timeout_sec, t1 - t0,
                        '(actually and in this very specific case) livestatus should take at least the requested WaitTimeout (%s sec) to complete ; response=%s' %
                        (wait_timeout_sec, response))
        goodresponse = "200          13\n[[0, 0, '']]\n"
        self.assertEqual(goodresponse, response)
        data = self.livestatus_broker.from_q.get(block=False)
        self.assertIsInstance(data, ExternalCommand)
        self.assertEqual(command, data.cmd_line)

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


<commit_before>#!/usr/bin/python
# Core Configuration File For Commands and Options

from __future__ import print_function
import sys
import logging

#from sniff import sniffredis
#from sniff import sniffmongo
#from sniff import sniffcouch

from termcolor import colored

from dbattacks import mongoattacks
from dbattacks import couchattacks
from dbattacks import redisattacks
from dbattacks import hbaseattacks
from dbattacks import cassattacks
from dbattacks import utils

import webattacks
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
logging.getLogger("ctype.runtime").setLevel(logging.ERROR)

def Config(args):

    mas = False
    paramcheck = []
    specify_params = []
    available = ['mongo', 'couch', 'redis']
    post_status = False
    target = args['ip']
    port = args['port']
    url = args['webapp']
    seldb = args['enum']
    filename = args['file']

    try:

        # Checks whether Host is up
        if args['ip']:
            utils.host_up(target)

        # Credentials
        screen = args['screen'] if args['screen'] else False
        creds = args['auth'] if args['auth'] else False
        authall = args['authall'] if args['authall'] else False
        mass = args['mass'] if args['mass'] else False
        db = args['db'] if args['db'] else 'admin'
        column_select = args['c'] if args['c'] else False
        dump = True if args['dump'] else False
        post_status = True if args['post'] else False
        limit = int(args['limit']) if args['limit'] else 0
        write = args['write'] if args['write'] else False

        # Scan for General DB Targets
        if args['scan']:
            utils.scan_target(target)

        # Web Attacks
        # This is argument is not working correctly - Need to fix - th3r3p0
        # if args['url']:
        #	seldb=args['webapp'] if args['webapp'] else False
        #	if seldb == 'mongo':
        #		filename=['payload/js_inject.txt','payload/js_time']

        # Dictionary Attacks

        if args['dict']:
            seldb = args['dict']
            if args['file']:
                if seldb == 'mongo':
                    if args['port'] or args['db']:
                        pass
                    else:
                        port = 27017
                        db = 'admin'
                    # mongoattacks.mongo_web_interface(target,port,creds,screen)
                    mongoattacks.dict_mongo(filename, target, port, db)
                elif seldb == 'couch':
                    if args['port']:
                        pass
                    else:
                        port = 5984
                    couchattacks.dict_couch(filename, target, port)
                elif seldb == 'redis':
                    if args['port']:
                        pass
                    else:
                        port = 6379
                    redisattacks.dict_redis(filename, target, port)
            else:
                print(colored("[-] Specify File Name", 'red'))

        # Enumeration Check
        if args['enum']:
            seldb = args['enum']
            if seldb == 'mongo':
                if port:
                    pass
                else:
                    port = 27017
                # mongo_web_scan(target)
                try:
                    conn = mongoattacks.mongo_conn(target, port, mass)
                    mongoattacks.mongo_enum(
                        conn, creds, authall, db, column_select, dump, limit, write)
                except Exception as e:
                    print(colored(e, 'red'))
            elif seldb == 'couch':
                if port:
                    pass
                else:
                    port = 5984
                try:
                    # print post_status
                    if db == 'admin':
                        db = ""
                    couch = couchattacks.couch_conn(target, port)
                    couchattacks.couch_enum(
                        couch, target, port, creds, db, column_select, post_status)
                except Exception:
                    print(Exception)
                    print(colored("[-] Enumeration Failed \n", 'red'))
            elif seldb == 'redis':
                if port:
                    pass
                else:
                    port = 6379
                try:
                    r_server = redisattacks.redis_conn(target, port)
                    redisattacks.redis_enum(r_server, creds)
                except Exception:
                    print(colored(e, 'red'))
            elif seldb == 'cassandra':
                if port:
                    pass
                else:
                    port = 9160
                    creds = False
                if db == 'admin':
                    db = False
                cassattacks.cassa_enum(target, port, db, dump)
            elif seldb == 'hbase':
                if port:
                    pass
                else:
                    port = 8080
                hbaseattacks.hbase_enum(port)
            else:
                print(colored("[-] No Support for the Specified DB", 'red'))

        # Mass Scan Settings

        if args['mass'] in available:
            select = args['mass']
            if args['file']:
                mas = True
                mass_scan(args['mass'], args['file'])
            else:
                print(colored("[-] Plse specify File name \n", 'red'))

        # Database Select (Currently available for Mongo,Couch)
        if args['db']:
            db_select = args['db']
            column_select = args['c']
        else:
            db_select = ""
            if args['post'] == 'enable':
                post_status = True

        if args['param']:
            paramcheck = args['param']
            specify_params = paramcheck.split(',')
        else:
            pass

        specify_params = args['param']

        # Scans for WebAPP Attacks
        if args['webapp']:
            webattacks.nosqlweb.attack(url)

        # Redis DOS (2.6+)
        if args['exhaust']:
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_exhaust(target, port)

        # Redis RCE Check
        if args['remotecheck']:
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_rce(target, port)

        # Redis File Enumeration Check
        if args['filecheck']:
            filename = args['filecheck']
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_file_enum(target, filename, port, creds)

        # Shodan IP Grabber
        if args['shodan']:
            utils.shodan_frame(args['shodan'])

        # Sniffing Module
        if args['sniff'] == 'mongo':
            sniffmongo.sniff_mongo()
        if args['sniff'] == 'redis':
            sniffredis.sniff_redis()
        if args['sniff'] == 'couch':
            sniffcouch.sniff_couch()

        # Clone Database Currently Available for Mongo,Couch and Redis
        if args['clone'] == 'couch':
            couchattacks.clone_couch(target)
        if args['clone'] == 'redis':
            redisattacks.clone_redis(target)
    except KeyboardInterrupt:
        print(colored("[-] Cntrl+C Shutting Down", 'red'))
        sys.exit(0)
<commit_msg>Updated removed web attacks for now<commit_after>#!/usr/bin/python
# Core Configuration File For Commands and Options

from __future__ import print_function
import sys
import logging

#from sniff import sniffredis
#from sniff import sniffmongo
#from sniff import sniffcouch

from termcolor import colored

from dbattacks import mongoattacks
from dbattacks import couchattacks
from dbattacks import redisattacks
from dbattacks import hbaseattacks
from dbattacks import cassattacks
from dbattacks import utils

#import webattacks
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
logging.getLogger("ctype.runtime").setLevel(logging.ERROR)

def Config(args):

    mas = False
    paramcheck = []
    specify_params = []
    available = ['mongo', 'couch', 'redis']
    post_status = False
    target = args['ip']
    port = args['port']
    url = args['webapp']
    seldb = args['enum']
    filename = args['file']

    try:

        # Checks whether Host is up
        if args['ip']:
            utils.host_up(target)

        # Credentials
        screen = args['screen'] if args['screen'] else False
        creds = args['auth'] if args['auth'] else False
        authall = args['authall'] if args['authall'] else False
        mass = args['mass'] if args['mass'] else False
        db = args['db'] if args['db'] else 'admin'
        column_select = args['c'] if args['c'] else False
        dump = True if args['dump'] else False
        post_status = True if args['post'] else False
        limit = int(args['limit']) if args['limit'] else 0
        write = args['write'] if args['write'] else False

        # Scan for General DB Targets
        if args['scan']:
            utils.scan_target(target)

        # Web Attacks
        # This is argument is not working correctly - Need to fix - th3r3p0
        # if args['url']:
        #	seldb=args['webapp'] if args['webapp'] else False
        #	if seldb == 'mongo':
        #		filename=['payload/js_inject.txt','payload/js_time']

        # Dictionary Attacks

        if args['dict']:
            seldb = args['dict']
            if args['file']:
                if seldb == 'mongo':
                    if args['port'] or args['db']:
                        pass
                    else:
                        port = 27017
                        db = 'admin'
                    # mongoattacks.mongo_web_interface(target,port,creds,screen)
                    mongoattacks.dict_mongo(filename, target, port, db)
                elif seldb == 'couch':
                    if args['port']:
                        pass
                    else:
                        port = 5984
                    couchattacks.dict_couch(filename, target, port)
                elif seldb == 'redis':
                    if args['port']:
                        pass
                    else:
                        port = 6379
                    redisattacks.dict_redis(filename, target, port)
            else:
                print(colored("[-] Specify File Name", 'red'))

        # Enumeration Check
        if args['enum']:
            seldb = args['enum']
            if seldb == 'mongo':
                if port:
                    pass
                else:
                    port = 27017
                # mongo_web_scan(target)
                try:
                    conn = mongoattacks.mongo_conn(target, port, mass)
                    mongoattacks.mongo_enum(
                        conn, creds, authall, db, column_select, dump, limit, write)
                except Exception as e:
                    print(colored(e, 'red'))
            elif seldb == 'couch':
                if port:
                    pass
                else:
                    port = 5984
                try:
                    # print post_status
                    if db == 'admin':
                        db = ""
                    couch = couchattacks.couch_conn(target, port)
                    couchattacks.couch_enum(
                        couch, target, port, creds, db, column_select, post_status)
                except Exception:
                    print(Exception)
                    print(colored("[-] Enumeration Failed \n", 'red'))
            elif seldb == 'redis':
                if port:
                    pass
                else:
                    port = 6379
                try:
                    r_server = redisattacks.redis_conn(target, port)
                    redisattacks.redis_enum(r_server, creds)
                except Exception:
                    print(colored(e, 'red'))
            elif seldb == 'cassandra':
                if port:
                    pass
                else:
                    port = 9160
                    creds = False
                if db == 'admin':
                    db = False
                cassattacks.cassa_enum(target, port, db, dump)
            elif seldb == 'hbase':
                if port:
                    pass
                else:
                    port = 8080
                hbaseattacks.hbase_enum(port)
            else:
                print(colored("[-] No Support for the Specified DB", 'red'))

        # Mass Scan Settings

        if args['mass'] in available:
            select = args['mass']
            if args['file']:
                mas = True
                mass_scan(args['mass'], args['file'])
            else:
                print(colored("[-] Plse specify File name \n", 'red'))

        # Database Select (Currently available for Mongo,Couch)
        if args['db']:
            db_select = args['db']
            column_select = args['c']
        else:
            db_select = ""
            if args['post'] == 'enable':
                post_status = True

        if args['param']:
            paramcheck = args['param']
            specify_params = paramcheck.split(',')
        else:
            pass

        specify_params = args['param']

        # Scans for WebAPP Attacks
        if args['webapp']:
            webattacks.nosqlweb.attack(url)

        # Redis DOS (2.6+)
        if args['exhaust']:
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_exhaust(target, port)

        # Redis RCE Check
        if args['remotecheck']:
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_rce(target, port)

        # Redis File Enumeration Check
        if args['filecheck']:
            filename = args['filecheck']
            if port:
                pass
            else:
                port = 6379
            redisattacks.redis_file_enum(target, filename, port, creds)

        # Shodan IP Grabber
        if args['shodan']:
            utils.shodan_frame(args['shodan'])

        # Sniffing Module
        if args['sniff'] == 'mongo':
            sniffmongo.sniff_mongo()
        if args['sniff'] == 'redis':
            sniffredis.sniff_redis()
        if args['sniff'] == 'couch':
            sniffcouch.sniff_couch()

        # Clone Database Currently Available for Mongo,Couch and Redis
        if args['clone'] == 'couch':
            couchattacks.clone_couch(target)
        if args['clone'] == 'redis':
            redisattacks.clone_redis(target)
    except KeyboardInterrupt:
        print(colored("[-] Cntrl+C Shutting Down", 'red'))
        sys.exit(0)
<|endoftext|>