<commit_before>import numpy as np

from sisl import geom, Atom, Geometry

__all__ = ['nanoribbon', 'graphene_nanoribbon', 'agnr', 'zgnr']

def nanoribbon(bond, atom, width, kind='armchair'):
    r""" Construction of a nanoribbon unit cell of type armchair or zigzag.
    The geometry is oriented along the x-axis.

    Parameters
    ----------
    bond : float
       bond length between atoms in the honeycomb lattice
    atom : Atom
       atom (or atoms) in the honeycomb lattice
    width : int
       number of atoms in the transverse direction
    kind : {'armchair', 'zigzag'}
       type of ribbon

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    graphene_nanoribbon : graphene nanoribbon
    agnr : armchair graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    if atom is None:
        atom = Atom(Z=6, R=bond * 1.01)

    # Width characterization
    if not isinstance(width, int):
        raise ValueError("nanoribbon: the width needs to be a postive integer!")
    width = max(width, 1)
    n, m = width // 2, width % 2

    ribbon = geom.honeycomb(bond, atom, orthogonal=True)

    kind = kind.lower()
    if kind.startswith('a'):
        # Construct armchair GNR
        if m == 1:
            ribbon = ribbon.repeat(n + 1, 1)
            ribbon = ribbon.remove(3 * (n + 1)).remove(0)
        else:
            ribbon = ribbon.repeat(n, 1)
    elif kind.startswith('z'):
        # Construct zigzag GNR
        ribbon = ribbon.rotate(90, [0, 0, -1])
        if m == 1:
            ribbon = ribbon.tile(n + 1, 0)
            ribbon = ribbon.remove(-1).remove(-1)
        else:
            ribbon = ribbon.tile(n, 0)
        # Invert y-coordinates
        ribbon.xyz[:, 1] *= -1
        # Set lattice vectors strictly orthogonal
        ribbon.cell[:] = np.diag([ribbon.cell[1, 0], -ribbon.cell[0, 1], ribbon.cell[2, 2]])
        # Sort along x-axis
        ribbon = ribbon.sort()
    else:
        raise ValueError("nanoribbon: kind must be armchair or zigzag")

    # Separate ribbons along y-axis
    ribbon.cell[1, 1] += 10.

    # Movie inside unit cell
    ribbon = ribbon.move([-np.min(ribbon.xyz[:, 0]), -np.min(ribbon.xyz[:, 1]), 0])

    return ribbon

def graphene_nanoribbon(width, bond=1.42, atom=None, kind='armchair'):
    r""" Construction of a graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom
       atom (or atoms) in the honeycomb lattice
    kind : {'armchair', 'zigzag'}
       type of ribbon

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    agnr : armchair graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    return nanoribbon(bond, atom, width, kind=kind)

def agnr(width, bond=1.42, atom=None):
    r""" Construction of an armchair graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom
       atom (or atoms) in the honeycomb lattice

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    graphene_nanoribbon : graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    return graphene_nanoribbon(width, bond, atom, kind='armchair')

def zgnr(width, bond=1.42, atom=None):
    r""" Construction of a zigzag graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom
       atom (or atoms) in the honeycomb lattice

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    graphene_nanoribbon : graphene nanoribbon
    agnr : armchair graphene nanoribbon
    """
    return graphene_nanoribbon(width, bond, atom, kind='zigzag')
<commit_msg>maint: atom defaults for nanoribbon geometry<commit_after>import numpy as np

from sisl import geom, Atom, Geometry

__all__ = ['nanoribbon', 'graphene_nanoribbon', 'agnr', 'zgnr']

def nanoribbon(bond, atom, width, kind='armchair'):
    r""" Construction of a nanoribbon unit cell of type armchair or zigzag.
    The geometry is oriented along the x-axis.

    Parameters
    ----------
    bond : float
       bond length between atoms in the honeycomb lattice
    atom : Atom
       atom (or atoms) in the honeycomb lattice
    width : int
       number of atoms in the transverse direction
    kind : {'armchair', 'zigzag'}
       type of ribbon

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    graphene_nanoribbon : graphene nanoribbon
    agnr : armchair graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    if not isinstance(bond, float):
        raise ValueError("nanoribbon: bond needs to be a float!")
    if not isinstance(atom, (Atom, list, tuple)):
        raise ValueError("nanoribbon: atom needs to be an instance of Atom (or list of Atoms)!")
    if not isinstance(width, int):
        raise ValueError("nanoribbon: the width needs to be a postive integer!")
    # Width characterization
    width = max(width, 1)
    n, m = width // 2, width % 2

    ribbon = geom.honeycomb(bond, atom, orthogonal=True)

    kind = kind.lower()
    if kind.startswith('a'):
        # Construct armchair GNR
        if m == 1:
            ribbon = ribbon.repeat(n + 1, 1)
            ribbon = ribbon.remove(3 * (n + 1)).remove(0)
        else:
            ribbon = ribbon.repeat(n, 1)
    elif kind.startswith('z'):
        # Construct zigzag GNR
        ribbon = ribbon.rotate(90, [0, 0, -1])
        if m == 1:
            ribbon = ribbon.tile(n + 1, 0)
            ribbon = ribbon.remove(-1).remove(-1)
        else:
            ribbon = ribbon.tile(n, 0)
        # Invert y-coordinates
        ribbon.xyz[:, 1] *= -1
        # Set lattice vectors strictly orthogonal
        ribbon.cell[:] = np.diag([ribbon.cell[1, 0], -ribbon.cell[0, 1], ribbon.cell[2, 2]])
        # Sort along x-axis
        ribbon = ribbon.sort()
    else:
        raise ValueError("nanoribbon: kind must be armchair or zigzag")

    # Separate ribbons along y-axis
    ribbon.cell[1, 1] += 10.

    # Movie inside unit cell
    ribbon = ribbon.move([-np.min(ribbon.xyz[:, 0]), -np.min(ribbon.xyz[:, 1]), 0])

    return ribbon

