<commit_before>from .utils import config, checks, formats
import discord
from discord.ext import commands
import discord.utils
import random, json, asyncio
from urllib.parse import unquote
import os

class Radio:
    """The radio-bot related commands."""

    def __init__(self, bot):
        self.bot = bot
        if not discord.opus.is_loaded():
            discord.opus.load_opus('/usr/local/lib/libopus.so') #FreeBSD path

        self.player = None
        #self.stopped = True
        self.q = asyncio.Queue()
        self.play_next_song = asyncio.Event()
        self.stopped = asyncio.Event()
        self.current_song = None
        self.songs_dir = 'radio/'
        self.songs = []
        self.update_song_list()

    @property
    def is_playing(self):
        return self.player is not None and self.player.is_playing() and not self.stopped.is_set()

    def toggle_next_song(self):
        self.bot.loop.call_soon_threadsafe(self.bot.change_status(None))
        if not self.stopped.is_set():
            self.bot.loop.call_soon_threadsafe(self.play_next_song.set)

    def update_song_list(self):
        self.songs = self.bot.pycopy.list_files(self.songs_dir)

    @commands.command(pass_context=True)
    async def join(self, ctx, *, channel_name : str):
        """Зайти на указанный голосовой канал."""
        if self.bot.is_voice_connected():
            await ctx.invoke(self.leave)

        check = lambda c: c.name == channel_name and c.type == discord.ChannelType.voice
        channel = discord.utils.find(check, ctx.message.server.channels)
        if channel is None:
            await self.bot.say('Нет такого голосового канала.')
        await self.bot.join_voice_channel(channel)

    @commands.command(pass_context=True)
    async def leave(self, ctx):
        """Покинуть текущий голосовой каналю"""
        await ctx.invoke(self.stop)
        await self.bot.voice.disconnect()

    @commands.command()
    async def pause(self):
        """Приостановить воспроизведение."""
        if self.player is not None:
            self.player.pause()

    @commands.command()
    async def resume(self):
        """Продолжить воспроизведение."""
        if self.player is not None and not self.is_playing:
            self.player.resume()

    @commands.command()
    async def skip(self):
        """Перейти к следующей песне в очереди."""
        if self.player is not None and self.is_playing:
            self.player.stop()
            self.toggle_next_song()

    @commands.command()
    async def stop(self):
        """Остановить воспроизведение."""
        if self.is_playing:
            self.stopped.set()
            self.player.stop()
            self.player = None

    @commands.command(pass_context=True)
    async def play(self, ctx):
        """Начать воспроизведение песен из очереди."""
        self.stopped.clear()
        if self.player is not None:# and not self.stopped.is_set():
            if not self.is_playing:
                await ctx.invoke(self.resume)
                return
            else:
                await self.bot.say('Уже играю песенку.')
                return

        while True:
            if not self.bot.is_voice_connected():
                author_channel = ctx.message.author.voice_channel
                if author_channel is not None:
                    await ctx.invoke(self.join, channel_name=author_channel.name)
                else:
                    await self.bot.say('Не выбран голосовой канал.')
                    return

            if self.stopped.is_set():
                break
            if self.q.empty():
                await self.q.put(random.choice(self.songs))
            self.play_next_song.clear()
            self.current = await self.q.get()
            self.player = self.bot.voice.create_ffmpeg_player(
                self.bot.pycopy.direct_link(self.songs_dir + self.current),
                after=self.toggle_next_song,
                #options="-loglevel debug -report",
                headers = dict(self.bot.pycopy.session.headers))
            #self.stopped.clear()
            self.player.start()
            song_name = unquote(self.current.split('/')[-1])
            await self.bot.change_status(discord.Game(name=song_name))

            await self.play_next_song.wait()

    @commands.command(aliases=['c'])
    async def current(self):
        """Что там играет?"""
        if self.is_playing:
            song_name = unquote(self.current.split('/')[-1])
            await self.bot.say(song_name)

    @commands.command()
    async def update(self):
        """Обновить список песен."""
        self.update_song_list()
        await self.bot.say("Найдено {} песенок".format(len(self.songs)))

    @commands.command()
    async def list(self):
        """Вывести список всех доступных песен."""
        song_list = ""
        id = 1
        for song in self.songs:
            song_list += "{}. {}\n".format(id, song)
            id += 1
            if len(song_list) > 1800:
                await self.bot.say(song_list)
                song_list = ''
        await self.bot.say(song_list)

    @commands.command()
    async def add(self, song_num : int):
        """Добавить в конец очереди песню с данным номером."""
        await self.q.put(self.songs[song_num-1])
        await self.bot.say("{} будет следующей песенкой".format(self.songs[song_num-1]))

def setup(bot):
    bot.add_cog(Radio(bot))
<commit_msg>Update radio.py<commit_after>from .utils import config, checks, formats
import discord
from discord.ext import commands
import discord.utils
import random, json, asyncio
from urllib.parse import unquote
import os

class Radio:
    """The radio-bot related commands."""

    def __init__(self, bot):
        self.bot = bot
        if not discord.opus.is_loaded():
            discord.opus.load_opus('/usr/local/lib/libopus.so') #FreeBSD path

        self.player = None
        #self.stopped = True
        self.q = asyncio.Queue()
        self.play_next_song = asyncio.Event()
        self.stopped = asyncio.Event()
        self.current_song = None
        self.songs_dir = 'radio/'
        self.songs = []
        self.update_song_list()

    @property
    def is_playing(self):
        return self.player is not None and self.player.is_playing() and not self.stopped.is_set()

    def toggle_next_song(self):
        self.bot.loop.call_soon_threadsafe(self.bot.change_status(None))
        if not self.stopped.is_set():
            self.bot.loop.call_soon_threadsafe(self.play_next_song.set)

    def update_song_list(self):
        self.songs = self.bot.pycopy.list_files(self.songs_dir)

    @commands.command(pass_context=True)
    async def join(self, ctx, *, channel_name : str):
        """Зайти на указанный голосовой канал."""
        if self.bot.is_voice_connected():
            await ctx.invoke(self.leave)

        check = lambda c: c.name == channel_name and c.type == discord.ChannelType.voice
        channel = discord.utils.find(check, ctx.message.server.channels)
        if channel is None:
            await self.bot.say('Нет такого голосового канала.')
        await self.bot.join_voice_channel(channel)

    @commands.command(pass_context=True)
    async def leave(self, ctx):
        """Покинуть текущий голосовой каналю"""
        await ctx.invoke(self.stop)
        await self.bot.voice.disconnect()

    @commands.command()
    async def pause(self):
        """Приостановить воспроизведение."""
        if self.player is not None:
            self.player.pause()

    @commands.command()
    async def resume(self):
        """Продолжить воспроизведение."""
        if self.player is not None and not self.is_playing:
            self.player.resume()

    @commands.command()
    async def skip(self):
        """Перейти к следующей песне в очереди."""
        if self.player is not None and self.is_playing:
            self.player.stop()
            self.toggle_next_song()

    @commands.command()
    async def stop(self):
        """Остановить воспроизведение."""
        if self.is_playing:
            self.stopped.set()
            self.player.stop()
            self.player = None

    @commands.command(pass_context=True)
    async def play(self, ctx):
        """Начать воспроизведение песен из очереди."""
        self.stopped.clear()
        if self.player is not None:# and not self.stopped.is_set():
            if not self.is_playing:
                await ctx.invoke(self.resume)
                return
            else:
                await self.bot.say('Уже играю песенку.')
                return

        while True:
            if not self.bot.is_voice_connected():
                author_channel = ctx.message.author.voice_channel
                if author_channel is not None:
                    await ctx.invoke(self.join, channel_name=author_channel.name)
                else:
                    await self.bot.say('Не выбран голосовой канал.')
                    return

            if self.stopped.is_set():
                break
            if self.q.empty():
                await self.q.put(random.choice(self.songs))
            self.play_next_song.clear()
            self.current = await self.q.get()
            self.player = self.bot.voice.create_ffmpeg_player(
                self.bot.pycopy.direct_link(self.songs_dir + self.current),
                after=self.toggle_next_song,
                #options="-loglevel debug -report",
                headers = dict(self.bot.pycopy.session.headers))
            #self.stopped.clear()
            self.player.start()
            song_name = unquote(self.current.split('/')[-1])
            await self.bot.change_status(discord.Game(name=song_name))

            await self.play_next_song.wait()
        await self.bot.say('Leaving player loop')

    @commands.command(aliases=['c'])
    async def current(self):
        """Что там играет?"""
        if self.is_playing:
            song_name = unquote(self.current.split('/')[-1])
            await self.bot.say(song_name)

    @commands.command()
    async def update(self):
        """Обновить список песен."""
        self.update_song_list()
        await self.bot.say("Найдено {} песенок".format(len(self.songs)))

    @commands.command()
    async def list(self):
        """Вывести список всех доступных песен."""
        song_list = ""
        id = 1
        for song in self.songs:
            song_list += "{}. {}\n".format(id, song)
            id += 1
            if len(song_list) > 1800:
                await self.bot.say(song_list)
                song_list = ''
        await self.bot.say(song_list)

    @commands.command()
    async def add(self, song_num : int):
        """Добавить в конец очереди песню с данным номером."""
        await self.q.put(self.songs[song_num-1])
        await self.bot.say("{} будет следующей песенкой".format(self.songs[song_num-1]))

def setup(bot):
    bot.add_cog(Radio(bot))
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import division
from __future__ import print_function
from alsaaudio import Mixer
from random import choice
from serial import Serial
from urllib2 import urlopen, HTTPError
import argparse
import json
import os
import subprocess as sp
import sys
import time

# This is a list of sample songs that will randomly play if the
# user is misidentified or does not exist!
DEFAULT_SONGS = os.listdir("/home/pi/random")

SONG_EXTS = (
    ".mp3", ".mp4", ".m4a", ".m4p",
    ".flac", ".ogg", ".oga", ".wav",
    ".wma"
)

DING_SONG = "/home/pi/ding.mp3"

MPLAYER_FIFO = "/tmp/mplayer.fifo"

FNULL = open(os.devnull, 'w')

class MockSerial:

    def __init__(self, fi=sys.stdin):
        self.fi = fi

    def readline(self):
        return self.fi.readline()

    def flushInput(self):
        return self.fi.flush()

def quiet_hours():
    ' Returns True if the current time is within RIT quiet hours '
    currtime = time.localtime()
    if currtime.tm_wday > 4:
        return (currtime.tm_hour + 23) % 24 < 6
    else:
        return (currtime.tm_hour + 1) % 24 < 8

def read_ibutton(varID, cache={}):
    '''
    Use Nick Depinet's LDAP service to convert iButtons to usernames

    Caches values when possible (iButtons don't really change)
    '''
    if varID in cache:
        return cache[varID]
    try:
        data = urlopen('http://www.csh.rit.edu:56124/?ibutton=' + varID)
        usernameData = json.load(data)
    except HTTPError as error:
        # Need to check its an 404, 503, 500, 403 etc.
        print(error.read())
    except ValueError as error:
        # Got malformed JSON somehow
        print(error)
    else:
        cache[varID] = usernameData['username'][0], usernameData['homeDir'][0]
        return cache[varID]
    return "", ""

def get_user_song(homedir):
    '''
    Load one of the following files:
    ~/harold.mp3
    ~/harold/*, of one of the supported file types
    '''
    if homedir:
        print("Home:", homedir)
        hdir = os.path.join(homedir, "harold")
        hfile = os.path.join(homedir, "harold.mp3")
        if os.path.isdir(hdir):
            playlist = [os.path.join(hdir, f)
                        for f in os.listdir(hdir)
                        if os.path.isfile(os.path.join(hdir, f))
                        and f.endswith(SONG_EXTS)]
            return choice(playlist or DEFAULT_SONGS)
        elif os.path.isfile(hfile):
            return hfile
    return choice(DEFAULT_SONGS)

class Harold(object):

    def __init__(self, mplfifo, ser, mpout, beep=True):
        self.playing = False
        self.mixer = Mixer(control='PCM')
        self.fifo = mplfifo
        self.ser = ser
        self.mpout = mpout
        self.beep = beep

    def write(self, *args, **kwargs):
        delay = kwargs.pop("delay", 0.5)
        kws = {"file": self.fifo}
        kws.update(kwargs)
        print(*args, **kws)
        time.sleep(delay)

    def __call__(self):
        if not self.playing:
            # Lower the volume during quiet hours... Don't piss off the RA!
            self.mixer.setvolume(85 if quiet_hours() else 100)
            varID = self.ser.readline()
            print(varID)
            # mplayer will play any files sent to the FIFO file.
            if self.beep:
                self.write("loadfile", DING_SONG)
            if "ready" not in varID:
                # Get the username from the ibutton
                username, homedir = read_ibutton(varID)

                # Print the user's name (Super handy for debugging...)
                print("New User: '" + username + "'")

                song = get_user_song(homedir)
                print("Now playing '" + song + "'...\n")
                self.write("loadfile '" + song.replace("'", "\\'") + "'",
                           delay=0.0)

                self.write("get_time_length")
                line = self.mpout.readline()
                while not line.startswith("ANS_LENGTH="):
                    line = self.mpout.readline()
                duration = float(line.strip().split("=")[-1])

                self.starttime = time.time()
                self.endtime = time.time() + min(30, duration)
                self.playing = True

        elif time.time() >= self.endtime:
            self.write("stop")
            self.playing = False
            self.ser.flushInput()
            print("Stopped\n")

        elif time.time() >= self.starttime+25:
            # Fade out the music at the end.
            vol = int(self.mixer.getvolume()[0])
            if vol > 60:
                vol -= 1 + (100 - vol)/3.
                self.mixer.setvolume(int(vol))
                time.sleep(0.1)

def main():
    # Handle some command line arguments
    parser = argparse.ArgumentParser(description="Start Harold system")
    parser.add_argument("--debug", "-d", action="store_true",
                        help="Use debug mode (stdin)")
    parser.add_argument("--serial", "-s",
                        default="/dev/ttyACM0",
                        help="Serial port to use", metavar="PORT")
    parser.add_argument("--rate", "-r",
                        default=9600, type=int,
                        help="Serial device to use")
    parser.add_argument("--fifo", "-f",
                        default="/tmp/mplayer.fifo",
                        help="FIFO to communicate to mplayer with")
    parser.add_argument("--nobeep", "-n", action="store_true",
                        help="Disable beep")
    args = parser.parse_args()
    try:
        os.mkfifo(args.fifo)
    except OSError as e:
        import errno
        if e.errno != errno.EEXIST:
            raise
    cmd = ["mplayer", "-idle", "-slave", "-input", "file="+args.fifo]
    mplayer = sp.Popen(cmd, stdout=sp.PIPE, stderr=FNULL)
    try:
        with open(args.fifo, "w", 0) as mplfifo:
            if args.debug:
                ser = MockSerial()
            else:
                ser = Serial(args.serial, args.rate)
                ser.flushInput()
            harold = Harold(mplfifo, ser, mplayer.stdout, not args.nobeep)
            while True:
                harold()
    except KeyboardInterrupt:
        print("Shutting down")
    finally:
        mplayer.kill()
        os.remove(args.fifo)

if __name__ == '__main__':
    main()
<commit_msg>Absolute random filenames<commit_after>#!/usr/bin/env python
from __future__ import division
from __future__ import print_function
from alsaaudio import Mixer
from random import choice
from serial import Serial
from urllib2 import urlopen, HTTPError
import argparse
import json
import os
import subprocess as sp
import sys
import time

# This is a list of sample songs that will randomly play if the
# user is misidentified or does not exist!
DEFAULT_SONGS = map(lambda f: os.path.join("/home/pi/random", f),
                    os.listdir("/home/pi/random"))

SONG_EXTS = (
    ".mp3", ".mp4", ".m4a", ".m4p",
    ".flac", ".ogg", ".oga", ".wav",
    ".wma"
)

DING_SONG = "/home/pi/ding.mp3"

MPLAYER_FIFO = "/tmp/mplayer.fifo"

FNULL = open(os.devnull, 'w')

class MockSerial:

    def __init__(self, fi=sys.stdin):
        self.fi = fi

    def readline(self):
        return self.fi.readline()

    def flushInput(self):
        return self.fi.flush()

def quiet_hours():
    ' Returns True if the current time is within RIT quiet hours '
    currtime = time.localtime()
    if currtime.tm_wday > 4:
        return (currtime.tm_hour + 23) % 24 < 6
    else:
        return (currtime.tm_hour + 1) % 24 < 8

def read_ibutton(varID, cache={}):
    '''
    Use Nick Depinet's LDAP service to convert iButtons to usernames

    Caches values when possible (iButtons don't really change)
    '''
    if varID in cache:
        return cache[varID]
    try:
        data = urlopen('http://www.csh.rit.edu:56124/?ibutton=' + varID)
        usernameData = json.load(data)
    except HTTPError as error:
        # Need to check its an 404, 503, 500, 403 etc.
        print(error.read())
    except ValueError as error:
        # Got malformed JSON somehow
        print(error)
    else:
        cache[varID] = usernameData['username'][0], usernameData['homeDir'][0]
        return cache[varID]
    return "", ""

def get_user_song(homedir):
    '''
    Load one of the following files:
    ~/harold.mp3
    ~/harold/*, of one of the supported file types
    '''
    if homedir:
        print("Home:", homedir)
        hdir = os.path.join(homedir, "harold")
        hfile = os.path.join(homedir, "harold.mp3")
        if os.path.isdir(hdir):
            playlist = [os.path.join(hdir, f)
                        for f in os.listdir(hdir)
                        if os.path.isfile(os.path.join(hdir, f))
                        and f.endswith(SONG_EXTS)]
            return choice(playlist or DEFAULT_SONGS)
        elif os.path.isfile(hfile):
            return hfile
    return choice(DEFAULT_SONGS)

class Harold(object):

    def __init__(self, mplfifo, ser, mpout, beep=True):
        self.playing = False
        self.mixer = Mixer(control='PCM')
        self.fifo = mplfifo
        self.ser = ser
        self.mpout = mpout
        self.beep = beep

    def write(self, *args, **kwargs):
        delay = kwargs.pop("delay", 0.5)
        kws = {"file": self.fifo}
        kws.update(kwargs)
        print(*args, **kws)
        time.sleep(delay)

    def __call__(self):
        if not self.playing:
            # Lower the volume during quiet hours... Don't piss off the RA!
            self.mixer.setvolume(85 if quiet_hours() else 100)
            varID = self.ser.readline()
            print(varID)
            # mplayer will play any files sent to the FIFO file.
            if self.beep:
                self.write("loadfile", DING_SONG)
            if "ready" not in varID:
                # Get the username from the ibutton
                username, homedir = read_ibutton(varID)

                # Print the user's name (Super handy for debugging...)
                print("New User: '" + username + "'")

                song = get_user_song(homedir)
                print("Now playing '" + song + "'...\n")
                self.write("loadfile '" + song.replace("'", "\\'") + "'",
                           delay=0.0)

                self.write("get_time_length")
                line = self.mpout.readline()
                while not line.startswith("ANS_LENGTH="):
                    line = self.mpout.readline()
                duration = float(line.strip().split("=")[-1])

                self.starttime = time.time()
                self.endtime = time.time() + min(30, duration)
                self.playing = True

        elif time.time() >= self.endtime:
            self.write("stop")
            self.playing = False
            self.ser.flushInput()
            print("Stopped\n")

        elif time.time() >= self.starttime+25:
            # Fade out the music at the end.
            vol = int(self.mixer.getvolume()[0])
            if vol > 60:
                vol -= 1 + (100 - vol)/3.
                self.mixer.setvolume(int(vol))
                time.sleep(0.1)