def graphene_nanoribbon(width, bond=1.42, atom=None, kind='armchair'):
    r""" Construction of a graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom, optional
       atom (or atoms) in the honeycomb lattice. Defaults to `Atom(6)`
    kind : {'armchair', 'zigzag'}
       type of ribbon

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    agnr : armchair graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    if atom is None:
        atom = Atom(Z=6, R=bond * 1.01)
    return nanoribbon(bond, atom, width, kind=kind)

def agnr(width, bond=1.42, atom=None):
    r""" Construction of an armchair graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom, optional
       atom (or atoms) in the honeycomb lattice. Defaults to `Atom(6)`

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    graphene_nanoribbon : graphene nanoribbon
    zgnr : zigzag graphene nanoribbon
    """
    return graphene_nanoribbon(width, bond, atom, kind='armchair')

def zgnr(width, bond=1.42, atom=None):
    r""" Construction of a zigzag graphene nanoribbon

    Parameters
    ----------
    width : int
       number of atoms in the transverse direction
    bond : float, optional
       CC bond length. Defaults to 1.42
    atom : Atom, optional
       atom (or atoms) in the honeycomb lattice. Defaults to `Atom(6)`

    See Also
    --------
    honeycomb : honeycomb lattices
    graphene : graphene geometry
    nanoribbon : honeycomb nanoribbon
    graphene_nanoribbon : graphene nanoribbon
    agnr : armchair graphene nanoribbon
    """
    return graphene_nanoribbon(width, bond, atom, kind='zigzag')
<|endoftext|>


<commit_before>import collections
import re
from os import path
from xml.etree import ElementTree

import requests

from provider import get_logger, Provider, ProviderException, Status

logger = get_logger('slf')

class Slf(Provider):
    provider_code = 'slf'
    provider_name = 'slf.ch'
    provider_url = 'http://www.slf.ch'

    provider_urls = {
        'default': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_EN',
        'en': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_EN',
        'de': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_DE',
        'fr': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_FR',
        'it': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_IT'
    }

    station_metadata_name_regexp = r'(.{3}-[0-9]{1}) : (.*? - .*?) \(IMIS\)'
    station_name_regexp = r'(.*?)([0-9]{3,4}) m'

    Measure = collections.namedtuple(
        'Measure', ['key', 'wind_direction', 'wind_average', 'wind_maximum', 'temperature'])

    def parse_data(self, line) -> Measure:
        values = line.split(';')
        return self.Measure(values[0], values[7], values[5], values[6], values[3])

    def has_wind_data(self, data):
        if not data:
            return False
        measure = self.parse_data(data[-1])
        if measure.key and measure.wind_average and measure.wind_maximum:
            return True

    def process_data(self):
        try:
            logger.info("Processing SLF data...")

            slf_metadata = {}
            with open(path.join(path.dirname(__file__), 'slf/SLF Messtationen Standorte.kml')) as kml_file:
                tree = ElementTree.parse(kml_file)
            ns = {'gis': 'http://www.opengis.net/kml/2.2'}
            for placemark in tree.getroot().findall('.//gis:Placemark', namespaces=ns):
                id, name = re.search(
                    self.station_metadata_name_regexp, placemark.find('gis:name', namespaces=ns).text).groups()
                lon, lat = placemark.find('gis:Point/gis:coordinates', namespaces=ns).text.split(',')
                slf_metadata[id.replace('-', '')] = {
                    'name': name,
                    'lat': float(lat),
                    'lon': float(lon),
                }

            result = requests.get("http://odb.slf.ch/odb/api/v1/stations",
                                  timeout=(self.connect_timeout, self.read_timeout))

            try:
                slf_stations = result.json()
            except:
                raise Exception("Unable to get SLF station list")

            for slf_station in slf_stations:
                station_id = None
                try:
                    slf_id = slf_station['id']
                    result = requests.get("http://odb.slf.ch/odb/api/v1/measurement?id={id}".format(id=slf_id),
                                          timeout=(self.connect_timeout, self.read_timeout))
                    data = result.json()
                    if not self.has_wind_data(data):
                        continue

                    name, altitude = re.search(self.station_name_regexp, slf_station['name']).groups()
                    metadata_name, lat, lon = None, None, None
                    if slf_id in slf_metadata:
                        metadata_name = slf_metadata[slf_id]['name']
                        lat = slf_metadata[slf_id]['lat']
                        lon = slf_metadata[slf_id]['lon']
                    else:
                        logger.warn('No metadata found for station {id}/{name}'.format(id=slf_id, name=name))

                    station = self.save_station(
                        slf_id,
                        metadata_name or name,
                        metadata_name,
                        lat,
                        lon,
                        Status.GREEN,
                        altitude=altitude,
                        url=self.provider_urls)
                    station_id = station['_id']

                    measures = [self.parse_data(line) for line in data]

                    measures_collection = self.measures_collection(station_id)
                    new_measures = []
                    for slf_measure in measures:
                        key = int(slf_measure.key)
                        if not self.has_measure(measures_collection, key):
                            try:
                                measure = self.create_measure(
                                    key,
                                    slf_measure.wind_direction,
                                    slf_measure.wind_average,
                                    slf_measure.wind_maximum,
                                    temperature=slf_measure.temperature,
                                )
                                new_measures.append(measure)
                            except ProviderException as e:
                                logger.warn("Error while processing measure '{0}' for station '{1}': {2}"
                                            .format(key, station_id, e))
                            except Exception as e:
                                logger.exception("Error while processing measure '{0}' for station '{1}': {2}"
                                                 .format(key, station_id, e))
                                self.raven_client.captureException()

                    self.insert_new_measures(measures_collection, station, new_measures, logger)

                except ProviderException as e:
                    logger.warn("Error while processing station '{0}': {1}".format(station_id, e))
                except Exception as e:
                    logger.exception("Error while processing station '{0}': {1}".format(station_id, e))
                    self.raven_client.captureException()

        except Exception as e:
            logger.exception("Error while processing SLF: {0}".format(e))
            self.raven_client.captureException()

        logger.info("Done !")

Slf().process_data()
<commit_msg>Fixed exception<commit_after>import collections
import re
from os import path
from xml.etree import ElementTree

import requests

from provider import get_logger, Provider, ProviderException, Status

logger = get_logger('slf')

class Slf(Provider):
    provider_code = 'slf'
    provider_name = 'slf.ch'
    provider_url = 'http://www.slf.ch'

    provider_urls = {
        'default': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_EN',
        'en': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_EN',
        'de': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_DE',
        'fr': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_FR',
        'it': 'http://www.slf.ch/schneeinfo/messwerte/wt-daten/index_IT'
    }

    station_metadata_name_regexp = r'(.{3}-[0-9]{1}) : (.*? - .*?) \(IMIS\)'
    station_name_regexp = r'(.*?)([0-9]{3,4}) m'

    Measure = collections.namedtuple(
        'Measure', ['key', 'wind_direction', 'wind_average', 'wind_maximum', 'temperature'])

    def parse_data(self, line) -> Measure:
        values = line.split(';')
        return self.Measure(values[0], values[7], values[5], values[6], values[3])

    def has_wind_data(self, data):
        if not data:
            return False
        measure = self.parse_data(data[-1])
        if measure.key and measure.wind_average and measure.wind_maximum:
            return True

    def process_data(self):
        try:
            logger.info('Processing SLF data...')

            slf_metadata = {}
            with open(path.join(path.dirname(__file__), 'slf/SLF Messtationen Standorte.kml')) as kml_file:
                tree = ElementTree.parse(kml_file)
            ns = {'gis': 'http://www.opengis.net/kml/2.2'}
            for placemark in tree.getroot().findall('.//gis:Placemark', namespaces=ns):
                id, name = re.search(
                    self.station_metadata_name_regexp, placemark.find('gis:name', namespaces=ns).text).groups()
                lon, lat = placemark.find('gis:Point/gis:coordinates', namespaces=ns).text.split(',')
                slf_metadata[id.replace('-', '')] = {
                    'name': name,
                    'lat': float(lat),
                    'lon': float(lon),
                }

            result = requests.get('http://odb.slf.ch/odb/api/v1/stations',
                                  timeout=(self.connect_timeout, self.read_timeout))
            slf_stations = result.json()

            for slf_station in slf_stations:
                station_id = None
                try:
                    slf_id = slf_station['id']
                    result = requests.get("http://odb.slf.ch/odb/api/v1/measurement?id={id}".format(id=slf_id),
                                          timeout=(self.connect_timeout, self.read_timeout))
                    data = result.json()
                    if not self.has_wind_data(data):
                        continue

                    name, altitude = re.search(self.station_name_regexp, slf_station['name']).groups()
                    metadata_name, lat, lon = None, None, None
                    if slf_id in slf_metadata:
                        metadata_name = slf_metadata[slf_id]['name']
                        lat = slf_metadata[slf_id]['lat']
                        lon = slf_metadata[slf_id]['lon']
                    else:
                        logger.warn('No metadata found for station {id}/{name}'.format(id=slf_id, name=name))

                    station = self.save_station(
                        slf_id,
                        metadata_name or name,
                        metadata_name,
                        lat,
                        lon,
                        Status.GREEN,
                        altitude=altitude,
                        url=self.provider_urls)
                    station_id = station['_id']

                    measures = [self.parse_data(line) for line in data]

                    measures_collection = self.measures_collection(station_id)
                    new_measures = []
                    for slf_measure in measures:
                        key = int(slf_measure.key)
                        if not self.has_measure(measures_collection, key):
                            try:
                                measure = self.create_measure(
                                    key,
                                    slf_measure.wind_direction,
                                    slf_measure.wind_average,
                                    slf_measure.wind_maximum,
                                    temperature=slf_measure.temperature,
                                )
                                new_measures.append(measure)
                            except ProviderException as e:
                                logger.warn("Error while processing measure '{0}' for station '{1}': {2}"
                                            .format(key, station_id, e))
                            except Exception as e:
                                logger.exception("Error while processing measure '{0}' for station '{1}': {2}"
                                                 .format(key, station_id, e))
                                self.raven_client.captureException()

                    self.insert_new_measures(measures_collection, station, new_measures, logger)

                except ProviderException as e:
                    logger.warn("Error while processing station '{0}': {1}".format(station_id, e))
                except Exception as e:
                    logger.exception("Error while processing station '{0}': {1}".format(station_id, e))
                    self.raven_client.captureException()

        except Exception as e:
            logger.exception("Error while processing SLF: {0}".format(e))
            self.raven_client.captureException()

        logger.info("Done !")

Slf().process_data()
<|endoftext|>


<commit_before>"""integration tests for BridgeDB ."""

from __future__ import print_function
from twisted.trial import unittest

import smtplib
from smtpd import SMTPServer
import asyncore
import threading
import Queue
import random
import time
import random

# ------------- SMTP Client Config
SMTP_DEBUG_LEVEL = 0  # set to 1 to see SMTP message exchange
BRIDGEDB_SMTP_SERVER_ADDRESS = "localhost"
BRIDGEDB_SMTP_SERVER_PORT = 6725
FROM_ADDRESS_TEMPLATE = "test%d@127.0.0.1" # %d is parameterised with a random integer to make the sender unique
MIN_FROM_ADDRESS = 1 # minimum value used to parameterise FROM_ADDRESS_TEMPLATE
MAX_FROM_ADDRESS = 10**8  # max value used to parameterise FROM_ADDRESS_TEMPLATE. Needs to be pretty big to reduce the chance of collisions
TO_ADDRESS = "bridges@torproject.org"
MESSAGE_TEMPLATE = """From: %s
To: %s
Subject: testing