def main():
    # Handle some command line arguments
    parser = argparse.ArgumentParser(description="Start Harold system")
    parser.add_argument("--debug", "-d", action="store_true",
                        help="Use debug mode (stdin)")
    parser.add_argument("--serial", "-s",
                        default="/dev/ttyACM0",
                        help="Serial port to use", metavar="PORT")
    parser.add_argument("--rate", "-r",
                        default=9600, type=int,
                        help="Serial device to use")
    parser.add_argument("--fifo", "-f",
                        default="/tmp/mplayer.fifo",
                        help="FIFO to communicate to mplayer with")
    parser.add_argument("--nobeep", "-n", action="store_true",
                        help="Disable beep")
    args = parser.parse_args()
    try:
        os.mkfifo(args.fifo)
    except OSError as e:
        import errno
        if e.errno != errno.EEXIST:
            raise
    cmd = ["mplayer", "-idle", "-slave", "-input", "file="+args.fifo]
    mplayer = sp.Popen(cmd, stdout=sp.PIPE, stderr=FNULL)
    try:
        with open(args.fifo, "w", 0) as mplfifo:
            if args.debug:
                ser = MockSerial()
            else:
                ser = Serial(args.serial, args.rate)
                ser.flushInput()
            harold = Harold(mplfifo, ser, mplayer.stdout, not args.nobeep)
            while True:
                harold()
    except KeyboardInterrupt:
        print("Shutting down")
    finally:
        mplayer.kill()
        os.remove(args.fifo)

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


<commit_before># -*- coding: utf-8 -*-
""" \file saml/schema/element.py
\brief Defines base element classes for the schema package.

\author Ryan Leckey (mehcode) leckey.ryan@gmail.com

\copyright Copyright © 2012 Concordus Applications, Inc.
           \n \n
           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:
           \n \n
           The above copyright notice and this permission notice shall be
           included in all copies or substantial portions of the Software.
           \n \n
           THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
           EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
           MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
           NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
           BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
           ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
           CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
           SOFTWARE.
"""
import inspect
from lxml.builder import ElementMaker
from lxml import etree
from . import attribute

class Element(object):
    """Represents a generic element on an XML element.
    """

    class Meta:
        """Base meta object if none is provided in derived object.
        """
        pass

    @staticmethod
    def deserialize(text):
        """Deserialize the passed text as an element.
        """
        pass

    @classmethod
    def _get_ordered_members(cls):
        """"""
        sorted = []

        # Add members from base classes, recursively
        for base in cls.__bases__:
            try:
                sorted.extend(base._get_ordered_members())
            except AttributeError:
                # Reached object, we're done
                pass

        # Add members from calling class
        for name, value in cls.__dict__.items():
            if isinstance(value, Element):
                if hasattr(value, '_meta') and hasattr(value._meta, 'index'):
                    sorted.append((value._meta.index, name, value))
                else:
                    sorted.append((0, name, value))
            if isinstance(value, Simple):
                sorted.append((value.index, name, value))
            elif isinstance(value, attribute.Attribute):
                sorted.append((0, name, value))

        return sorted

    @classmethod
    def serialize(cls, obj):
        """Serializes the passed element into an XML representation."""
        # Instantiate an element maker context
        E = ElementMaker(
            namespace=obj._meta.namespace[1],
            nsmap={obj._meta.namespace[0]: obj._meta.namespace[1]}
        )

        # Instantiate an XML element with its name
        xml = E(obj._meta.name)

        # Append text content (if available)
        if hasattr(obj, "text"):
            xml.text = str(obj.text)

        # Iterate through ordered list of members to serialize and append
        for unused, name, value in cls._get_ordered_members():
            # Was this member provided ?
            attr = obj.__dict__.get(name)
            if attr is not None and value is not None:
                if isinstance(value, attribute.Attribute):
                    # Set as an attribute if it is derived from
                    # the base Attribute class.
                    xml.set(value.name, value.serialize(attr))
                elif isinstance(value, Element):
                    # Serialize and append as an element if it is derived
                    # from the base Element class
                    try:
                        # First; try appending several as an iterable
                        for item in attr:
                            xml.append(value.serialize(item))
                    except TypeError:
                        # Didn't work; this must be a singular element
                        xml.append(value.serialize(attr))

        # Return serialized XML
        return xml

#    def serialize(self):
#        """Serialize the current element as text.
#        """
#        # Instantiate an element maker tailored for this element
#        E = ElementMaker(
#            namespace=self._meta.namespace[1],
#            nsmap={self._meta.namespace[0]: self._meta.namespace[1]})
#
#        # Instantiate an XML element with its name
#        xml = E(self._meta.name)
#
#        # Append text, if available
#        if hasattr(self, "text"):
#            xml.text = str(self.text)
#
#        # Append elements
#        for index, name, value in self._get_ordered_members():
#            # Does this exist ?
#            attr = self.__dict__.get(name)
#            if attr is not None and value is not None:
#                # Attempt to set this as an attribute
#                try:
#                    xml.set(value.name, value.serialize(attr))
#                    continue
#                except:
#                    pass
#
#                # Attempt to set this as a simple element
#                try:
#                    xml.append(etree.XML(value.serialize(attr)))
#                    continue
#                except:
#                    pass
#
#                # Loop through and append all elements as XML
#                try:
#                    for item in attr:
#                        # FIXME: This etree.XML(..) is stupid
#                        xml.append(etree.XML(item.serialize()))
#                    continue
#                except:
#                    pass
#
#                # Or can this be directly serialized as XML ?
#                try:
#                    # FIXME: This etree.XML(..) is stupid
#                    xml.append(etree.XML(attr.serialize()))
#                    continue
#                except:
#                    # We have no idea what this... just fail silently
#                    pass
#
#        # Return the XML element as string
#        return etree.tostring(xml)

    def __init__(self, text=None, **kwargs):
        """Instantiates an element reference.
        """
        ## Metadata about the element
        self._meta = self.__class__.Meta()

        # Declare any defaults, if present
        for name, value in inspect.getmembers(self.__class__):
            # If there was a default provided
            if hasattr(value, 'default') and value.default is not None:
                # Then supply it
                try:
                    # Either it is a callable that will generate a
                    # default
                    self.__dict__[name] = value.default()
                except:
                    # Or simply a scalar
                    self.__dict__[name] = value.default

        # Default value for meta elements
        self._meta.__dict__.update(inspect.getmembers(self._meta.__class__))
        self._meta.name = self.__class__.__name__

        # Iterate through and update object dictionaries from passed values
        for name, value in kwargs.items():
            segs = name.split("__")
            if segs[0] == 'meta':
                # Was passed as meta__X so set X of meta to value
                self._meta.__dict__[segs[1]] = value
            else:
                # Was passed normally, so set normally
                self.__dict__[name] = value

        # Did we receive any text ?
        if text is not None:
            self.text = text

class Simple(object):
    """
    Represents a 'simple' element that acts like an attribute of the object but
    is serialized and deserialized as an XML element.
    """
    def __init__(self, name, index=None, namespace=None, default=None):
        """Initialize a simple element with constraints on its values."""
        ## Name of the simple.
        self.name = name

        ## Namespace of the simple.
        self.namespace = namespace

        ## Index of the simple.
        self.index = index if index is not None else 0

        ## A default value that can be used.
        self.default = default

#    def serialize(self, obj):
#        """Serialize the current element as text.
#        """
#        # Instantiate an element maker tailored for this element
#        E = ElementMaker(
#            namespace=self.namespace[1],
#            nsmap={self.namespace[0]: self.namespace[1]})
#
#        # Instantiate an XML element with its name
#        xml = E(self.name)
#
#        # Append content
#        xml.text = str(obj)
#
#        # Return constructed XML block
#        return etree.tostring(xml)
<commit_msg>Wrote simple element serializer.<commit_after># -*- coding: utf-8 -*-
""" \file saml/schema/element.py
\brief Defines base element classes for the schema package.

\author Ryan Leckey (mehcode) leckey.ryan@gmail.com

\copyright Copyright © 2012 Concordus Applications, Inc.
           \n \n
           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:
           \n \n
           The above copyright notice and this permission notice shall be
           included in all copies or substantial portions of the Software.
           \n \n
           THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
           EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
           MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
           NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
           BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
           ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
           CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
           SOFTWARE.
"""
import inspect
from lxml.builder import ElementMaker
from lxml import etree
from . import attribute

class Element(object):
    """Represents a generic element on an XML element.
    """

    class Meta:
        """Base meta object if none is provided in derived object.
        """
        pass

    @classmethod
    def _get_ordered_members(cls):
        """"""
        sorted = []

        # Add members from base classes, recursively
        for base in cls.__bases__:
            try:
                sorted.extend(base._get_ordered_members())
            except AttributeError:
                # Reached object, we're done
                pass

        # Add members from calling class
        for name, value in cls.__dict__.items():
            if isinstance(value, Element):
                if hasattr(value, '_meta') and hasattr(value._meta, 'index'):
                    sorted.append((value._meta.index, name, value))
                else:
                    sorted.append((0, name, value))
            if isinstance(value, Simple):
                sorted.append((value.index, name, value))
            elif isinstance(value, attribute.Attribute):
                sorted.append((0, name, value))

        return sorted

    @classmethod
    def serialize(cls, obj):
        """Serializes the passed element into an XML representation."""
        # Instantiate an element maker context
        E = ElementMaker(
            namespace=obj._meta.namespace[1],
            nsmap={obj._meta.namespace[0]: obj._meta.namespace[1]}
        )

        # Instantiate an XML element with its name
        xml = E(obj._meta.name)

        # Append text content (if available)
        if hasattr(obj, "text"):
            xml.text = str(obj.text)

        # Iterate through ordered list of members to serialize and append
        for unused, name, value in cls._get_ordered_members():
            # Was this member provided ?
            attr = obj.__dict__.get(name)
            if attr is not None and value is not None:
                if isinstance(value, attribute.Attribute):
                    # Set as an attribute if it is derived from
                    # the base Attribute class.
                    xml.set(value.name, value.serialize(attr))
                elif isinstance(value, Element):
                    # Serialize and append as an element if it is derived
                    # from the base Element class
                    try:
                        # First; try appending several as an iterable
                        for item in attr:
                            xml.append(value.serialize(item))
                    except TypeError:
                        # Didn't work; this must be a singular element
                        xml.append(value.serialize(attr))
                elif isinstance(value, Simple):
                    # Serialize and append the simple element
                    xml.append(value.serialize(value, attr))

        # Return serialized XML
        return xml

    @staticmethod
    def deserialize(text):
        """Deserialize the passed text as an element.
        """
        pass

    def __init__(self, text=None, **kwargs):
        """Instantiates an element reference.
        """
        ## Metadata about the element
        self._meta = self.__class__.Meta()

        # Declare any defaults, if present
        for name, value in inspect.getmembers(self.__class__):
            # If there was a default provided
            if hasattr(value, 'default') and value.default is not None:
                # Then supply it
                try:
                    # Either it is a callable that will generate a
                    # default
                    self.__dict__[name] = value.default()
                except:
                    # Or simply a scalar
                    self.__dict__[name] = value.default

        # Default value for meta elements
        self._meta.__dict__.update(inspect.getmembers(self._meta.__class__))
        self._meta.name = self.__class__.__name__

        # Iterate through and update object dictionaries from passed values
        for name, value in kwargs.items():
            segs = name.split("__")
            if segs[0] == 'meta':
                # Was passed as meta__X so set X of meta to value
                self._meta.__dict__[segs[1]] = value
            else:
                # Was passed normally, so set normally
                self.__dict__[name] = value

        # Did we receive any text ?
        if text is not None:
            self.text = text

class Simple(object):
    """
    Represents a 'simple' element that acts like an attribute of the object but
    is serialized and deserialized as an XML element.
    """

    @classmethod
    def serialize(cls, obj, value):
        """Serializes the passed element as a 'simple' element."""
        # Instantiate an element maker context
        E = ElementMaker(
            namespace=obj.namespace[1],
            nsmap={obj.namespace[0]: obj.namespace[1]}
        )

        # Instantiate an XML element with its name
        xml = E(obj.name)

        # Append "simple" content
        xml.text = str(value)

        # Return constructed element
        return xml

    def __init__(self, name, index=None, namespace=None, default=None):
        """Initialize a simple element with constraints on its values."""
        ## Name of the simple.
        self.name = name

        ## Namespace of the simple.
        self.namespace = namespace

        ## Index of the simple.
        self.index = index if index is not None else 0

        ## A default value that can be used.
        self.default = default
<|endoftext|>


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

"""
    :copyright: (C) 2010-2013 by Contrail Consortium.
"""

from os.path import exists, join
from os import makedirs, remove

from conpaas.core.expose import expose
from conpaas.core.https.server import HttpJsonResponse, HttpErrorResponse
from conpaas.core.https.server import HttpFileDownloadResponse
from conpaas.services.xtreemfs.agent import role
from conpaas.core.agent import BaseAgent, AgentException
from conpaas.core.misc import run_cmd

from threading import Lock
import pickle
import base64

class XtreemFSAgent(BaseAgent):
    def __init__(self,
                 config_parser, # config file
                 **kwargs):     # anything you can't send in config_parser
                                # (hopefully the new service won't need anything extra)
        BaseAgent.__init__(self, config_parser)
        role.init(config_parser)
        self.gen_string = config_parser.get('agent','STRING_TO_GENERATE')

        self.VAR_TMP = config_parser.get('agent','VAR_TMP')

        self.dir_file = join(self.VAR_TMP,'dir.pickle')
        self.mrc_file = join(self.VAR_TMP,'mrc.pickle')
        self.osd_file = join(self.VAR_TMP,'osd.pickle')
        self.mrc_lock = Lock()
        self.dir_lock = Lock()
        self.osd_lock = Lock()

        self.DIR = role.DIR
        self.MRC = role.MRC
        self.OSD = role.OSD

    def _create(self, post_params, class_file, pClass):
        if exists(class_file):
            return HttpErrorResponse(AgentException(
                AgentException.E_CONFIG_EXISTS).message)
        try:
            if type(post_params) != dict:
                raise TypeError()
            self.logger.debug('Creating class')
            p = pClass(**post_params)
            self.logger.debug('Created class')
        except (ValueError, TypeError) as e:
            ex = AgentException(AgentException.E_ARGS_INVALID, detail=str(e))
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)
        except Exception as e:
            ex = AgentException(AgentException.E_UNKNOWN, detail=e)
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)
        else:
            try:
                self.logger.debug('Openning file %s' % class_file)
                fd = open(class_file, 'w')
                pickle.dump(p, fd)
                fd.close()
            except Exception as e:
                ex = AgentException(AgentException.E_CONFIG_COMMIT_FAILED,
                    detail=e)
                self.logger.exception(ex.message)
                return HttpErrorResponse(ex.message)
            else:
                self.logger.debug('Created class file')
                return HttpJsonResponse()

    def _stop(self, kwargs, class_file, pClass):
        self.logger.debug("_stop(kwargs=%s, class_file=%s, pClass=%s)" %
                (kwargs, class_file, pClass))

        if not exists(class_file):
            self.logger.error("class_file '%s' does not exist" % class_file)

            return HttpErrorResponse(AgentException(
                AgentException.E_CONFIG_NOT_EXIST).message)

        try:
            try:
                fd = open(class_file, 'r')
                p = pickle.load(fd)
                self.logger.debug('dump file %s loaded' %class_file)
                fd.close()
            except Exception as e:
                ex = AgentException(
                    AgentException.E_CONFIG_READ_FAILED, detail=e)
                self.logger.exception(ex.message)
                return HttpErrorResponse(ex.message)

            if 'drain' in kwargs:
                p.stop(kwargs['drain'])
            else:
                p.stop()

            self.logger.debug("Removing class_file '%s'" % class_file)
            remove(class_file)
            return HttpJsonResponse()

        except Exception as e:
            ex = AgentException(AgentException.E_UNKNOWN, detail=e)
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)

    @expose('POST')
    def createDIR(self, kwargs):
        try:
            kwargs = self._DIR_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        with self.dir_lock:
            return self._create(kwargs, self.dir_file, self.DIR)

    @expose('POST')
    def createMRC(self, kwargs):
        try:
            kwargs = self._MRC_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        else:
            with self.mrc_lock:
                return self._create(kwargs, self.mrc_file, self.MRC)

    @expose('POST')
    def createOSD(self, kwargs):
        try:
            kwargs = self._OSD_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        else:
            with self.osd_lock:
                return self._create(kwargs, self.osd_file, self.OSD)

    @expose('POST')
    def stopDIR(self, kwargs):
        """Kill the DIR service"""
        with self.dir_lock:
            return self._stop(kwargs, self.dir_file, self.DIR)

    @expose('POST')
    def stopMRC(self, kwargs):
        """Kill the MRC service"""
        with self.mrc_lock:
            return self._stop(kwargs, self.mrc_file, self.MRC)

    @expose('POST')
    def stopOSD(self, kwargs):
        """Kill the OSD service"""
        with self.osd_lock:
            return self._stop(kwargs, self.osd_file, self.OSD)

    def _DIR_get_params(self, kwargs):
        ret = {}
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        return ret

    def _MRC_get_params(self, kwargs):
        ret = {}
        if 'dir_serviceHost' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'dir service host')
        ret['dir_serviceHost'] = kwargs.pop('dir_serviceHost')
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        return ret

    def _OSD_get_params(self, kwargs):
        ret = {}
        if 'dir_serviceHost' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'dir service host')
        ret['dir_serviceHost'] = kwargs.pop('dir_serviceHost')
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        ret['mkfs'] = kwargs['mkfs']
        return ret

    @expose('POST')
    def startup(self, kwargs):
        self.state = 'RUNNING'
        self.logger.info('Agent started up')
        return HttpJsonResponse()

    @expose('POST')
    def set_snapshot(self, kwargs):
        if 'archive_url' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'archive_url')

        self.logger.info('set_snapshot: restoring archive')

        archive_cmd = "wget --ca-certificate /etc/cpsagent/certs/ca_cert.pem "
        archive_cmd += kwargs['archive_url'] + " -O - | tar xz -C /"

        self.logger.debug('set_snapshot: %s' % archive_cmd)

        out, err = run_cmd(archive_cmd)

        self.logger.debug('set_snapshot: stdout %s' % out)
        self.logger.debug('set_snapshot: stderr %s' % err)

        self.logger.info('set_snapshot: archive restored successfully')
        return HttpJsonResponse()

    @expose('POST')
    def get_snapshot(self, kwargs):
        # pack data from:
        #     /etc/cpsagent/certs/ (original agent certificates)
        #     /etc/xos/xtreemfs/ (config files, SSL-certificates, policies)
        #     /var/lib/xtreemfs/ (save everything, after shutting down the services)
        #     /var/log/xtreemfs/ (xtreemfs log files)
        filename = "/root/snapshot.tar.gz"
        dirs = "var/lib/xtreemfs/ etc/xos/xtreemfs/ var/log/xtreemfs/ etc/cpsagent/certs"

        err, out = run_cmd("tar -czf %s %s" % (filename, dirs), "/")
        if err:
            self.logger.exception(err)
            return HttpErrorResponse(err)

        return HttpFileDownloadResponse("snapshot.tar.gz", filename)

    @expose('POST')
    def set_certificates(self, kwargs):
        self.logger.debug('set_snapshot called')
    	# write certificates and truststore to the xtreemfs config directory
    	path = "/etc/xos/xtreemfs/truststore/certs"
    	if not exists(path):
  			makedirs(path)
    	open(path + "/dir.p12", 'wb').write(base64.b64decode(kwargs['dir']))
    	open(path + "/mrc.p12", 'wb').write(base64.b64decode(kwargs['mrc']))
    	open(path + "/osd.p12", 'wb').write(base64.b64decode(kwargs['osd']))
    	open(path + "/trusted.jks", 'wb').write(base64.b64decode(kwargs['truststore']))
    	return HttpJsonResponse()