get bridges"""

# ------------- SMTP Server Setup
# Setup an SMTP server which we use to check for responses
# from bridgedb. This needs to be done before sending the actual mail
LOCAL_SMTP_SERVER_ADDRESS = 'localhost'
LOCAL_SMTP_SERVER_PORT = 2525 # Must be the same as bridgedb's EMAIL_SMTP_PORT

class EmailServer(SMTPServer):
    def process_message(self, peer, mailfrom, rcpttos, data):
        ''' Overridden from SMTP server, called whenever a message is received'''
        self.message_queue.put(data)

    def thread_proc(self):
        ''' This function runs in thread, and will continue looping
        until the _stop Event object is set by the stop() function'''
        while self._stop.is_set() == False:
            asyncore.loop(timeout=0.0, count=1)
        # must close, or asyncore will hold on to the socket and subsequent tests will fail with 'Address not in use'
        self.close()

    def start(self):
        self.message_queue = Queue.Queue()
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self.thread_proc)
        self._thread.setDaemon(True) # ensures that if any tests do fail, then threads will exit when the parent exits
        self._thread.start()

    @classmethod
    def startServer(cls):
        #print("Starting SMTP server on %s:%s" % (LOCAL_SMTP_SERVER_ADDRESS, LOCAL_SMTP_SERVER_PORT))
        server = EmailServer((LOCAL_SMTP_SERVER_ADDRESS, LOCAL_SMTP_SERVER_PORT), None)
        server.start()
        return server

    def stop(self):
        # signal thread_proc to stop
        self._stop.set()
        # wait for thread_proc to return (shouldn't take long)
        self._thread.join()
        assert self._thread.is_alive() == False, "Thread is alive and kicking"

    def getAndCheckMessageContains(self, text, timeoutInSecs=2.0):
        #print("Checking for reponse")
        message = self.message_queue.get(block=True, timeout=timeoutInSecs)
        assert message.find(text) != -1, "Message did not contain text \"%s\". Full message is:\n %s" % (text, message)

    def checkNoMessageReceived(self, timeoutInSecs=2.0):
        try:
            self.message_queue.get(block=True, timeout=timeoutInSecs)
        except Queue.Empty:
            return True
        assert False, "Found a message in the queue, but expected none"

def sendMail(fromAddress):
    #print("Connecting to %s:%d" % (BRIDGEDB_SMTP_SERVER_ADDRESS, BRIDGEDB_SMTP_SERVER_PORT))
    client = smtplib.SMTP(BRIDGEDB_SMTP_SERVER_ADDRESS, BRIDGEDB_SMTP_SERVER_PORT)
    client.set_debuglevel(SMTP_DEBUG_LEVEL)

    #print("Sending mail TO:%s, FROM:%s" % (TO_ADDRESS, fromAddress))
    result = client.sendmail(fromAddress, TO_ADDRESS, MESSAGE_TEMPLATE % (fromAddress, TO_ADDRESS))
    assert result == {}, "Failed to send mail"
    client.quit()

class SMTPTests(unittest.TestCase):
    def setUp(self):
        ''' Called at the start of each test, ensures that the SMTP server is running'''
        self.server = EmailServer.startServer()

    def tearDown(self):
        ''' Called after each test, ensures that the SMTP server is cleaned up'''
        self.server.stop()

    def test_getBridges(self):
        # send the mail to bridgedb, choosing a random email address
        sendMail(fromAddress=FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS))

        # then check that our local SMTP server received a response
        # and that response contained some bridges
        self.server.getAndCheckMessageContains("Here are your bridges")

    def test_getBridges_rateLimitExceeded(self):
        # send the mail to bridgedb, choosing a random email address
        FROM_ADDRESS = FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS)
        sendMail(FROM_ADDRESS)

        # then check that our local SMTP server received a response
        # and that response contained some bridges
        self.server.getAndCheckMessageContains("Here are your bridges")

	# send another request from the same email address
        sendMail(FROM_ADDRESS)

        # this time, the email response should not contain any bridges
        self.server.getAndCheckMessageContains("You have exceeded the rate limit. Please slow down!")

        # then we send another request from the same email address
        sendMail(FROM_ADDRESS)

        # now there should be no response at all (wait 1 second to make sure)
        self.server.checkNoMessageReceived(timeoutInSecs=1.0)

    def test_getBridges_stressTest(self):
        ''' Sends a large number of emails in a short period of time, and checks that
            a response is received for each message '''
        NUM_MAILS = 100
        for i in range(NUM_MAILS):
            # Note: if by chance two emails with the same FROM_ADDRESS are generated, this test will fail
            # Setting 'MAX_FROM_ADDRESS' to be a high value reduces the probability of this occuring, but does not rule it out
            sendMail(fromAddress=FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS))

        for i in range(NUM_MAILS):
            self.server.getAndCheckMessageContains("Here are your bridges")

<commit_msg>Whitespace fixes for bridgedb.test.test_smtp.<commit_after>"""integration tests for BridgeDB ."""

from __future__ import print_function
from twisted.trial import unittest

import smtplib
from smtpd import SMTPServer
import asyncore
import threading
import Queue
import random
import time
import random

# ------------- SMTP Client Config
SMTP_DEBUG_LEVEL = 0  # set to 1 to see SMTP message exchange
BRIDGEDB_SMTP_SERVER_ADDRESS = "localhost"
BRIDGEDB_SMTP_SERVER_PORT = 6725
FROM_ADDRESS_TEMPLATE = "test%d@127.0.0.1" # %d is parameterised with a random integer to make the sender unique
MIN_FROM_ADDRESS = 1 # minimum value used to parameterise FROM_ADDRESS_TEMPLATE
MAX_FROM_ADDRESS = 10**8  # max value used to parameterise FROM_ADDRESS_TEMPLATE. Needs to be pretty big to reduce the chance of collisions
TO_ADDRESS = "bridges@torproject.org"
MESSAGE_TEMPLATE = """From: %s
To: %s
Subject: testing

get bridges"""

# ------------- SMTP Server Setup
# Setup an SMTP server which we use to check for responses
# from bridgedb. This needs to be done before sending the actual mail
LOCAL_SMTP_SERVER_ADDRESS = 'localhost'
LOCAL_SMTP_SERVER_PORT = 2525 # Must be the same as bridgedb's EMAIL_SMTP_PORT

class EmailServer(SMTPServer):
    def process_message(self, peer, mailfrom, rcpttos, data):
        ''' Overridden from SMTP server, called whenever a message is received'''
        self.message_queue.put(data)

    def thread_proc(self):
        ''' This function runs in thread, and will continue looping
        until the _stop Event object is set by the stop() function'''
        while self._stop.is_set() == False:
            asyncore.loop(timeout=0.0, count=1)
        # must close, or asyncore will hold on to the socket and subsequent tests will fail with 'Address not in use'
        self.close()

    def start(self):
        self.message_queue = Queue.Queue()
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self.thread_proc)
        self._thread.setDaemon(True) # ensures that if any tests do fail, then threads will exit when the parent exits
        self._thread.start()

    @classmethod
    def startServer(cls):
        #print("Starting SMTP server on %s:%s" % (LOCAL_SMTP_SERVER_ADDRESS, LOCAL_SMTP_SERVER_PORT))
        server = EmailServer((LOCAL_SMTP_SERVER_ADDRESS, LOCAL_SMTP_SERVER_PORT), None)
        server.start()
        return server

    def stop(self):
        # signal thread_proc to stop
        self._stop.set()
        # wait for thread_proc to return (shouldn't take long)
        self._thread.join()
        assert self._thread.is_alive() == False, "Thread is alive and kicking"

    def getAndCheckMessageContains(self, text, timeoutInSecs=2.0):
        #print("Checking for reponse")
        message = self.message_queue.get(block=True, timeout=timeoutInSecs)
        assert message.find(text) != -1, "Message did not contain text \"%s\". Full message is:\n %s" % (text, message)

    def checkNoMessageReceived(self, timeoutInSecs=2.0):
        try:
            self.message_queue.get(block=True, timeout=timeoutInSecs)
        except Queue.Empty:
            return True
        assert False, "Found a message in the queue, but expected none"

def sendMail(fromAddress):
    #print("Connecting to %s:%d" % (BRIDGEDB_SMTP_SERVER_ADDRESS, BRIDGEDB_SMTP_SERVER_PORT))
    client = smtplib.SMTP(BRIDGEDB_SMTP_SERVER_ADDRESS, BRIDGEDB_SMTP_SERVER_PORT)
    client.set_debuglevel(SMTP_DEBUG_LEVEL)

    #print("Sending mail TO:%s, FROM:%s" % (TO_ADDRESS, fromAddress))
    result = client.sendmail(fromAddress, TO_ADDRESS, MESSAGE_TEMPLATE % (fromAddress, TO_ADDRESS))
    assert result == {}, "Failed to send mail"
    client.quit()

class SMTPTests(unittest.TestCase):
    def setUp(self):
        ''' Called at the start of each test, ensures that the SMTP server is running'''
        self.server = EmailServer.startServer()

    def tearDown(self):
        ''' Called after each test, ensures that the SMTP server is cleaned up'''
        self.server.stop()

    def test_getBridges(self):
        # send the mail to bridgedb, choosing a random email address
        sendMail(fromAddress=FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS))

        # then check that our local SMTP server received a response
        # and that response contained some bridges
        self.server.getAndCheckMessageContains("Here are your bridges")

    def test_getBridges_rateLimitExceeded(self):
        # send the mail to bridgedb, choosing a random email address
        FROM_ADDRESS = FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS)
        sendMail(FROM_ADDRESS)

        # then check that our local SMTP server received a response
        # and that response contained some bridges
        self.server.getAndCheckMessageContains("Here are your bridges")

        # send another request from the same email address
        sendMail(FROM_ADDRESS)

        # this time, the email response should not contain any bridges
        self.server.getAndCheckMessageContains("You have exceeded the rate limit. Please slow down!")

        # then we send another request from the same email address
        sendMail(FROM_ADDRESS)

        # now there should be no response at all (wait 1 second to make sure)
        self.server.checkNoMessageReceived(timeoutInSecs=1.0)

    def test_getBridges_stressTest(self):
        ''' Sends a large number of emails in a short period of time, and checks that
            a response is received for each message '''
        NUM_MAILS = 100
        for i in range(NUM_MAILS):
            # Note: if by chance two emails with the same FROM_ADDRESS are generated, this test will fail
            # Setting 'MAX_FROM_ADDRESS' to be a high value reduces the probability of this occuring, but does not rule it out
            sendMail(fromAddress=FROM_ADDRESS_TEMPLATE % random.randint(MIN_FROM_ADDRESS, MAX_FROM_ADDRESS))

        for i in range(NUM_MAILS):
            self.server.getAndCheckMessageContains("Here are your bridges")
<|endoftext|>


<commit_before>__version__ = "9.1.0"
<commit_msg>Update version to 8.1.0<commit_after>__version__ = "9.1.0"<|endoftext|>


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

import abc
import copy
import os
import six
import smtplib

from jinja2.exceptions import TemplateNotFound
from oslo_config import cfg
from oslo_log import log

import storyboard.db.api.base as db_base
from storyboard.db.api.subscriptions import subscription_get_all_subscriber_ids
import storyboard.db.models as models
from storyboard.plugin.email.base import EmailPluginBase
from storyboard.plugin.email.factory import EmailFactory
from storyboard.plugin.email import smtp_client as smtp
from storyboard.plugin.event_worker import WorkerTaskBase

LOG = log.getLogger(__name__)

CONF = cfg.CONF

@six.add_metaclass(abc.ABCMeta)
class EmailWorkerBase(EmailPluginBase, WorkerTaskBase):
    """An abstract email construction worker.

    Abstract class that encapsulates common functionality needed in
    building emails off of our event queue.
    """

    __metaclass__ = abc.ABCMeta

    def handle(self, session, author, method, url, path, query_string, status,
               resource, resource_id, sub_resource=None, sub_resource_id=None,
               resource_before=None, resource_after=None):
        """Handle an event.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The HTTP query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

        # We only care about a subset of resource types.
        if resource not in ['task', 'project_group', 'project', 'story',
                            'branch', 'milestone', 'tag']:
            return

        # We only care about PUT, POST, and DELETE requests that do not
        # result in errors or redirects.
        if method == 'GET' or status >= 300:
            return

        # We only care if the current resource has subscribers.
        if resource == 'task' and method == 'DELETE':
            # FIXME(SotK): Workaround the fact that the task won't be in the
            # database anymore if it has been deleted.
            # We should archive instead of delete to solve this.
            # NOTE: People who are only subscribed to the task (not the story)
            # won't get an email.
            subscribers = self.get_subscribers(session, 'story',
                                               resource_before['story_id'])
        else:
            subscribers = self.get_subscribers(session, resource, resource_id)
        if not subscribers:
            return

        # Pass our values on to the handler.
        self.handle_email(session=session,
                          author=author,
                          subscribers=subscribers,
                          method=method,
                          url=url,
                          status=status,
                          path=path,
                          query_string=query_string,
                          resource=resource,
                          resource_id=resource_id,
                          sub_resource=sub_resource,
                          sub_resource_id=sub_resource_id,
                          resource_before=resource_before,
                          resource_after=resource_after)

    @abc.abstractmethod
    def handle_email(self, session, author, subscribers, method, url, path,
                     query_string, status, resource, resource_id,
                     sub_resource=None, sub_resource_id=None,
                     resource_before=None, resource_after=None):
        """Handle an email notification for the given subscribers.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param subscribers: A list of subscribers that should receive an email.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

    def get_subscribers(self, session, resource, resource_id):
        """Get a list of users who are subscribed to the resource,
        whose email address is valid, and whose email preferences indicate
        that they'd like to receive non-digest email.
        """
        subscribers = []

        # Resolve all the subscriber ID's.
        subscriber_ids = subscription_get_all_subscriber_ids(resource,
                                                             resource_id,
                                                             session=session)
        users = db_base.model_query(models.User, session) \
            .filter(models.User.id.in_(subscriber_ids)).all()

        for user in users:
            if not self.get_preference('plugin_email_enable', user) == 'true':
                continue
            subscribers.append(user)

        return subscribers

    def get_preference(self, name, user):
        if name not in user.preferences:
            return None
        return user.preferences[name].cast_value

    def get_changed_properties(self, original, new):
        """Shallow comparison diff.

        This method creates a shallow comparison between two dicts,
        and returns two dicts containing only the changed properties from
        before and after. It intentionally excludes created_at and updated_at,
        as those aren't true 'values' so to say.
        """

        # Clone our value arrays, since we might return them.
        before = copy.copy(original) or None
        after = copy.copy(new) or None

        # Strip out protected parameters
        for protected in ['created_at', 'updated_at']:
            if before and protected in before:
                del before[protected]
            if after and protected in after:
                del after[protected]

        # Sanity check, null values.
        if not before or not after:
            return before, after

        # Collect all the keys
        before_keys = set(before.keys())
        after_keys = set(after.keys())
        keys = before_keys | after_keys

        # Run the comparison.
        for key in keys:
            if key not in before:
                before[key] = None
            if key not in after:
                after[key] = None

            if after[key] == before[key]:
                del after[key]
                del before[key]

        return before, after

class SubscriptionEmailWorker(EmailWorkerBase):
    """This worker plugin generates individual event emails for users who
    have indicated that they wish to receive emails, but don't want digests.
    """

    def handle_email(self, session, author, subscribers, method, url, path,
                     query_string, status, resource, resource_id,
                     sub_resource=None, sub_resource_id=None,
                     resource_before=None, resource_after=None):
        """Send an email for a specific event.

        We assume that filtering logic has already occurred when this method
        is invoked.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param subscribers: A list of subscribers that should receive an email.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

        email_config = CONF.plugin_email

        # Retrieve the template names.
        (subject_template, text_template, html_template) = \
            self.get_templates(method=method,
                               resource_name=resource,
                               sub_resource_name=sub_resource)

        # Build our factory. If an HTML template exists, add it. If it can't
        # find the template, skip.
        try:
            factory = EmailFactory(sender=email_config.sender,
                                   subject=subject_template,
                                   text_template=text_template)
        except TemplateNotFound:
            LOG.error("Templates not found [%s, %s]" % (subject_template,
                                                        text_template))
            return

        # Try to add an HTML template
        try:
            factory.add_text_template(html_template, 'html')
        except TemplateNotFound:
            LOG.debug('Template %s not found' % (html_template,))

        # If there's a reply-to in our config, add that.
        if email_config.reply_to:
            factory.add_header('Reply-To', email_config.reply_to)

        # Resolve the resource instance
        resource_instance = self.resolve_resource_by_name(session, resource,
                                                          resource_id)
        sub_resource_instance = self.resolve_resource_by_name(session,
                                                              sub_resource,
                                                              sub_resource_id)

        # Figure out the diff between old and new.
        before, after = self.get_changed_properties(resource_before,
                                                    resource_after)

        # For each subscriber, create the email and send it.
        with smtp.get_smtp_client() as smtp_client:
            for subscriber in subscribers:

                # Make sure this subscriber's preferences indicate they want
                # email and they're not receiving digests.
                if not self.get_preference('plugin_email_enable', subscriber) \
                        or self.get_preference('plugin_email_digest',
                                               subscriber):
                    continue

                try:
                    # Build an email.
                    message = factory.build(recipient=subscriber.email,
                                            author=author,
                                            resource=resource_instance,
                                            sub_resource=sub_resource_instance,
                                            url=url,
                                            query_string=query_string,
                                            before=before,
                                            after=after)
                    # Send the email.
                    from_addr = message.get('From')
                    to_addrs = message.get('To')

                    try:
                        smtp_client.sendmail(from_addr=from_addr,
                                             to_addrs=to_addrs,
                                             msg=message.as_string())
                    except smtplib.SMTPException as e:
                        LOG.error('Cannot send email, discarding: %s' % (e,))
                except Exception as e:
                    # Skip, keep going.
                    LOG.error("Cannot schedule email: %s" % (e.message,))

    def get_templates(self, method, resource_name, sub_resource_name=None):
        """Return the email templates for the given resource.

        This method builds the names of templates for a provided resource
        action. The template folder structure is as follows:

        /{{resource_name}}/{{method}}_subject.txt
        /{{resource_name}}/{{method}}.txt
        /{{resource_name}}/{{method}}.html (optional)

        For subresources, it is as follows:
        /{{resource_name}}/{{subresource_name}}/{{method}}_subject.txt
        /{{resource_name}}/{{subresource_name}}/{{method}}.txt
        /{{resource_name}}/{{subresource_name}}/{{method}}.html (optional)
        """
        ## TODO(krotscheck): Templates can also resolve by user language.

        if sub_resource_name:
            base_template = os.path.join(resource_name, sub_resource_name)
        else:
            base_template = resource_name

        base_file = '%s' % (method,)

        subject_template = os.path.join(base_template,
                                        '%s_subject.txt' % (base_file,))
        text_template = os.path.join(base_template,
                                     '%s.txt' % (base_file,))
        html_template = os.path.join(base_template,
                                     '%s.html' % (base_file,))

        return subject_template, text_template, html_template