<commit_msg>Fix related to issues #31, #51, #63<commit_after># -*- coding: utf-8 -*-

"""
    :copyright: (C) 2010-2013 by Contrail Consortium.
"""

from os.path import exists, join
from os import makedirs, remove

from conpaas.core.expose import expose
from conpaas.core.https.server import HttpJsonResponse, HttpErrorResponse
from conpaas.core.https.server import HttpFileDownloadResponse
from conpaas.services.xtreemfs.agent import role
from conpaas.core.agent import BaseAgent, AgentException
from conpaas.core.misc import run_cmd

from threading import Lock
import pickle
import base64

class XtreemFSAgent(BaseAgent):
    def __init__(self,
                 config_parser, # config file
                 **kwargs):     # anything you can't send in config_parser
                                # (hopefully the new service won't need anything extra)
        BaseAgent.__init__(self, config_parser)
        role.init(config_parser)
        self.gen_string = config_parser.get('agent','STRING_TO_GENERATE')

        self.VAR_TMP = config_parser.get('agent','VAR_TMP')

        self.dir_file = join(self.VAR_TMP,'dir.pickle')
        self.mrc_file = join(self.VAR_TMP,'mrc.pickle')
        self.osd_file = join(self.VAR_TMP,'osd.pickle')
        self.mrc_lock = Lock()
        self.dir_lock = Lock()
        self.osd_lock = Lock()

        self.DIR = role.DIR
        self.MRC = role.MRC
        self.OSD = role.OSD

    def _create(self, post_params, class_file, pClass):
        if exists(class_file):
            return HttpErrorResponse(AgentException(
                AgentException.E_CONFIG_EXISTS).message)
        try:
            if type(post_params) != dict:
                raise TypeError()
            self.logger.debug('Creating class')
            p = pClass(**post_params)
            self.logger.debug('Created class')
        except (ValueError, TypeError) as e:
            ex = AgentException(AgentException.E_ARGS_INVALID, detail=str(e))
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)
        except Exception as e:
            ex = AgentException(AgentException.E_UNKNOWN, detail=e)
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)
        else:
            try:
                self.logger.debug('Openning file %s' % class_file)
                fd = open(class_file, 'w')
                pickle.dump(p, fd)
                fd.close()
            except Exception as e:
                ex = AgentException(AgentException.E_CONFIG_COMMIT_FAILED,
                    detail=e)
                self.logger.exception(ex.message)
                return HttpErrorResponse(ex.message)
            else:
                self.logger.debug('Created class file')
                return HttpJsonResponse()

    def _stop(self, kwargs, class_file, pClass):
        self.logger.debug("_stop(kwargs=%s, class_file=%s, pClass=%s)" %
                (kwargs, class_file, pClass))

        if not exists(class_file):
            self.logger.error("class_file '%s' does not exist" % class_file)

            return HttpErrorResponse(AgentException(
                AgentException.E_CONFIG_NOT_EXIST).message)

        try:
            try:
                fd = open(class_file, 'r')
                p = pickle.load(fd)
                self.logger.debug('dump file %s loaded' %class_file)
                fd.close()
            except Exception as e:
                ex = AgentException(
                    AgentException.E_CONFIG_READ_FAILED, detail=e)
                self.logger.exception(ex.message)
                return HttpErrorResponse(ex.message)

            if 'drain' in kwargs:
                p.stop(kwargs['drain'])
            else:
                p.stop()

            self.logger.debug("Removing class_file '%s'" % class_file)
            remove(class_file)
            return HttpJsonResponse()

        except Exception as e:
            ex = AgentException(AgentException.E_UNKNOWN, detail=e)
            self.logger.exception(e)
            return HttpErrorResponse(ex.message)

    @expose('POST')
    def createDIR(self, kwargs):
        try:
            kwargs = self._DIR_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        with self.dir_lock:
            return self._create(kwargs, self.dir_file, self.DIR)

    @expose('POST')
    def createMRC(self, kwargs):
        try:
            kwargs = self._MRC_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        else:
            with self.mrc_lock:
                return self._create(kwargs, self.mrc_file, self.MRC)

    @expose('POST')
    def createOSD(self, kwargs):
        try:
            kwargs = self._OSD_get_params(kwargs)
        except AgentException as e:
            return HttpErrorResponse(e.message)
        else:
            with self.osd_lock:
                return self._create(kwargs, self.osd_file, self.OSD)

    @expose('POST')
    def stopDIR(self, kwargs):
        """Kill the DIR service"""
        with self.dir_lock:
            return self._stop(kwargs, self.dir_file, self.DIR)

    @expose('POST')
    def stopMRC(self, kwargs):
        """Kill the MRC service"""
        with self.mrc_lock:
            return self._stop(kwargs, self.mrc_file, self.MRC)

    @expose('POST')
    def stopOSD(self, kwargs):
        """Kill the OSD service"""
        with self.osd_lock:
            return self._stop(kwargs, self.osd_file, self.OSD)

    def _DIR_get_params(self, kwargs):
        ret = {}
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        return ret

    def _MRC_get_params(self, kwargs):
        ret = {}
        if 'dir_serviceHost' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'dir service host')
        ret['dir_serviceHost'] = kwargs.pop('dir_serviceHost')
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        return ret

    def _OSD_get_params(self, kwargs):
        ret = {}
        if 'dir_serviceHost' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'dir service host')
        ret['dir_serviceHost'] = kwargs.pop('dir_serviceHost')
        if 'uuid' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'uuid')
        ret['uuid'] = kwargs.pop('uuid')
        ret['mkfs'] = kwargs['mkfs']
        if 'device_name' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'device_name')
        ret['device_name'] = kwargs.pop('device_name')
        return ret

    @expose('POST')
    def startup(self, kwargs):
        self.state = 'RUNNING'
        self.logger.info('Agent started up')
        return HttpJsonResponse()

    @expose('POST')
    def set_snapshot(self, kwargs):
        if 'archive_url' not in kwargs:
            raise AgentException(
                AgentException.E_ARGS_MISSING, 'archive_url')

        self.logger.info('set_snapshot: restoring archive')

        archive_cmd = "wget --ca-certificate /etc/cpsagent/certs/ca_cert.pem "
        archive_cmd += kwargs['archive_url'] + " -O - | tar xz -C /"

        self.logger.debug('set_snapshot: %s' % archive_cmd)

        out, err = run_cmd(archive_cmd)

        self.logger.debug('set_snapshot: stdout %s' % out)
        self.logger.debug('set_snapshot: stderr %s' % err)

        self.logger.info('set_snapshot: archive restored successfully')
        return HttpJsonResponse()

    @expose('POST')
    def get_snapshot(self, kwargs):
        # pack data from:
        #     /etc/cpsagent/certs/ (original agent certificates)
        #     /etc/xos/xtreemfs/ (config files, SSL-certificates, policies)
        #     /var/lib/xtreemfs/ (save everything, after shutting down the services)
        #     /var/log/xtreemfs/ (xtreemfs log files)
        filename = "/root/snapshot.tar.gz"
        dirs = "var/lib/xtreemfs/ etc/xos/xtreemfs/ var/log/xtreemfs/ etc/cpsagent/certs"

        err, out = run_cmd("tar -czf %s %s" % (filename, dirs), "/")
        if err:
            self.logger.exception(err)
            return HttpErrorResponse(err)

        return HttpFileDownloadResponse("snapshot.tar.gz", filename)

    @expose('POST')
    def set_certificates(self, kwargs):
        self.logger.debug('set_snapshot called')
        # write certificates and truststore to the xtreemfs config directory
        path = "/etc/xos/xtreemfs/truststore/certs"
        if not exists(path):
            makedirs(path)
        open(path + "/dir.p12", 'wb').write(base64.b64decode(kwargs['dir']))
        open(path + "/mrc.p12", 'wb').write(base64.b64decode(kwargs['mrc']))
        open(path + "/osd.p12", 'wb').write(base64.b64decode(kwargs['osd']))
        open(path + "/trusted.jks", 'wb').write(base64.b64decode(kwargs['truststore']))
        return HttpJsonResponse()
<|endoftext|>


<commit_before>from __future__ import print_function, division, absolute_import#, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import distributed
import time
from dask import delayed, array
from rfpipe import source, search, util, candidates
from dask.base import tokenize
import numpy as np

import logging
logger = logging.getLogger(__name__)
vys_timeout_default = 10

def pipeline_scan(st, segments=None, host=None, cl=None, cfile=None,
                  vys_timeout=vys_timeout_default, throttle=False):
    """ Given rfpipe state and dask distributed client, run search pipline.
    throttle option will check workers for memory before submitting.
    """

    read_overhead = 8.

    if cl is None:
        if host is None:
            cl = distributed.Client(n_workers=1, threads_per_worker=16,
                                    resources={"READER": 1, "MEMORY": 24,
                                               "CORES": 16},
                                    local_dir="/lustre/evla/test/realfast/scratch")
        else:
            cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    futures = []
    if not isinstance(segments, list):
        segments = range(st.nsegment)

    # set up computations
    t0 = time.time()
    timeout = st.metadata.inttime*st.metadata.nints
    elapsedtime = time.time() - t0
    if throttle:
        # submit if workers ready. timeout of scan length.
        segment = 0
        while (len(futures) < len(segments)) and (elapsedtime < timeout):
            if worker_ready(cl, read_overhead*st.vismem*1e9):
                futures.append(pipeline_seg(st, segment, cl=cl,
                               cfile=cfile, vys_timeout=vys_timeout))
                segment += 1

            else:
                time.sleep(0.1)
                elapsedtime = time.time() - t0

        if elapsedtime > timeout:
            logger.info("Throttled submission timed out. {0}/{1} segments submitted."
                        .format(len(futures), len(segments)))

    else:
        for segment in segments:
            futures.append(pipeline_seg(st, segment, cl=cl, cfile=cfile,
                                        vys_timeout=vys_timeout))

    return futures  # list of dicts