<commit_msg>Add support for email threading by story_id<commit_after># Copyright (c) 2015 Hewlett-Packard Development Company, L.P.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing permissions and
# limitations under the License.

import abc
import copy
import os
import six
import smtplib

from email.utils import make_msgid
from jinja2.exceptions import TemplateNotFound
from oslo_config import cfg
from oslo_log import log
from socket import getfqdn

import storyboard.db.api.base as db_base
from storyboard.db.api.subscriptions import subscription_get_all_subscriber_ids
import storyboard.db.models as models
from storyboard.plugin.email.base import EmailPluginBase
from storyboard.plugin.email.factory import EmailFactory
from storyboard.plugin.email import smtp_client as smtp
from storyboard.plugin.event_worker import WorkerTaskBase

LOG = log.getLogger(__name__)

CONF = cfg.CONF

@six.add_metaclass(abc.ABCMeta)
class EmailWorkerBase(EmailPluginBase, WorkerTaskBase):
    """An abstract email construction worker.

    Abstract class that encapsulates common functionality needed in
    building emails off of our event queue.
    """

    __metaclass__ = abc.ABCMeta

    def handle(self, session, author, method, url, path, query_string, status,
               resource, resource_id, sub_resource=None, sub_resource_id=None,
               resource_before=None, resource_after=None):
        """Handle an event.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The HTTP query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

        # We only care about a subset of resource types.
        if resource not in ['task', 'project_group', 'project', 'story',
                            'branch', 'milestone', 'tag']:
            return

        # We only care about PUT, POST, and DELETE requests that do not
        # result in errors or redirects.
        if method == 'GET' or status >= 300:
            return

        # We only care if the current resource has subscribers.
        if resource == 'task' and method == 'DELETE':
            # FIXME(SotK): Workaround the fact that the task won't be in the
            # database anymore if it has been deleted.
            # We should archive instead of delete to solve this.
            # NOTE: People who are only subscribed to the task (not the story)
            # won't get an email.
            subscribers = self.get_subscribers(session, 'story',
                                               resource_before['story_id'])
        else:
            subscribers = self.get_subscribers(session, resource, resource_id)
        if not subscribers:
            return

        # Pass our values on to the handler.
        self.handle_email(session=session,
                          author=author,
                          subscribers=subscribers,
                          method=method,
                          url=url,
                          status=status,
                          path=path,
                          query_string=query_string,
                          resource=resource,
                          resource_id=resource_id,
                          sub_resource=sub_resource,
                          sub_resource_id=sub_resource_id,
                          resource_before=resource_before,
                          resource_after=resource_after)

    @abc.abstractmethod
    def handle_email(self, session, author, subscribers, method, url, path,
                     query_string, status, resource, resource_id,
                     sub_resource=None, sub_resource_id=None,
                     resource_before=None, resource_after=None):
        """Handle an email notification for the given subscribers.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param subscribers: A list of subscribers that should receive an email.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

    def get_subscribers(self, session, resource, resource_id):
        """Get a list of users who are subscribed to the resource,
        whose email address is valid, and whose email preferences indicate
        that they'd like to receive non-digest email.
        """
        subscribers = []

        # Resolve all the subscriber ID's.
        subscriber_ids = subscription_get_all_subscriber_ids(resource,
                                                             resource_id,
                                                             session=session)
        users = db_base.model_query(models.User, session) \
            .filter(models.User.id.in_(subscriber_ids)).all()

        for user in users:
            if not self.get_preference('plugin_email_enable', user) == 'true':
                continue
            subscribers.append(user)

        return subscribers

    def get_preference(self, name, user):
        if name not in user.preferences:
            return None
        return user.preferences[name].cast_value

    def get_changed_properties(self, original, new):
        """Shallow comparison diff.

        This method creates a shallow comparison between two dicts,
        and returns two dicts containing only the changed properties from
        before and after. It intentionally excludes created_at and updated_at,
        as those aren't true 'values' so to say.
        """

        # Clone our value arrays, since we might return them.
        before = copy.copy(original) or None
        after = copy.copy(new) or None

        # Strip out protected parameters
        for protected in ['created_at', 'updated_at']:
            if before and protected in before:
                del before[protected]
            if after and protected in after:
                del after[protected]

        # Sanity check, null values.
        if not before or not after:
            return before, after

        # Collect all the keys
        before_keys = set(before.keys())
        after_keys = set(after.keys())
        keys = before_keys | after_keys

        # Run the comparison.
        for key in keys:
            if key not in before:
                before[key] = None
            if key not in after:
                after[key] = None

            if after[key] == before[key]:
                del after[key]
                del before[key]

        return before, after

class SubscriptionEmailWorker(EmailWorkerBase):
    """This worker plugin generates individual event emails for users who
    have indicated that they wish to receive emails, but don't want digests.
    """

    def handle_email(self, session, author, subscribers, method, url, path,
                     query_string, status, resource, resource_id,
                     sub_resource=None, sub_resource_id=None,
                     resource_before=None, resource_after=None):
        """Send an email for a specific event.

        We assume that filtering logic has already occurred when this method
        is invoked.

        :param session: An event-specific SQLAlchemy session.
        :param author: The author's user record.
        :param subscribers: A list of subscribers that should receive an email.
        :param method: The HTTP Method.
        :param url: The Referer header from the request.
        :param path: The full HTTP Path requested.
        :param query_string: The query string from the request.
        :param status: The returned HTTP Status of the response.
        :param resource: The resource type.
        :param resource_id: The ID of the resource.
        :param sub_resource: The subresource type.
        :param sub_resource_id: The ID of the subresource.
        :param resource_before: The resource state before this event occurred.
        :param resource_after: The resource state after this event occurred.
        """

        email_config = CONF.plugin_email

        # Retrieve the template names.
        (subject_template, text_template, html_template) = \
            self.get_templates(method=method,
                               resource_name=resource,
                               sub_resource_name=sub_resource)

        # Build our factory. If an HTML template exists, add it. If it can't
        # find the template, skip.
        try:
            factory = EmailFactory(sender=email_config.sender,
                                   subject=subject_template,
                                   text_template=text_template)
        except TemplateNotFound:
            LOG.error("Templates not found [%s, %s]" % (subject_template,
                                                        text_template))
            return

        # Try to add an HTML template
        try:
            factory.add_text_template(html_template, 'html')
        except TemplateNotFound:
            LOG.debug('Template %s not found' % (html_template,))

        # If there's a reply-to in our config, add that.
        if email_config.reply_to:
            factory.add_header('Reply-To', email_config.reply_to)

        # Resolve the resource instance
        resource_instance = self.resolve_resource_by_name(session, resource,
                                                          resource_id)
        sub_resource_instance = self.resolve_resource_by_name(session,
                                                              sub_resource,
                                                              sub_resource_id)

        # Set In-Reply-To message id for 'task' and 'story' resources
        if resource == 'task' and method == 'DELETE':
            # FIXME(pedroalvarez): Workaround the fact that the task won't be
            # in the database anymore if it has been deleted.
            # We should archive instead of delete to solve this.
            story_id = resource_before['story_id']
            created_at = self.resolve_resource_by_name(session, 'story',
                story_id).created_at
        elif resource == 'task':
            story_id = resource_instance.story.id
            created_at = resource_instance.story.created_at
        elif resource == 'story':
            story_id = resource_instance.id
            created_at = resource_instance.created_at

        if story_id and created_at:
            thread_id = "<storyboard.story.%s.%s@%s>" % (
                created_at.strftime("%Y%m%d%H%M"),
                story_id,
                getfqdn()
            )
        else:
            thread_id = make_msgid()

        factory.add_header("In-Reply-To", thread_id)

        # Figure out the diff between old and new.
        before, after = self.get_changed_properties(resource_before,
                                                    resource_after)

        # For each subscriber, create the email and send it.
        with smtp.get_smtp_client() as smtp_client:
            for subscriber in subscribers:

                # Make sure this subscriber's preferences indicate they want
                # email and they're not receiving digests.
                if not self.get_preference('plugin_email_enable', subscriber) \
                        or self.get_preference('plugin_email_digest',
                                               subscriber):
                    continue

                try:
                    # Build an email.
                    message = factory.build(recipient=subscriber.email,
                                            author=author,
                                            resource=resource_instance,
                                            sub_resource=sub_resource_instance,
                                            url=url,
                                            query_string=query_string,
                                            before=before,
                                            after=after)
                    # Send the email.
                    from_addr = message.get('From')
                    to_addrs = message.get('To')

                    try:
                        smtp_client.sendmail(from_addr=from_addr,
                                             to_addrs=to_addrs,
                                             msg=message.as_string())
                    except smtplib.SMTPException as e:
                        LOG.error('Cannot send email, discarding: %s' % (e,))
                except Exception as e:
                    # Skip, keep going.
                    LOG.error("Cannot schedule email: %s" % (e.message,))

    def get_templates(self, method, resource_name, sub_resource_name=None):
        """Return the email templates for the given resource.

        This method builds the names of templates for a provided resource
        action. The template folder structure is as follows:

        /{{resource_name}}/{{method}}_subject.txt
        /{{resource_name}}/{{method}}.txt
        /{{resource_name}}/{{method}}.html (optional)

        For subresources, it is as follows:
        /{{resource_name}}/{{subresource_name}}/{{method}}_subject.txt
        /{{resource_name}}/{{subresource_name}}/{{method}}.txt
        /{{resource_name}}/{{subresource_name}}/{{method}}.html (optional)
        """
        ## TODO(krotscheck): Templates can also resolve by user language.

        if sub_resource_name:
            base_template = os.path.join(resource_name, sub_resource_name)
        else:
            base_template = resource_name

        base_file = '%s' % (method,)

        subject_template = os.path.join(base_template,
                                        '%s_subject.txt' % (base_file,))
        text_template = os.path.join(base_template,
                                     '%s.txt' % (base_file,))
        html_template = os.path.join(base_template,
                                     '%s.html' % (base_file,))

        return subject_template, text_template, html_template