def worker_ready(cl, memory_required):
    """ Use worker state to decide if new reader call can be submitted.
    memory_required is the size of the read in bytes
    """

    for vals in itervalues(cl.scheduler_info()['workers']):
        # look for at least one worker with required memory
        if (('READER' in vals['resources']) and
           (vals['memory_limit']-vals['memory'] > memory_required)):
            return True

    logger.info("No worker found with required memory of {0} GB"
                .format(memory_required/1e9))

    return False

def pipeline_seg(st, segment, cl, cfile=None,
                 vys_timeout=vys_timeout_default):
    """ Submit pipeline processing of a single segment to scheduler.
    Can use distributed client or compute locally.

    Uses distributed resources parameter to control scheduling of GPUs.
    Pipeline produces jobs per DM/dt.
    Returns a dict with values as futures of certain jobs (data, collection).
    """

    logger.info('Building dask for observation {0}, scan {1}, segment {2}.'
                .format(st.metadata.datasetId, st.metadata.scan, segment))

    futures = {}

# new style read *TODO: note hack on resources*
    data = lazy_read_segment(st, segment, cfile, vys_timeout)
    data = cl.compute(data, resources={tuple(data.__dask_keys__()[0][0][0]):
                                       {'READER': 1}})  # get future
# old style read
#    data = cl.submit(source.read_segment, st, segment, cfile, vys_timeout,
#                     resources={'READER': 1})

    futures['data'] = data  # save future

    data_prep = cl.submit(source.data_prep, st, segment, data,
                          resources={'CORES': st.prefs.nthread})

    saved = []
    if st.fftmode == "fftw":
        searchresources = {'MEMORY': 2*st.immem+2*st.vismem,
                           'CORES': st.prefs.nthread}
        imgranges = [[(min(st.get_search_ints(segment, dmind, dtind)),
                     max(st.get_search_ints(segment, dmind, dtind)))
                      for dtind in range(len(st.dtarr))]
                     for dmind in range(len(st.dmarr))]
        wisdom = cl.submit(search.set_wisdom, st.npixx, st.npixy,
                           resources={'CORES': 1})

        for dmind in range(len(st.dmarr)):
            delay = cl.submit(util.calc_delay, st.freq, st.freq.max(),
                              st.dmarr[dmind], st.inttime,
                              resources={'CORES': 1})
            for dtind in range(len(st.dtarr)):
                data_corr = cl.submit(search.dedisperseresample, data_prep,
                                      delay, st.dtarr[dtind],
                                      parallel=st.prefs.nthread > 1,
                                      resources={'MEMORY': 2*st.vismem,
                                                 'CORES': st.prefs.nthread})

                im0, im1 = imgranges[dmind][dtind]
                integrationlist = [list(range(im0, im1)[i:i+st.chunksize])
                                   for i in range(0, im1-im0, st.chunksize)]
                for integrations in integrationlist:
                    saved.append(cl.submit(search.search_thresh_fftw, st,
                                           segment, data_corr, dmind, dtind,
                                           integrations=integrations,
                                           wisdom=wisdom,
                                           resources=searchresources))

    elif st.fftmode == "cuda":
        for dmind in range(len(st.dmarr)):
            saved.append(cl.submit(search.dedisperse_image_cuda, st, segment,
                                   data_prep, dmind,
                                   resources={'GPU': 1,
                                              'CORES': st.prefs.nthread}))

    canddatalist = cl.submit(mergelists, saved,
                             resources={'CORES': 1})
    candcollection = cl.submit(candidates.calc_features, canddatalist,
                               resources={'CORES': 1})
    futures['candcollection'] = candcollection

    return futures

def pipeline_seg_delayed(st, segment, host=None, cl=None, cfile=None,
                         vys_timeout=vys_timeout_default):
    """ As above, but uses dask delayed to compose graph before computing.
    """

    logger.info('Building dask for observation {0}, scan {1}, segment {2}.'
                .format(st.metadata.datasetId, st.metadata.scan, segment))

    if cl is None:
        if host is None:
            cl = distributed.Client(n_workers=1, threads_per_worker=16,
                                    local_dir="/lustre/evla/test/realfast/scratch")
        else:
            cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    imgranges = [[(min(st.get_search_ints(segment, dmind, dtind)),
                  max(st.get_search_ints(segment, dmind, dtind)))
                  for dtind in range(len(st.dtarr))]
                 for dmind in range(len(st.dmarr))]

    futures = {}

    # plan, if using fftw
    wisdom = delayed(search.set_wisdom, pure=True)(st.npixx, st.npixy)

    # will retry to get around thread collision during read (?)
    data = delayed(source.read_segment, pure=True)(st, segment,
                                                   timeout=vys_timeout,
                                                   cfile=cfile)

    data_prep = delayed(source.data_prep, pure=True)(st, data)

    futures['data'] = cl.compute(data_prep)

    saved = []
    for dmind in range(len(st.dmarr)):
        delay = delayed(util.calc_delay, pure=True)(st.freq, st.freq.max(),
                                                    st.dmarr[dmind], st.inttime)

        for dtind in range(len(st.dtarr)):
            data_corr = delayed(search.dedisperseresample, pure=True)(data_prep, delay,
                                                                      st.dtarr[dtind],
                                                                      parallel = st.prefs.nthread > 1)

            im0, im1 = imgranges[dmind][dtind]
            integrationlist = [list(range(im0, im1)[i:i+st.chunksize])
                               for i in range(0, im1-im0, st.chunksize)]
            for integrations in integrationlist:
                saved.append(delayed(search.search_thresh,
                                     pure=True)(st, segment, data_corr, dmind,
                                                dtind,
                                                integrations=integrations,
                                                wisdom=wisdom))

    canddatalist = delayed(mergelists, pure=True)(saved)
    candcollection = delayed(candidates.calc_features, pure=True)(canddatalist)

    futures['candcollection'] = cl.compute(candcollection)

    return futures

def mergelists(futlists):
    """ Take list of lists and return single list
    ** TODO: could put logic here to find islands, peaks, etc?
    """

    return [fut for futlist in futlists for fut in futlist]

def lazy_read_segment(st, segment, cfile=None,
                      timeout=vys_timeout_default):
    """ rfpipe read_segment as a dask array.
    equivalent to making delayed version of function and then:
    arr = dask.array.from_delayed(dd, st.datashape, np.complex64).
    """

    shape = st.datashape
    chunks = ((shape[0],), (shape[1],), (shape[2],), (shape[3],))

    name = 'read_segment-' + tokenize([st, segment])
    dask = {(name, 0, 0, 0, 0): (source.read_segment, st, segment,
                                 cfile, timeout)}

    return array.Array(dask=dask, name=name, chunks=chunks, dtype=np.complex64)
<commit_msg>added placeholder function for further throttling based on disk usage<commit_after>from __future__ import print_function, division, absolute_import#, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import distributed
import time
from dask import delayed, array
from rfpipe import source, search, util, candidates
from dask.base import tokenize
import numpy as np

import logging
logger = logging.getLogger(__name__)
vys_timeout_default = 10

def pipeline_scan(st, segments=None, host=None, cl=None, cfile=None,
                  vys_timeout=vys_timeout_default, throttle=False):
    """ Given rfpipe state and dask distributed client, run search pipline.
    throttle option will check workers for memory before submitting.
    """

    read_overhead = 8.

    if cl is None:
        if host is None:
            cl = distributed.Client(n_workers=1, threads_per_worker=16,
                                    resources={"READER": 1, "MEMORY": 24,
                                               "CORES": 16},
                                    local_dir="/lustre/evla/test/realfast/scratch")
        else:
            cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    futures = []
    if not isinstance(segments, list):
        segments = range(st.nsegment)

    # set up computations
    t0 = time.time()
    timeout = st.metadata.inttime*st.metadata.nints
    elapsedtime = time.time() - t0
    if throttle:
        # submit if workers ready. timeout of scan length.
        segment = 0
        while (len(futures) < len(segments)) and (elapsedtime < timeout):
            if (worker_ready(cl, read_overhead*st.vismem*1e9) and disk_ready):
                futures.append(pipeline_seg(st, segment, cl=cl,
                               cfile=cfile, vys_timeout=vys_timeout))
                segment += 1

            else:
                time.sleep(0.1)
                elapsedtime = time.time() - t0

        if elapsedtime > timeout:
            logger.info("Throttled submission timed out. {0}/{1} segments submitted."
                        .format(len(futures), len(segments)))

    else:
        for segment in segments:
            futures.append(pipeline_seg(st, segment, cl=cl, cfile=cfile,
                                        vys_timeout=vys_timeout))

    return futures  # list of dicts

def worker_ready(cl, memory_required):
    """ Use worker state to decide if new reader call can be submitted.
    memory_required is the size of the read in bytes
    """

    for vals in itervalues(cl.scheduler_info()['workers']):
        # look for at least one worker with required memory
        if (('READER' in vals['resources']) and
           (vals['memory_limit']-vals['memory'] > memory_required)):
            return True

    logger.info("No worker found with required memory of {0} GB"
                .format(memory_required/1e9))

    return False

def disk_ready(cl):
  """ Is disk usage low enough to be able to handle spilled data?
  """

  return True

def pipeline_seg(st, segment, cl, cfile=None,
                 vys_timeout=vys_timeout_default):
    """ Submit pipeline processing of a single segment to scheduler.
    Can use distributed client or compute locally.

    Uses distributed resources parameter to control scheduling of GPUs.
    Pipeline produces jobs per DM/dt.
    Returns a dict with values as futures of certain jobs (data, collection).
    """

    logger.info('Building dask for observation {0}, scan {1}, segment {2}.'
                .format(st.metadata.datasetId, st.metadata.scan, segment))

    futures = {}

# new style read *TODO: note hack on resources*
    data = lazy_read_segment(st, segment, cfile, vys_timeout)
    data = cl.compute(data, resources={tuple(data.__dask_keys__()[0][0][0]):
                                       {'READER': 1}})  # get future
# old style read
#    data = cl.submit(source.read_segment, st, segment, cfile, vys_timeout,
#                     resources={'READER': 1})

    futures['data'] = data  # save future

    data_prep = cl.submit(source.data_prep, st, segment, data,
                          resources={'CORES': st.prefs.nthread})

    saved = []
    if st.fftmode == "fftw":
        searchresources = {'MEMORY': 2*st.immem+2*st.vismem,
                           'CORES': st.prefs.nthread}
        imgranges = [[(min(st.get_search_ints(segment, dmind, dtind)),
                     max(st.get_search_ints(segment, dmind, dtind)))
                      for dtind in range(len(st.dtarr))]
                     for dmind in range(len(st.dmarr))]
        wisdom = cl.submit(search.set_wisdom, st.npixx, st.npixy,
                           resources={'CORES': 1})

        for dmind in range(len(st.dmarr)):
            delay = cl.submit(util.calc_delay, st.freq, st.freq.max(),
                              st.dmarr[dmind], st.inttime,
                              resources={'CORES': 1})
            for dtind in range(len(st.dtarr)):
                data_corr = cl.submit(search.dedisperseresample, data_prep,
                                      delay, st.dtarr[dtind],
                                      parallel=st.prefs.nthread > 1,
                                      resources={'MEMORY': 2*st.vismem,
                                                 'CORES': st.prefs.nthread})

                im0, im1 = imgranges[dmind][dtind]
                integrationlist = [list(range(im0, im1)[i:i+st.chunksize])
                                   for i in range(0, im1-im0, st.chunksize)]
                for integrations in integrationlist:
                    saved.append(cl.submit(search.search_thresh_fftw, st,
                                           segment, data_corr, dmind, dtind,
                                           integrations=integrations,
                                           wisdom=wisdom,
                                           resources=searchresources))

    elif st.fftmode == "cuda":
        for dmind in range(len(st.dmarr)):
            saved.append(cl.submit(search.dedisperse_image_cuda, st, segment,
                                   data_prep, dmind,
                                   resources={'GPU': 1,
                                              'CORES': st.prefs.nthread}))

    canddatalist = cl.submit(mergelists, saved,
                             resources={'CORES': 1})
    candcollection = cl.submit(candidates.calc_features, canddatalist,
                               resources={'CORES': 1})
    futures['candcollection'] = candcollection

    return futures

def pipeline_seg_delayed(st, segment, host=None, cl=None, cfile=None,
                         vys_timeout=vys_timeout_default):
    """ As above, but uses dask delayed to compose graph before computing.
    """

    logger.info('Building dask for observation {0}, scan {1}, segment {2}.'
                .format(st.metadata.datasetId, st.metadata.scan, segment))

    if cl is None:
        if host is None:
            cl = distributed.Client(n_workers=1, threads_per_worker=16,
                                    local_dir="/lustre/evla/test/realfast/scratch")
        else:
            cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    imgranges = [[(min(st.get_search_ints(segment, dmind, dtind)),
                  max(st.get_search_ints(segment, dmind, dtind)))
                  for dtind in range(len(st.dtarr))]
                 for dmind in range(len(st.dmarr))]

    futures = {}

    # plan, if using fftw
    wisdom = delayed(search.set_wisdom, pure=True)(st.npixx, st.npixy)

    # will retry to get around thread collision during read (?)
    data = delayed(source.read_segment, pure=True)(st, segment,
                                                   timeout=vys_timeout,
                                                   cfile=cfile)

    data_prep = delayed(source.data_prep, pure=True)(st, data)

    futures['data'] = cl.compute(data_prep)

    saved = []
    for dmind in range(len(st.dmarr)):
        delay = delayed(util.calc_delay, pure=True)(st.freq, st.freq.max(),
                                                    st.dmarr[dmind], st.inttime)

        for dtind in range(len(st.dtarr)):
            data_corr = delayed(search.dedisperseresample, pure=True)(data_prep, delay,
                                                                      st.dtarr[dtind],
                                                                      parallel = st.prefs.nthread > 1)

            im0, im1 = imgranges[dmind][dtind]
            integrationlist = [list(range(im0, im1)[i:i+st.chunksize])
                               for i in range(0, im1-im0, st.chunksize)]
            for integrations in integrationlist:
                saved.append(delayed(search.search_thresh,
                                     pure=True)(st, segment, data_corr, dmind,
                                                dtind,
                                                integrations=integrations,
                                                wisdom=wisdom))

    canddatalist = delayed(mergelists, pure=True)(saved)
    candcollection = delayed(candidates.calc_features, pure=True)(canddatalist)

    futures['candcollection'] = cl.compute(candcollection)

    return futures

def mergelists(futlists):
    """ Take list of lists and return single list
    ** TODO: could put logic here to find islands, peaks, etc?
    """

    return [fut for futlist in futlists for fut in futlist]

def lazy_read_segment(st, segment, cfile=None,
                      timeout=vys_timeout_default):
    """ rfpipe read_segment as a dask array.
    equivalent to making delayed version of function and then:
    arr = dask.array.from_delayed(dd, st.datashape, np.complex64).
    """

    shape = st.datashape
    chunks = ((shape[0],), (shape[1],), (shape[2],), (shape[3],))

    name = 'read_segment-' + tokenize([st, segment])
    dask = {(name, 0, 0, 0, 0): (source.read_segment, st, segment,
                                 cfile, timeout)}

    return array.Array(dask=dask, name=name, chunks=chunks, dtype=np.complex64)
<|endoftext|>


<commit_before>from discord.ext import commands
from collections import Counter

from .utils import checks

import logging
import discord
import datetime
import psutil
import os

log = logging.getLogger()

LOGGING_CHANNEL = '309632009427222529'

class Stats:
    """Bot usage statistics."""

    def __init__(self, bot):
        self.bot = bot
        self.process = psutil.Process()

    async def on_command(self, command, ctx):
        self.bot.commands_used[ctx.command.qualified_name] += 1
        message = ctx.message
        destination = None
        if message.channel.is_private:
            destination = 'Private Message'
        else:
            destination = '#{0.channel.name} ({0.server.name})'.format(message)

        log.info('{0.timestamp}: {0.author.name} in {1}: {0.content}'.format(message, destination))

    async def on_socket_response(self, msg):
        self.bot.socket_stats[msg.get('t')] += 1

    @commands.command(hidden=True)
    @checks.is_owner()
    async def commandstats(self):
        p = commands.Paginator()
        counter = self.bot.commands_used
        width = len(max(counter, key=len))
        total = sum(counter.values())

        fmt = '{0:<{width}}: {1}'
        p.add_line(fmt.format('Total', total, width=width))
        for key, count in counter.most_common():
            p.add_line(fmt.format(key, count, width=width))

        for page in p.pages:
            await self.bot.say(page)

    @commands.command(hidden=True)
    async def socketstats(self):
        delta = datetime.datetime.utcnow() - self.bot.uptime
        minutes = delta.total_seconds() / 60
        total = sum(self.bot.socket_stats.values())
        cpm = total / minutes

        fmt = '%s socket events observed (%.2f/minute):\n%s'
        await self.bot.say(fmt % (total, cpm, self.bot.socket_stats))

    def get_bot_uptime(self, *, brief=False):
        now = datetime.datetime.utcnow()
        delta = now - self.bot.uptime
        hours, remainder = divmod(int(delta.total_seconds()), 3600)
        minutes, seconds = divmod(remainder, 60)
        days, hours = divmod(hours, 24)

        if not brief:
            if days:
                fmt = '{d} days, {h} hours, {m} minutes, and {s} seconds'
            else:
                fmt = '{h} hours, {m} minutes, and {s} seconds'
        else:
            fmt = '{h}h {m}m {s}s'
            if days:
                fmt = '{d}d ' + fmt

        return fmt.format(d=days, h=hours, m=minutes, s=seconds)

    @commands.command()
    async def uptime(self):
        """Tells you how long the bot has been up for."""
        await self.bot.say('Uptime: **{}**'.format(self.get_bot_uptime()))

    @commands.command(aliases=['stats'])
    async def about(self):
        """Tells you information about the bot itself."""
        cmd = r'git show -s HEAD~3..HEAD --format="[{}](https://github.com/Rapptz/RoboDanny/commit/%H) %s (%cr)"'
        if os.name == 'posix':
            cmd = cmd.format(r'\`%h\`')
        else:
            cmd = cmd.format(r'`%h`')

        revision = os.popen(cmd).read().strip()
        embed = discord.Embed(description='Latest Changes:\n' + revision)
        embed.title = 'Official Bot Server Invite'
        embed.url = 'https://discord.gg/0118rJdtd1rVJJfuI'
        embed.colour = 0x738bd7 # blurple

        try:
            owner = self._owner
        except AttributeError:
            owner = self._owner = await self.bot.get_user_info('80088516616269824')

        embed.set_author(name=str(owner), icon_url=owner.avatar_url)

        # statistics
        total_members = sum(len(s.members) for s in self.bot.servers)
        total_online  = sum(1 for m in self.bot.get_all_members() if m.status != discord.Status.offline)
        unique_members = set(self.bot.get_all_members())
        unique_online = sum(1 for m in unique_members if m.status != discord.Status.offline)
        channel_types = Counter(c.type for c in self.bot.get_all_channels())
        voice = channel_types[discord.ChannelType.voice]
        text = channel_types[discord.ChannelType.text]

        members = '%s total\n%s online\n%s unique\n%s unique online' % (total_members, total_online, len(unique_members), unique_online)
        embed.add_field(name='Members', value=members)
        embed.add_field(name='Channels', value='{} total\n{} text\n{} voice'.format(text + voice, text, voice))
        memory_usage = self.process.memory_full_info().uss / 1024**2
        cpu_usage = self.process.cpu_percent() / psutil.cpu_count()
        embed.add_field(name='Process', value='{:.2f} MiB\n{:.2f}% CPU'.format(memory_usage, cpu_usage))
        embed.set_footer(text='Made with discord.py', icon_url='http://i.imgur.com/5BFecvA.png')
        embed.timestamp = self.bot.uptime

        embed.add_field(name='Servers', value=len(self.bot.servers))
        embed.add_field(name='Commands Run', value=sum(self.bot.commands_used.values()))
        embed.add_field(name='Uptime', value=self.get_bot_uptime(brief=True))

        await self.bot.say(embed=embed)

    async def send_server_stat(self, e, server):
        e.add_field(name='Name', value=server.name)
        e.add_field(name='ID', value=server.id)
        e.add_field(name='Owner', value='{0} (ID: {0.id})'.format(server.owner))

        bots = sum(m.bot for m in server.members)
        total = server.member_count
        online = sum(m.status is discord.Status.online for m in server.members)
        e.add_field(name='Members', value=str(total))
        e.add_field(name='Bots', value='{} ({:.2%})'.format(bots, bots / total))
        e.add_field(name='Online', value='{} ({:.2%})'.format(online, online / total))
        if server.icon:
            e.set_thumbnail(url=server.icon_url)

        if server.me:
            e.timestamp = server.me.joined_at

        ch = self.bot.get_channel(LOGGING_CHANNEL)
        await self.bot.send_message(ch, embed=e)

    async def on_server_join(self, server):
        e = discord.Embed(colour=0x53dda4, title='New Server') # green colour
        await self.send_server_stat(e, server)

    async def on_server_remove(self, server):
        e = discord.Embed(colour=0xdd5f53, title='Left Server') # red colour
        await self.send_server_stat(e, server)