<|endoftext|>


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

from __future__ import absolute_import, unicode_literals, division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys

# Third-party
import astropy.units as u
import numpy as np
import pytest

from ...misc.units import UnitSystem
from ..particles import *
from ..orbits import *

usys = UnitSystem(u.kpc, u.Myr, u.M_sun)

def test_orbitcollection_init():
    t = np.linspace(0., 10., 100)*u.yr
    r = np.random.random((100,10,3))*u.kpc
    v = np.random.random((100,10,3))*u.kpc/u.yr
    m = np.random.random(10)*u.M_sun

    oc = OrbitCollection(t=t, r=r, v=v, m=m)

    # Should fail -- different length units, no usys:
    t = np.linspace(0., 10., 100)*u.yr
    r = np.random.random((100,10,3))*u.kpc
    v = np.random.random((100,10,3))*u.km/u.s
    m = np.random.random(10)*u.M_sun
    with pytest.raises(ValueError):
        oc = OrbitCollection(t=t, r=r, v=v, m=m)

    # should pass bc we give a usys
    oc = OrbitCollection(t=t, r=r, v=v, m=m, unit_system=usys)
<commit_msg>added test for to() in orbitcollection<commit_after># coding: utf-8
""" """

from __future__ import absolute_import, unicode_literals, division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys

# Third-party
import astropy.units as u
import numpy as np
import pytest

from ...misc.units import UnitSystem
from ..particles import *
from ..orbits import *

usys = UnitSystem(u.kpc, u.Myr, u.M_sun)

def test_orbitcollection_init():
    t = np.linspace(0., 10., 100)*u.yr
    r = np.random.random((100,10,3))*u.kpc
    v = np.random.random((100,10,3))*u.kpc/u.yr
    m = np.random.random(10)*u.M_sun

    oc = OrbitCollection(t=t, r=r, v=v, m=m)

    # Should fail -- different length units, no usys:
    t = np.linspace(0., 10., 100)*u.yr
    r = np.random.random((100,10,3))*u.kpc
    v = np.random.random((100,10,3))*u.km/u.s
    m = np.random.random(10)*u.M_sun
    with pytest.raises(ValueError):
        oc = OrbitCollection(t=t, r=r, v=v, m=m)

    # should pass bc we give a usys
    oc = OrbitCollection(t=t, r=r, v=v, m=m, unit_system=usys)

def test_to():
    t = np.arange(0., 100, 0.1)*u.Myr
    r = np.random.random(size=(len(t),10,3))*u.kpc
    v = np.random.random(size=(len(t),10,3))*u.kpc/u.Myr
    m = np.random.random(10)*u.M_sun

    pc = OrbitCollection(t=t, r=r, v=v, m=m, unit_system=usys)

    usys2 = UnitSystem(u.km, u.s, u.kg)
    pc2 = pc.to(usys2)

    assert np.allclose(pc2._t, t.to(u.s).value)
    assert np.allclose(pc2._r, r.to(u.km).value)
    assert np.allclose(pc2._v, v.to(u.km/u.s).value)
    assert np.allclose(pc2._m, m.to(u.kg).value)<|endoftext|>


<commit_before># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import ConfigParser

from lib.cuckoo.common.exceptions import CuckooMachineError