def setup(bot):
    if not hasattr(bot, 'commands_used'):
        bot.commands_used = Counter()

    if not hasattr(bot, 'socket_stats'):
        bot.socket_stats = Counter()

    bot.add_cog(Stats(bot))
<commit_msg>Track on_error and on_command_error.<commit_after>from discord.ext import commands
from collections import Counter

from .utils import checks

import logging
import discord
import datetime
import traceback
import psutil
import os

log = logging.getLogger()

LOGGING_CHANNEL = '309632009427222529'

class Stats:
    """Bot usage statistics."""

    def __init__(self, bot):
        self.bot = bot
        self.process = psutil.Process()

    async def on_command(self, command, ctx):
        self.bot.commands_used[ctx.command.qualified_name] += 1
        message = ctx.message
        destination = None
        if message.channel.is_private:
            destination = 'Private Message'
        else:
            destination = '#{0.channel.name} ({0.server.name})'.format(message)

        log.info('{0.timestamp}: {0.author.name} in {1}: {0.content}'.format(message, destination))

    async def on_socket_response(self, msg):
        self.bot.socket_stats[msg.get('t')] += 1

    @commands.command(hidden=True)
    @checks.is_owner()
    async def commandstats(self):
        p = commands.Paginator()
        counter = self.bot.commands_used
        width = len(max(counter, key=len))
        total = sum(counter.values())

        fmt = '{0:<{width}}: {1}'
        p.add_line(fmt.format('Total', total, width=width))
        for key, count in counter.most_common():
            p.add_line(fmt.format(key, count, width=width))

        for page in p.pages:
            await self.bot.say(page)

    @commands.command(hidden=True)
    async def socketstats(self):
        delta = datetime.datetime.utcnow() - self.bot.uptime
        minutes = delta.total_seconds() / 60
        total = sum(self.bot.socket_stats.values())
        cpm = total / minutes

        fmt = '%s socket events observed (%.2f/minute):\n%s'
        await self.bot.say(fmt % (total, cpm, self.bot.socket_stats))

    def get_bot_uptime(self, *, brief=False):
        now = datetime.datetime.utcnow()
        delta = now - self.bot.uptime
        hours, remainder = divmod(int(delta.total_seconds()), 3600)
        minutes, seconds = divmod(remainder, 60)
        days, hours = divmod(hours, 24)

        if not brief:
            if days:
                fmt = '{d} days, {h} hours, {m} minutes, and {s} seconds'
            else:
                fmt = '{h} hours, {m} minutes, and {s} seconds'
        else:
            fmt = '{h}h {m}m {s}s'
            if days:
                fmt = '{d}d ' + fmt

        return fmt.format(d=days, h=hours, m=minutes, s=seconds)

    @commands.command()
    async def uptime(self):
        """Tells you how long the bot has been up for."""
        await self.bot.say('Uptime: **{}**'.format(self.get_bot_uptime()))

    @commands.command(aliases=['stats'])
    async def about(self):
        """Tells you information about the bot itself."""
        cmd = r'git show -s HEAD~3..HEAD --format="[{}](https://github.com/Rapptz/RoboDanny/commit/%H) %s (%cr)"'
        if os.name == 'posix':
            cmd = cmd.format(r'\`%h\`')
        else:
            cmd = cmd.format(r'`%h`')

        revision = os.popen(cmd).read().strip()
        embed = discord.Embed(description='Latest Changes:\n' + revision)
        embed.title = 'Official Bot Server Invite'
        embed.url = 'https://discord.gg/0118rJdtd1rVJJfuI'
        embed.colour = 0x738bd7 # blurple

        try:
            owner = self._owner
        except AttributeError:
            owner = self._owner = await self.bot.get_user_info('80088516616269824')

        embed.set_author(name=str(owner), icon_url=owner.avatar_url)

        # statistics
        total_members = sum(len(s.members) for s in self.bot.servers)
        total_online  = sum(1 for m in self.bot.get_all_members() if m.status != discord.Status.offline)
        unique_members = set(self.bot.get_all_members())
        unique_online = sum(1 for m in unique_members if m.status != discord.Status.offline)
        channel_types = Counter(c.type for c in self.bot.get_all_channels())
        voice = channel_types[discord.ChannelType.voice]
        text = channel_types[discord.ChannelType.text]

        members = '%s total\n%s online\n%s unique\n%s unique online' % (total_members, total_online, len(unique_members), unique_online)
        embed.add_field(name='Members', value=members)
        embed.add_field(name='Channels', value='{} total\n{} text\n{} voice'.format(text + voice, text, voice))
        memory_usage = self.process.memory_full_info().uss / 1024**2
        cpu_usage = self.process.cpu_percent() / psutil.cpu_count()
        embed.add_field(name='Process', value='{:.2f} MiB\n{:.2f}% CPU'.format(memory_usage, cpu_usage))
        embed.set_footer(text='Made with discord.py', icon_url='http://i.imgur.com/5BFecvA.png')
        embed.timestamp = self.bot.uptime

        embed.add_field(name='Servers', value=len(self.bot.servers))
        embed.add_field(name='Commands Run', value=sum(self.bot.commands_used.values()))
        embed.add_field(name='Uptime', value=self.get_bot_uptime(brief=True))

        await self.bot.say(embed=embed)

    async def send_server_stat(self, e, server):
        e.add_field(name='Name', value=server.name)
        e.add_field(name='ID', value=server.id)
        e.add_field(name='Owner', value='{0} (ID: {0.id})'.format(server.owner))

        bots = sum(m.bot for m in server.members)
        total = server.member_count
        online = sum(m.status is discord.Status.online for m in server.members)
        e.add_field(name='Members', value=str(total))
        e.add_field(name='Bots', value='{} ({:.2%})'.format(bots, bots / total))
        e.add_field(name='Online', value='{} ({:.2%})'.format(online, online / total))
        if server.icon:
            e.set_thumbnail(url=server.icon_url)

        if server.me:
            e.timestamp = server.me.joined_at

        ch = self.bot.get_channel(LOGGING_CHANNEL)
        await self.bot.send_message(ch, embed=e)

    async def on_server_join(self, server):
        e = discord.Embed(colour=0x53dda4, title='New Server') # green colour
        await self.send_server_stat(e, server)

    async def on_server_remove(self, server):
        e = discord.Embed(colour=0xdd5f53, title='Left Server') # red colour
        await self.send_server_stat(e, server)

    async def on_command_error(self, error, ctx):
        ignored = (commands.NoPrivateMessage, commands.DisabledCommand, commands.CheckFailure,
                   commands.CommandNotFound, discord.HTTPException)
        error = getattr(error, 'original', error)

        if isinstance(error, ignored):
            return

        e = discord.Embed(title='Command Error', colour=0xcc3366)
        e.add_field(name='Name', value=ctx.command.qualified_name)
        e.add_field(name='Author', value='{0} (ID: {0.id})'.format(ctx.message.author))

        if ctx.message.server:
            fmt = 'Channel: {0} (ID: {0.id})\nGuild: {1} (ID: {1.id})'
        else:
            fmt = 'Channel: {0} (ID: {0.id})'

        e.add_field(name='Location', value=fmt.format(ctx.message.channel, ctx.message.server), inline=False)

        e.description = '```py\n%s\n```' % traceback.format_exc(chain=False)
        e.timestamp = datetime.datetime.utcnow()
        ch = self.bot.get_channel(LOGGING_CHANNEL)
        await self.bot.send_message(ch, embed=e)

old_on_error = commands.Bot.on_error

async def on_error(self, event, *args, **kwargs):
    e = discord.Embed(title='Event Error', colour=0xa32952)
    e.add_field(name='Event', value=event)
    e.description = '```py\n%s\n```' % traceback.format_exc()
    e.timestamp = datetime.datetime.utcnow()
    ch = self.get_channel(LOGGING_CHANNEL)
    try:
        await self.send_message(ch, embed=e)
    except:
        pass

def setup(bot):
    if not hasattr(bot, 'commands_used'):
        bot.commands_used = Counter()

    if not hasattr(bot, 'socket_stats'):
        bot.socket_stats = Counter()

    bot.add_cog(Stats(bot))
    commands.Bot.on_error = on_error

def teardown(bot):
    commands.Bot.on_error = old_on_error
<|endoftext|>