class Dictionary(dict):
    """Cuckko custom dict."""

    def __getattr__(self, key):
        return self.get(key, None)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class MachineManager(object):
    """Base abstract class for analysis machine manager."""

    def __init__(self):
        self.module_name = ""
        self.config_path = ""
        self.config = ConfigParser.ConfigParser()
        self.options = {}
        self.machines = []

    def initialize(self, module_name):
        """Read configuration.
        @param module_name: module name.
        """
        self.module_name = module_name
        self.config_path = os.path.join("conf", "%s.conf" % module_name)
        self.config.read(self.config_path)

        machines_list = self.config.get(self.module_name, "machines").strip().split(",")
        for machine_id in machines_list:
            machine = Dictionary()
            machine.id = machine_id
            machine.label = self.config.get(machine_id, "label")
            machine.platform = self.config.get(machine_id, "platform")
            machine.ip = self.config.get(machine_id, "ip")
            machine.locked = False
            self.machines.append(machine)

        # Checks if machines configured are really available.
        try:
            configured_vm = self._list()
            for machine in self.machines:
                if machine.label not in configured_vm:
                    raise CuckooMachineError("Configured VM %s was not detected or not in proper state" % machine.label)
        except NotImplementedError:
            pass

    def availables(self):
        """How many machines are free.
        @return: free machines count.
        """
        count = 0
        for machine in self.machines:
            if not machine.locked:
                count += 1

        return count

    def acquire(self, label=None, platform=None):
        """Acquire a machine to start analysis.
        @param label: machine name.
        @param platform: machine platform.
        @return: machine or None.
        """
        if label:
            for machine in self.machines:
                if machine.label == label and not machine.locked:
                    machine.locked = True
                    return machine
        elif platform:
            for machine in self.machines:
                if machine.platform == platform and not machine.locked:
                    machine.locked = True
                    return machine
        else:
            for machine in self.machines:
                if not machine.locked:
                    machine.locked = True
                    return machine

        return None

    def release(self, label=None):
        """Release a machine.
        @param label: machine name.
        """
        if label:
            for machine in self.machines:
                if machine.label == label:
                    machine.locked = False

    def running(self):
        """Returns running virutal machines.
        @return: running virtual machines list.
        """
        return [m for m in self.machines if m.locked]

    def start(self, label=None):
        """Start a machine.
        @param label: machine name.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

    def stop(self, label=None):
        """Stop a machine.
        @param label: machine name.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

    def _list(self):
        """Lists virtual machines configured.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Processing(object):
    """Base abstract class for processing module."""

    def __init__(self):
        self.analysis_path = ""
        self.logs_path = ""

    def set_path(self, analysis_path):
        """Set paths.
        @param analysis_path: analysis folder path.
        """
        self.analysis_path = analysis_path
        self.log_path = os.path.join(analysis_path, "analysis.log")
        self.conf_path = os.path.join(analysis_path, "analysis.conf")
        self.file_path = os.path.join(analysis_path, "binary")
        self.dropped_path = os.path.join(analysis_path, "files")
        self.logs_path = os.path.join(analysis_path, "logs")
        self.shots_path = os.path.join(analysis_path, "shots")
        self.pcap_path = os.path.join(analysis_path, "dump.pcap")

    def run(self):
        """Start processing.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Signature(object):
    """Base abstract class for signature."""

    name = ""
    description = ""
    severity = 1
    references = []
    alert = False
    enabled = True

    def __init__(self):
        self.data = []

    def run(self, results=None):
        """Start signature processing.
        @param results: analysis results.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Report(object):
    """Base abstract class for reporting module."""

    def __init__(self):
        self.analysis_path = ""
        self.reports_path = ""
        self.options = None

    def set_path(self, analysis_path):
        """Set analysis folder path.
        @param analysis_path: analysis folder path.
        """
        self.analysis_path = analysis_path
        self.reports_path = os.path.join(self.analysis_path, "reports")

        if not os.path.exists(self.reports_path):
            os.mkdir(self.reports_path)

    def set_options(self, options):
        """Set report options.
        @param options: report options dict.
        """
        self.options = options

    def run(self):
        """Start report processing.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError
<commit_msg>Fixed a bug in lib/cuckoo/common/abstracts MachineManager, because it did not take into account the actual CUCKOO_ROOT when forming the path to the config file<commit_after># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import ConfigParser

from lib.cuckoo.common.exceptions import CuckooMachineError
from lib.cuckoo.common.constants import CUCKOO_ROOT

class Dictionary(dict):
    """Cuckko custom dict."""

    def __getattr__(self, key):
        return self.get(key, None)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class MachineManager(object):
    """Base abstract class for analysis machine manager."""

    def __init__(self):
        self.module_name = ""
        self.config_path = ""
        self.config = ConfigParser.ConfigParser()
        self.options = {}
        self.machines = []

    def initialize(self, module_name):
        """Read configuration.
        @param module_name: module name.
        """
        self.module_name = module_name
        self.config_path = os.path.join(CUCKOO_ROOT, "conf", "%s.conf" % module_name)
        self.config.read(self.config_path)

        machines_list = self.config.get(self.module_name, "machines").strip().split(",")
        for machine_id in machines_list:
            machine = Dictionary()
            machine.id = machine_id
            machine.label = self.config.get(machine_id, "label")
            machine.platform = self.config.get(machine_id, "platform")
            machine.ip = self.config.get(machine_id, "ip")
            machine.locked = False
            self.machines.append(machine)

        # Checks if machines configured are really available.
        try:
            configured_vm = self._list()
            for machine in self.machines:
                if machine.label not in configured_vm:
                    raise CuckooMachineError("Configured VM %s was not detected or not in proper state" % machine.label)
        except NotImplementedError:
            pass

    def availables(self):
        """How many machines are free.
        @return: free machines count.
        """
        count = 0
        for machine in self.machines:
            if not machine.locked:
                count += 1

        return count

    def acquire(self, label=None, platform=None):
        """Acquire a machine to start analysis.
        @param label: machine name.
        @param platform: machine platform.
        @return: machine or None.
        """
        if label:
            for machine in self.machines:
                if machine.label == label and not machine.locked:
                    machine.locked = True
                    return machine
        elif platform:
            for machine in self.machines:
                if machine.platform == platform and not machine.locked:
                    machine.locked = True
                    return machine
        else:
            for machine in self.machines:
                if not machine.locked:
                    machine.locked = True
                    return machine

        return None

    def release(self, label=None):
        """Release a machine.
        @param label: machine name.
        """
        if label:
            for machine in self.machines:
                if machine.label == label:
                    machine.locked = False

    def running(self):
        """Returns running virutal machines.
        @return: running virtual machines list.
        """
        return [m for m in self.machines if m.locked]

    def start(self, label=None):
        """Start a machine.
        @param label: machine name.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

    def stop(self, label=None):
        """Stop a machine.
        @param label: machine name.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

    def _list(self):
        """Lists virtual machines configured.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Processing(object):
    """Base abstract class for processing module."""

    def __init__(self):
        self.analysis_path = ""
        self.logs_path = ""

    def set_path(self, analysis_path):
        """Set paths.
        @param analysis_path: analysis folder path.
        """
        self.analysis_path = analysis_path
        self.log_path = os.path.join(analysis_path, "analysis.log")
        self.conf_path = os.path.join(analysis_path, "analysis.conf")
        self.file_path = os.path.join(analysis_path, "binary")
        self.dropped_path = os.path.join(analysis_path, "files")
        self.logs_path = os.path.join(analysis_path, "logs")
        self.shots_path = os.path.join(analysis_path, "shots")
        self.pcap_path = os.path.join(analysis_path, "dump.pcap")

    def run(self):
        """Start processing.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Signature(object):
    """Base abstract class for signature."""

    name = ""
    description = ""
    severity = 1
    references = []
    alert = False
    enabled = True

    def __init__(self):
        self.data = []

    def run(self, results=None):
        """Start signature processing.
        @param results: analysis results.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError

class Report(object):
    """Base abstract class for reporting module."""

    def __init__(self):
        self.analysis_path = ""
        self.reports_path = ""
        self.options = None

    def set_path(self, analysis_path):
        """Set analysis folder path.
        @param analysis_path: analysis folder path.
        """
        self.analysis_path = analysis_path
        self.reports_path = os.path.join(self.analysis_path, "reports")

        if not os.path.exists(self.reports_path):
            os.mkdir(self.reports_path)

    def set_options(self, options):
        """Set report options.
        @param options: report options dict.
        """
        self.options = options

    def run(self):
        """Start report processing.
        @raise NotImplementedError: this method is abstract.
        """
        raise NotImplementedError
<|endoftext|>