<commit_before>import hoomd
import numpy
import pytest
import os
from copy import deepcopy

@pytest.fixture(scope="function")
def get_snapshot(device):
    def make_snapshot(n=10, particle_types=['A']):
        s = hoomd.snapshot.Snapshot(device.comm)
        if s.exists:
            s.configuration.box = [20, 20, 20, 0, 0, 0]
            s.particles.N = n
            s.particles.position[:] = numpy.random.uniform(-10, 10, size=(n, 3))
            s.particles.types = particle_types
        return s
    return make_snapshot

def test_initialization(device, simulation_factory, get_snapshot):
    with pytest.raises(TypeError):
        sim = hoomd.simulation.Simulation()

    sim = hoomd.simulation.Simulation(device)
    with pytest.raises(RuntimeError):
        sim.run(1)  # Before setting state

    sim = simulation_factory(get_snapshot())
    with pytest.raises(RuntimeError):
        sim.run(1)  # Before scheduling operations

    sim.operations.schedule()
    sim.run(1)

def test_run(simulation_factory, get_snapshot):
    sim = simulation_factory(get_snapshot())
    sim.operations.schedule()
    assert sim.timestep == 0
    steps = 0
    n_step_list = [1, 10, 100]
    for n_steps in n_step_list:
        steps += n_steps
        sim.run(n_steps)
        assert sim.timestep == steps
    assert sim.timestep == sum(n_step_list)

_tetrahedron_verts = [(0.5, 0.5, 0.5), (-0.5, -0.5, 0.5),
                      (0.5, -0.5, -0.5), (-0.5, 0.5, -0.5)]

_state_args = [((10, ['A']),
                hoomd.hpmc.integrate.Sphere,
                {'diameter': 1}, [10]),
               ((100, ['A']),
                hoomd.hpmc.integrate.ConvexPolyhedron,
                {'vertices': _tetrahedron_verts}, [10, 3]),
               ((5, ['A']),
                hoomd.hpmc.integrate.Ellipsoid,
                {'a': 0.2, 'b': 0.25, 'c': 0.5}, [1, 3, 4]),
               ((50, ['A']), hoomd.hpmc.integrate.ConvexSpheropolyhedron,
                {'vertices': _tetrahedron_verts, 'sweep_radius': 0.1}, [10])]

@pytest.fixture(scope="function", params=_state_args)
def state_args(request):
    return deepcopy(request.param)

class TemporaryFileContext():
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        os.remove(self.filename)

def test_state_from_gsd(simulation_factory, get_snapshot, device, state_args):
    filename = 'temporary_test_file.gsd'

    with TemporaryFileContext(filename) as file:
        snap_params, integrator, shape_dict, run_sequence = state_args
        sim = simulation_factory(get_snapshot(n=snap_params[0],
                                              particle_types=snap_params[1]))
        mc = integrator(2345)
        mc.shape['A'] = shape_dict

        sim.operations.add(mc)
        gsd_dumper = hoomd.dump.GSD(filename=file.filename,
                                    trigger=1,
                                    overwrite=True)
        gsd_logger = hoomd.logger.Logger()
        gsd_logger += mc
        gsd_dumper.log = gsd_logger
        sim.operations.add(gsd_dumper)
        sim.operations.schedule()
        pos_dict = {}
        initial_pos = sim.state.snapshot.particles.position
        count = 0
        for nsteps in run_sequence:
            sim.run(nsteps)
            count += nsteps
            pos_dict[count] = sim.state.snapshot.particles.position

        final_pos = sim.state.snapshot.particles.position
        sim.run(1)

        initial_pos_sim = hoomd.simulation.Simulation(device)
        initial_pos_sim.create_state_from_gsd(file.filename, frame=0)
        initial_pos_snap = initial_pos_sim.state.snapshot
        numpy.testing.assert_allclose(initial_pos,
                                      initial_pos_snap.particles.position)

        final_pos_sim = hoomd.simulation.Simulation(device)
        final_pos_sim.create_state_from_gsd(file.filename)
        final_pos_snap = final_pos_sim.state.snapshot
        numpy.testing.assert_allclose(final_pos,
                                      final_pos_snap.particles.position)

        for nsteps, positions in pos_dict.items():
            sim = hoomd.simulation.Simulation(device)
            sim.create_state_from_gsd(file.filename, frame=nsteps)
            snap = sim.state.snapshot
            numpy.testing.assert_allclose(positions, snap.particles.position)
<commit_msg>Check all particle data<commit_after>import hoomd
import numpy as np
import pytest
import os
from copy import deepcopy

@pytest.fixture(scope="function")
def get_snapshot(device):
    def make_snapshot(n=10, particle_types=['A']):
        s = hoomd.snapshot.Snapshot(device.comm)
        if s.exists:
            s.configuration.box = [20, 20, 20, 0, 0, 0]
            s.particles.N = n
            s.particles.position[:] = np.random.uniform(-10, 10, size=(n, 3))
            s.particles.types = particle_types
        return s
    return make_snapshot

def assert_equivalent_snapshots(snap1, snap2):
    particles1 = snap1.particles
    particles2 = snap2.particles
    particle_data = [(particles1.N, particles2.N),
                     (particles1.typeid, particles2.typeid),
                     (particles1.mass, particles2.mass),
                     (particles1.diameter, particles2.diameter),
                     (particles1.charge, particles2.charge),
                     (particles1.position, particles2.position),
                     (particles1.orientation, particles2.orientation),
                     (particles1.velocity, particles2.velocity),
                     (particles1.acceleration, particles2.acceleration),
                     (particles1.image, particles2.image),
                     (particles1.body, particles2.body),
                     (particles1.moment_inertia, particles2.moment_inertia),
                     (particles1.angmom, particles2.angmom)]

    assert particles1.types == particles2.types
    for snap_data1, snap_data2 in particle_data:
        np.testing.assert_allclose(snap_data1, snap_data2)

def test_initialization(device, simulation_factory, get_snapshot):
    with pytest.raises(TypeError):
        sim = hoomd.simulation.Simulation()

    sim = hoomd.simulation.Simulation(device)
    with pytest.raises(RuntimeError):
        sim.run(1)  # Before setting state

    sim = simulation_factory(get_snapshot())
    with pytest.raises(RuntimeError):
        sim.run(1)  # Before scheduling operations

    sim.operations.schedule()
    sim.run(1)

def test_run(simulation_factory, get_snapshot):
    sim = simulation_factory(get_snapshot())
    sim.operations.schedule()
    assert sim.timestep == 0
    steps = 0
    n_step_list = [1, 10, 100]
    for n_steps in n_step_list:
        steps += n_steps
        sim.run(n_steps)
        assert sim.timestep == steps
    assert sim.timestep == sum(n_step_list)

_tetrahedron_verts = [(0.5, 0.5, 0.5), (-0.5, -0.5, 0.5),
                      (0.5, -0.5, -0.5), (-0.5, 0.5, -0.5)]

_state_args = [((10, ['A']),
                hoomd.hpmc.integrate.Sphere,
                {'diameter': 1}, [10]),
               ((100, ['A']),
                hoomd.hpmc.integrate.ConvexPolyhedron,
                {'vertices': _tetrahedron_verts}, [10, 3]),
               ((5, ['A']),
                hoomd.hpmc.integrate.Ellipsoid,
                {'a': 0.2, 'b': 0.25, 'c': 0.5}, [1, 3, 4]),
               ((50, ['A']), hoomd.hpmc.integrate.ConvexSpheropolyhedron,
                {'vertices': _tetrahedron_verts, 'sweep_radius': 0.1}, [10])]

@pytest.fixture(scope="function", params=_state_args)
def state_args(request):
    return deepcopy(request.param)

class TemporaryFileContext():
    def __init__(self, filename):
        self.filename = filename

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        os.remove(self.filename)

def test_state_from_gsd(simulation_factory, get_snapshot, device, state_args):
    filename = 'temporary_test_file.gsd'
    snap_params, integrator, shape_dict, run_sequence = state_args

    with TemporaryFileContext(filename) as file:
        sim = simulation_factory(get_snapshot(n=snap_params[0],
                                              particle_types=snap_params[1]))
        mc = integrator(2345)
        mc.shape['A'] = shape_dict

        sim.operations.add(mc)
        gsd_dumper = hoomd.dump.GSD(filename=file.filename,
                                    trigger=1,
                                    overwrite=True)
        gsd_logger = hoomd.logger.Logger()
        gsd_logger += mc
        gsd_dumper.log = gsd_logger
        sim.operations.add(gsd_dumper)
        sim.operations.schedule()
        snapshot_dict = {}
        initial_snap = sim.state.snapshot

        count = 0
        for nsteps in run_sequence:
            sim.run(nsteps)
            count += nsteps
            snapshot_dict[count] = sim.state.snapshot

        final_snap = sim.state.snapshot
        sim.run(1)

        initial_snap_sim = hoomd.simulation.Simulation(device)
        initial_snap_sim.create_state_from_gsd(file.filename, frame=0)
        assert_equivalent_snapshots(initial_snap,
                                    initial_snap_sim.state.snapshot)

        final_snap_sim = hoomd.simulation.Simulation(device)
        final_snap_sim.create_state_from_gsd(file.filename)
        assert_equivalent_snapshots(final_snap, final_snap_sim.state.snapshot)

        for nsteps, snap in snapshot_dict.items():
            sim = hoomd.simulation.Simulation(device)
            sim.create_state_from_gsd(file.filename, frame=nsteps)
            snap = sim.state.snapshot
            assert_equivalent_snapshots(snap, sim.state.snapshot)
<|endoftext|>


<commit_before>""" Mailu admin app utilities
"""

try:
    import cPickle as pickle
except ImportError:
    import pickle

import dns.resolver
import dns.exception
import dns.flags
import dns.rdtypes
import dns.rdatatype
import dns.rdataclass

import hmac
import secrets
import time

from multiprocessing import Value
from mailu import limiter
from flask import current_app as app

import flask
import flask_login
import flask_migrate
import flask_babel
import ipaddress
import redis

from datetime import datetime, timedelta
from flask.sessions import SessionMixin, SessionInterface
from itsdangerous.encoding import want_bytes
from werkzeug.datastructures import CallbackDict
from werkzeug.middleware.proxy_fix import ProxyFix

# Login configuration
login = flask_login.LoginManager()
login.login_view = "sso.login"

@login.unauthorized_handler
def handle_needs_login():
    """ redirect unauthorized requests to login page """
    return flask.redirect(
        flask.url_for('sso.login')
    )

# DNS stub configured to do DNSSEC enabled queries
resolver = dns.resolver.Resolver()
resolver.use_edns(0, 0, 1232)
resolver.flags = dns.flags.AD | dns.flags.RD

def has_dane_record(domain, timeout=10):
    try:
        result = resolver.query(f'_25._tcp.{domain}', dns.rdatatype.TLSA,dns.rdataclass.IN, lifetime=timeout)
        if result.response.flags & dns.flags.AD:
            for record in result:
                if isinstance(record, dns.rdtypes.ANY.TLSA.TLSA):
                    record.validate()
                    if record.usage in [2,3] and record.selector in [0,1] and record.mtype in [0,1,2]:
                        return True
    except dns.resolver.NoNameservers:
        # If the DNSSEC data is invalid and the DNS resolver is DNSSEC enabled
        # we will receive this non-specific exception. The safe behaviour is to
        # accept to defer the email.
        app.logger.warn(f'Unable to lookup the TLSA record for {domain}. Is the DNSSEC zone okay on https://dnsviz.net/d/{domain}/dnssec/?')
        return app.config['DEFER_ON_TLS_ERROR']
    except dns.exception.Timeout:
        app.logger.warn(f'Timeout while resolving the TLSA record for {domain} ({timeout}s).')
    except (dns.resolver.NXDOMAIN, dns.name.EmptyLabel):
        pass # this is expected, not TLSA record is fine
    except Exception as e:
        app.logger.info(f'Error while looking up the TLSA record for {domain} {e}')
        pass

# Rate limiter
limiter = limiter.LimitWraperFactory()

def extract_network_from_ip(ip):
    n = ipaddress.ip_network(ip)
    if n.version == 4:
        return str(n.supernet(prefixlen_diff=(32-app.config["AUTH_RATELIMIT_IP_V4_MASK"])).network_address)
    else:
        return str(n.supernet(prefixlen_diff=(128-app.config["AUTH_RATELIMIT_IP_V6_MASK"])).network_address)

def is_exempt_from_ratelimits(ip):
    ip = ipaddress.ip_address(ip)
    return any(ip in cidr for cidr in app.config['AUTH_RATELIMIT_EXEMPTION'])

# Application translation
babel = flask_babel.Babel()

@babel.localeselector
def get_locale():
    """ selects locale for translation """
    if not app.config['SESSION_COOKIE_NAME'] in flask.request.cookies:
        return flask.request.accept_languages.best_match(app.config.translations.keys())
    language = flask.session.get('language')
    if not language in app.config.translations:
        language = flask.request.accept_languages.best_match(app.config.translations.keys())
        flask.session['language'] = language
    return language

# Proxy fixer
class PrefixMiddleware(object):
    """ fix proxy headers """
    def __init__(self):
        self.app = None

    def __call__(self, environ, start_response):
        return self.app(environ, start_response)

    def init_app(self, app):
        self.app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1)
        app.wsgi_app = self

proxy = PrefixMiddleware()

# Data migrate
migrate = flask_migrate.Migrate()

# session store (inspired by https://github.com/mbr/flask-kvsession)
class RedisStore:
    """ Stores session data in a redis db. """

    def __init__(self, redisstore):
        self.redis = redisstore

    def get(self, key):
        """ load item from store. """
        value = self.redis.get(key)
        if value is None:
            raise KeyError(key)
        return value

    def put(self, key, value, ttl=None):
        """ save item to store. """
        if ttl:
            self.redis.setex(key, int(ttl), value)
        else:
            self.redis.set(key, value)

    def delete(self, key):
        """ delete item from store. """
        self.redis.delete(key)

    def list(self, prefix=None):
        """ return list of keys starting with prefix """
        if prefix:
            prefix += b'*'
        return list(self.redis.scan_iter(match=prefix))

class DictStore:
    """ Stores session data in a python dict. """

    def __init__(self):
        self.dict = {}

    def get(self, key):
        """ load item from store. """
        return self.dict[key]

    def put(self, key, value, ttl=None):
        """ save item to store. """
        self.dict[key] = value

    def delete(self, key):
        """ delete item from store. """
        try:
            del self.dict[key]
        except KeyError:
            pass

    def list(self, prefix=None):
        """ return list of keys starting with prefix """
        if prefix is None:
            return list(self.dict.keys())
        return [key for key in self.dict if key.startswith(prefix)]

class MailuSession(CallbackDict, SessionMixin):
    """ Custom flask session storage. """

    # default modified to false
    modified = False

    def __init__(self, key=None, app=None):

        self.app = app or flask.current_app

        initial = None

        key = want_bytes(key)
        if parsed := self.app.session_config.parse_key(key, self.app):
            try:
                initial = pickle.loads(app.session_store.get(key))
            except (KeyError, EOFError, pickle.UnpicklingError):
                # either the cookie was manipulated or we did not find the
                # session in the backend or the pickled data is invalid.
                # => start new session
                pass
            else:
                (self._uid, self._sid, self._created) = parsed
                self._key = key

        if initial is None:
            # start new session
            self.new = True
            self._uid = None
            self._sid = None
            self._created = self.app.session_config.gen_created()
            self._key = None

        def _on_update(obj):
            obj.modified = True

        CallbackDict.__init__(self, initial, _on_update)

    @property
    def saved(self):
        """ this reflects if the session was saved. """
        return self._key is not None

    @property
    def sid(self):
        """ this reflects the session's id. """
        if self._sid is None or self._uid is None or self._created is None:
            return None
        return b''.join([self._uid, self._sid, self._created])

    def destroy(self):
        """ destroy session for security reasons. """
        if 'webmail_token' in self:
            self.app.session_store.delete(self['webmail_token'])
        self.delete()

        self._uid = None
        self._sid = None
        self._created = None

        self.clear()

        self.modified = True
        self.new = False

    def regenerate(self):
        """ generate new id for session to avoid `session fixation`. """
        self.delete()
        self._sid = None
        self.modified = True

    def delete(self):
        """ Delete stored session. """
        if self.saved:
            self.app.session_store.delete(self._key)
            self._key = None

    def save(self):
        """ Save session to store. """
        set_cookie = False
        # set uid from dict data
        if self._uid is None:
            self._uid = self.app.session_config.gen_uid(self.get('_user_id', ''))

        # create new session id for new or regenerated sessions and force setting the cookie
        if self._sid is None:
            self._sid = self.app.session_config.gen_sid()
            set_cookie = True
            if 'webmail_token' in self:
                app.session_store.put(self['webmail_token'],
                        self.sid,
                        app.config['PERMANENT_SESSION_LIFETIME'],
                )

        # get new session key
        key = self.sid

        # delete old session if key has changed
        if key != self._key:
            self.delete()

        # save session
        self.app.session_store.put(
            key,
            pickle.dumps(dict(self)),
            app.config['SESSION_TIMEOUT'],
        )

        self._key = key

        self.new = False
        self.modified = False

        return set_cookie

class MailuSessionConfig:
    """ Stores sessions crypto config """

    # default size of session key parts
    uid_bits = 64 # default if SESSION_KEY_BITS is not set in config
    sid_bits = 128 # for now. must be multiple of 8!
    time_bits = 32 # for now. must be multiple of 8!

    def __init__(self, app=None):

        if app is None:
            app = flask.current_app

        bits = app.config.get('SESSION_KEY_BITS', self.uid_bits)
        if not 64 <= bits <= 256:
            raise ValueError('SESSION_KEY_BITS must be between 64 and 256!')

        uid_bytes = bits//8 + (bits%8>0)
        sid_bytes = self.sid_bits//8

        key = want_bytes(app.secret_key)

        self._hmac    = hmac.new(hmac.digest(key, b'SESSION_UID_HASH', digest='sha256'), digestmod='sha256')
        self._uid_len = uid_bytes
        self._uid_b64 = len(self._encode(bytes(uid_bytes)))
        self._sid_len = sid_bytes
        self._sid_b64 = len(self._encode(bytes(sid_bytes)))
        self._key_min = self._uid_b64 + self._sid_b64
        self._key_max = self._key_min + len(self._encode(bytes(self.time_bits//8)))

    def gen_sid(self):
        """ Generate random session id. """
        return self._encode(secrets.token_bytes(self._sid_len))

    def gen_uid(self, uid):
        """ Generate hashed user id part of session key. """
        _hmac = self._hmac.copy()
        _hmac.update(want_bytes(uid))
        return self._encode(_hmac.digest()[:self._uid_len])

    def gen_created(self, now=None):
        """ Generate base64 representation of creation time. """
        return self._encode(int(now or time.time()).to_bytes(8, byteorder='big').lstrip(b'\0'))

    def parse_key(self, key, app=None, now=None):
        """ Split key into sid, uid and creation time. """

        if not (isinstance(key, bytes) and self._key_min <= len(key) <= self._key_max):
            return None

        uid = key[:self._uid_b64]
        sid = key[self._uid_b64:self._key_min]
        crt = key[self._key_min:]

        # validate if parts are decodeable
        created = self._decode(crt)
        if created is None or self._decode(uid) is None or self._decode(sid) is None:
            return None

        # validate creation time
        if now is None:
            now = int(time.time())
        created = int.from_bytes(created, byteorder='big')
        if not created <= now <= created + app.config['PERMANENT_SESSION_LIFETIME']:
            return None

        return (uid, sid, crt)

    def _encode(self, value):
        return secrets.base64.urlsafe_b64encode(value).rstrip(b'=')

    def _decode(self, value):
        try:
            return secrets.base64.urlsafe_b64decode(value + b'='*(4-len(value)%4))
        except secrets.binascii.Error:
            return None

class MailuSessionInterface(SessionInterface):
    """ Custom flask session interface. """

    def open_session(self, app, request):
        """ Load or create session. """
        return MailuSession(request.cookies.get(app.config['SESSION_COOKIE_NAME'], None), app)

    def save_session(self, app, session, response):
        """ Save modified session. """

        # If the session is modified to be empty, remove the cookie.
        # If the session is empty, return without setting the cookie.
        if not session:
            if session.modified:
                session.delete()
                response.delete_cookie(
                    app.session_cookie_name,
                    domain=self.get_cookie_domain(app),
                    path=self.get_cookie_path(app),
                )
            return

        # Add a "Vary: Cookie" header if the session was accessed
        if session.accessed:
            response.vary.add('Cookie')

        # save session and update cookie if necessary
        if session.save():
            response.set_cookie(
                app.session_cookie_name,
                session.sid,
                expires=datetime.now()+timedelta(seconds=app.config['PERMANENT_SESSION_LIFETIME']),
                httponly=self.get_cookie_httponly(app),
                domain=self.get_cookie_domain(app),
                path=self.get_cookie_path(app),
                secure=self.get_cookie_secure(app),
                samesite=self.get_cookie_samesite(app)
            )

class MailuSessionExtension:
    """ Server side session handling """

    @staticmethod
    def cleanup_sessions(app=None):
        """ Remove invalid or expired sessions. """

        app = app or flask.current_app
        now = int(time.time())

        count = 0
        for key in app.session_store.list():
            if not app.session_config.parse_key(key, app, now=now):
                app.session_store.delete(key)
                count += 1

        return count

    @staticmethod
    def prune_sessions(uid=None, keep=None, app=None):
        """ Remove sessions
            uid: remove all sessions (NONE) or sessions belonging to a specific user
            keep: keep listed sessions
        """

        keep = keep or set()
        app = app or flask.current_app

        prefix = None if uid is None else app.session_config.gen_uid(uid)

        count = 0
        for key in app.session_store.list(prefix):
            if key not in keep:
                app.session_store.delete(key)
                count += 1

        return count

    def init_app(self, app):
        """ Replace session management of application. """

        if app.config.get('MEMORY_SESSIONS'):
            # in-memory session store for use in development
            app.session_store = DictStore()

        else:
            # redis-based session store for use in production
            app.session_store = RedisStore(
                redis.StrictRedis().from_url(app.config['SESSION_STORAGE_URL'])
            )

            # clean expired sessions oonce on first use in case lifetime was changed
            def cleaner():
                with cleaned.get_lock():
                    if not cleaned.value:
                        cleaned.value = True
                        app.logger.info('cleaning session store')
                        MailuSessionExtension.cleanup_sessions(app)

            app.before_first_request(cleaner)

        app.session_config = MailuSessionConfig(app)
        app.session_interface = MailuSessionInterface()

cleaned = Value('i', False)
session = MailuSessionExtension()

# this is used by the webmail to authenticate IMAP/SMTP
def verify_temp_token(email, token):
    try:
        if token.startswith('token-'):
            sessid = app.session_store.get(token)
            if sessid:
                session = MailuSession(sessid, app)
                if session.get('_user_id', '') == email:
                    return True
    except:
        pass

def gen_temp_token(email, session):
    token = session.get('webmail_token', 'token-'+secrets.token_urlsafe())
    session['webmail_token'] = token
    app.session_store.put(token,
            session.sid,
            app.config['PERMANENT_SESSION_LIFETIME'],
    )
    return token
<commit_msg>syntax<commit_after>""" Mailu admin app utilities
"""

try:
    import cPickle as pickle
except ImportError:
    import pickle

import dns.resolver
import dns.exception
import dns.flags
import dns.rdtypes
import dns.rdatatype
import dns.rdataclass

import hmac
import secrets
import time

from multiprocessing import Value
from mailu import limiter
from flask import current_app as app

import flask
import flask_login
import flask_migrate
import flask_babel
import ipaddress
import redis

from datetime import datetime, timedelta
from flask.sessions import SessionMixin, SessionInterface
from itsdangerous.encoding import want_bytes
from werkzeug.datastructures import CallbackDict
from werkzeug.middleware.proxy_fix import ProxyFix

# Login configuration
login = flask_login.LoginManager()
login.login_view = "sso.login"

@login.unauthorized_handler
def handle_needs_login():
    """ redirect unauthorized requests to login page """
    return flask.redirect(
        flask.url_for('sso.login')
    )

# DNS stub configured to do DNSSEC enabled queries
resolver = dns.resolver.Resolver()
resolver.use_edns(0, 0, 1232)
resolver.flags = dns.flags.AD | dns.flags.RD

def has_dane_record(domain, timeout=10):
    try:
        result = resolver.query(f'_25._tcp.{domain}', dns.rdatatype.TLSA,dns.rdataclass.IN, lifetime=timeout)
        if result.response.flags & dns.flags.AD:
            for record in result:
                if isinstance(record, dns.rdtypes.ANY.TLSA.TLSA):
                    record.validate()
                    if record.usage in [2,3] and record.selector in [0,1] and record.mtype in [0,1,2]:
                        return True
    except dns.resolver.NoNameservers:
        # If the DNSSEC data is invalid and the DNS resolver is DNSSEC enabled
        # we will receive this non-specific exception. The safe behaviour is to
        # accept to defer the email.
        app.logger.warn(f'Unable to lookup the TLSA record for {domain}. Is the DNSSEC zone okay on https://dnsviz.net/d/{domain}/dnssec/?')
        return app.config['DEFER_ON_TLS_ERROR']
    except dns.exception.Timeout:
        app.logger.warn(f'Timeout while resolving the TLSA record for {domain} ({timeout}s).')
    except (dns.resolver.NXDOMAIN, dns.name.EmptyLabel):
        pass # this is expected, not TLSA record is fine
    except Exception as e:
        app.logger.info(f'Error while looking up the TLSA record for {domain} {e}')
        pass

# Rate limiter
limiter = limiter.LimitWraperFactory()

def extract_network_from_ip(ip):
    n = ipaddress.ip_network(ip)
    if n.version == 4:
        return str(n.supernet(prefixlen_diff=(32-app.config["AUTH_RATELIMIT_IP_V4_MASK"])).network_address)
    else:
        return str(n.supernet(prefixlen_diff=(128-app.config["AUTH_RATELIMIT_IP_V6_MASK"])).network_address)

def is_exempt_from_ratelimits(ip):
    ip = ipaddress.ip_address(ip)
    return any(ip in cidr for cidr in app.config['AUTH_RATELIMIT_EXEMPTION'])

# Application translation
babel = flask_babel.Babel()

@babel.localeselector
def get_locale():
    """ selects locale for translation """
    if not app.config['SESSION_COOKIE_NAME'] in flask.request.cookies:
        return flask.request.accept_languages.best_match(app.config.translations.keys())
    language = flask.session.get('language')
    if not language in app.config.translations:
        language = flask.request.accept_languages.best_match(app.config.translations.keys())
        flask.session['language'] = language
    return language

# Proxy fixer
class PrefixMiddleware(object):
    """ fix proxy headers """
    def __init__(self):
        self.app = None

    def __call__(self, environ, start_response):
        return self.app(environ, start_response)

    def init_app(self, app):
        self.app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1)
        app.wsgi_app = self

proxy = PrefixMiddleware()

# Data migrate
migrate = flask_migrate.Migrate()

# session store (inspired by https://github.com/mbr/flask-kvsession)
class RedisStore:
    """ Stores session data in a redis db. """

    def __init__(self, redisstore):
        self.redis = redisstore

    def get(self, key):
        """ load item from store. """
        value = self.redis.get(key)
        if value is None:
            raise KeyError(key)
        return value

    def put(self, key, value, ttl=None):
        """ save item to store. """
        if ttl:
            self.redis.setex(key, int(ttl), value)
        else:
            self.redis.set(key, value)

    def delete(self, key):
        """ delete item from store. """
        self.redis.delete(key)

    def list(self, prefix=None):
        """ return list of keys starting with prefix """
        if prefix:
            prefix += b'*'
        return list(self.redis.scan_iter(match=prefix))

class DictStore:
    """ Stores session data in a python dict. """

    def __init__(self):
        self.dict = {}

    def get(self, key):
        """ load item from store. """
        return self.dict[key]

    def put(self, key, value, ttl=None):
        """ save item to store. """
        self.dict[key] = value

    def delete(self, key):
        """ delete item from store. """
        try:
            del self.dict[key]
        except KeyError:
            pass

    def list(self, prefix=None):
        """ return list of keys starting with prefix """
        if prefix is None:
            return list(self.dict.keys())
        return [key for key in self.dict if key.startswith(prefix)]

class MailuSession(CallbackDict, SessionMixin):
    """ Custom flask session storage. """

    # default modified to false
    modified = False

    def __init__(self, key=None, app=None):

        self.app = app or flask.current_app

        initial = None

        key = want_bytes(key)
        if parsed := self.app.session_config.parse_key(key, self.app):
            try:
                initial = pickle.loads(app.session_store.get(key))
            except (KeyError, EOFError, pickle.UnpicklingError):
                # either the cookie was manipulated or we did not find the
                # session in the backend or the pickled data is invalid.
                # => start new session
                pass
            else:
                (self._uid, self._sid, self._created) = parsed
                self._key = key

        if initial is None:
            # start new session
            self.new = True
            self._uid = None
            self._sid = None
            self._created = self.app.session_config.gen_created()
            self._key = None

        def _on_update(obj):
            obj.modified = True

        CallbackDict.__init__(self, initial, _on_update)

    @property
    def saved(self):
        """ this reflects if the session was saved. """
        return self._key is not None

    @property
    def sid(self):
        """ this reflects the session's id. """
        if self._sid is None or self._uid is None or self._created is None:
            return None
        return b''.join([self._uid, self._sid, self._created])

    def destroy(self):
        """ destroy session for security reasons. """
        if 'webmail_token' in self:
            self.app.session_store.delete(self['webmail_token'])
        self.delete()

        self._uid = None
        self._sid = None
        self._created = None

        self.clear()

        self.modified = True
        self.new = False

    def regenerate(self):
        """ generate new id for session to avoid `session fixation`. """
        self.delete()
        self._sid = None
        self.modified = True

    def delete(self):
        """ Delete stored session. """
        if self.saved:
            self.app.session_store.delete(self._key)
            self._key = None

    def save(self):
        """ Save session to store. """
        set_cookie = False
        # set uid from dict data
        if self._uid is None:
            self._uid = self.app.session_config.gen_uid(self.get('_user_id', ''))

        # create new session id for new or regenerated sessions and force setting the cookie
        if self._sid is None:
            self._sid = self.app.session_config.gen_sid()
            set_cookie = True
            if 'webmail_token' in self:
                app.session_store.put(self['webmail_token'],
                        self.sid,
                        app.config['PERMANENT_SESSION_LIFETIME'],
                )

        # get new session key
        key = self.sid

        # delete old session if key has changed
        if key != self._key:
            self.delete()

        # save session
        self.app.session_store.put(
            key,
            pickle.dumps(dict(self)),
            app.config['SESSION_TIMEOUT'],
        )

        self._key = key

        self.new = False
        self.modified = False

        return set_cookie

class MailuSessionConfig:
    """ Stores sessions crypto config """

    # default size of session key parts
    uid_bits = 64 # default if SESSION_KEY_BITS is not set in config
    sid_bits = 128 # for now. must be multiple of 8!
    time_bits = 32 # for now. must be multiple of 8!

    def __init__(self, app=None):

        if app is None:
            app = flask.current_app

        bits = app.config.get('SESSION_KEY_BITS', self.uid_bits)
        if not 64 <= bits <= 256:
            raise ValueError('SESSION_KEY_BITS must be between 64 and 256!')

        uid_bytes = bits//8 + (bits%8>0)
        sid_bytes = self.sid_bits//8

        key = want_bytes(app.secret_key)

        self._hmac    = hmac.new(hmac.digest(key, b'SESSION_UID_HASH', digest='sha256'), digestmod='sha256')
        self._uid_len = uid_bytes
        self._uid_b64 = len(self._encode(bytes(uid_bytes)))
        self._sid_len = sid_bytes
        self._sid_b64 = len(self._encode(bytes(sid_bytes)))
        self._key_min = self._uid_b64 + self._sid_b64
        self._key_max = self._key_min + len(self._encode(bytes(self.time_bits//8)))

    def gen_sid(self):
        """ Generate random session id. """
        return self._encode(secrets.token_bytes(self._sid_len))

    def gen_uid(self, uid):
        """ Generate hashed user id part of session key. """
        _hmac = self._hmac.copy()
        _hmac.update(want_bytes(uid))
        return self._encode(_hmac.digest()[:self._uid_len])

    def gen_created(self, now=None):
        """ Generate base64 representation of creation time. """
        return self._encode(int(now or time.time()).to_bytes(8, byteorder='big').lstrip(b'\0'))

    def parse_key(self, key, app=None, now=None):
        """ Split key into sid, uid and creation time. """

        if not (isinstance(key, bytes) and self._key_min <= len(key) <= self._key_max):
            return None

        uid = key[:self._uid_b64]
        sid = key[self._uid_b64:self._key_min]
        crt = key[self._key_min:]

        # validate if parts are decodeable
        created = self._decode(crt)
        if created is None or self._decode(uid) is None or self._decode(sid) is None:
            return None

        # validate creation time
        if now is None:
            now = int(time.time())
        created = int.from_bytes(created, byteorder='big')
        if not created <= now <= created + app.config['PERMANENT_SESSION_LIFETIME']:
            return None

        return (uid, sid, crt)

    def _encode(self, value):
        return secrets.base64.urlsafe_b64encode(value).rstrip(b'=')

    def _decode(self, value):
        try:
            return secrets.base64.urlsafe_b64decode(value + b'='*(4-len(value)%4))
        except secrets.binascii.Error:
            return None

class MailuSessionInterface(SessionInterface):
    """ Custom flask session interface. """

    def open_session(self, app, request):
        """ Load or create session. """
        return MailuSession(request.cookies.get(app.config['SESSION_COOKIE_NAME'], None), app)

    def save_session(self, app, session, response):
        """ Save modified session. """

        # If the session is modified to be empty, remove the cookie.
        # If the session is empty, return without setting the cookie.
        if not session:
            if session.modified:
                session.delete()
                response.delete_cookie(
                    app.session_cookie_name,
                    domain=self.get_cookie_domain(app),
                    path=self.get_cookie_path(app),
                )
            return

        # Add a "Vary: Cookie" header if the session was accessed
        if session.accessed:
            response.vary.add('Cookie')

        # save session and update cookie if necessary
        if session.save():
            response.set_cookie(
                app.session_cookie_name,
                session.sid,
                expires=datetime.now()+timedelta(seconds=app.config['PERMANENT_SESSION_LIFETIME']),
                httponly=self.get_cookie_httponly(app),
                domain=self.get_cookie_domain(app),
                path=self.get_cookie_path(app),
                secure=self.get_cookie_secure(app),
                samesite=self.get_cookie_samesite(app)
            )

class MailuSessionExtension:
    """ Server side session handling """

    @staticmethod
    def cleanup_sessions(app=None):
        """ Remove invalid or expired sessions. """

        app = app or flask.current_app
        now = int(time.time())

        count = 0
        for key in app.session_store.list():
            if not app.session_config.parse_key(key, app, now=now):
                app.session_store.delete(key)
                count += 1

        return count

    @staticmethod
    def prune_sessions(uid=None, keep=None, app=None):
        """ Remove sessions
            uid: remove all sessions (NONE) or sessions belonging to a specific user
            keep: keep listed sessions
        """

        keep = keep or set()
        app = app or flask.current_app

        prefix = None if uid is None else app.session_config.gen_uid(uid)

        count = 0
        for key in app.session_store.list(prefix):
            if key not in keep:
                app.session_store.delete(key)
                count += 1

        return count

    def init_app(self, app):
        """ Replace session management of application. """

        if app.config.get('MEMORY_SESSIONS'):
            # in-memory session store for use in development
            app.session_store = DictStore()

        else:
            # redis-based session store for use in production
            app.session_store = RedisStore(
                redis.StrictRedis().from_url(app.config['SESSION_STORAGE_URL'])
            )

            # clean expired sessions oonce on first use in case lifetime was changed
            def cleaner():
                with cleaned.get_lock():
                    if not cleaned.value:
                        cleaned.value = True
                        app.logger.info('cleaning session store')
                        MailuSessionExtension.cleanup_sessions(app)

            app.before_first_request(cleaner)

        app.session_config = MailuSessionConfig(app)
        app.session_interface = MailuSessionInterface()

cleaned = Value('i', False)
session = MailuSessionExtension()

# this is used by the webmail to authenticate IMAP/SMTP
def verify_temp_token(email, token):
    try:
        if token.startswith('token-'):
            if sessid := app.session_store.get(token)
                session = MailuSession(sessid, app)
                if session.get('_user_id', '') == email:
                    return True
    except:
        pass

def gen_temp_token(email, session):
    token = session.get('webmail_token', 'token-'+secrets.token_urlsafe())
    session['webmail_token'] = token
    app.session_store.put(token,
            session.sid,
            app.config['PERMANENT_SESSION_LIFETIME'],
    )
    return token
<|endoftext|>


<commit_before>"""Analysis module for Databench."""

import os
import json
import time
import gevent
import logging
import zipstream
import subprocess
import geventwebsocket
import zmq.green as zmq
from flask import Blueprint, render_template, Response

class Analysis(object):
    """Databench's analysis class.

    This contains the analysis code. Every browser connection corresponds to
    and instance of this class.

    **Incoming messages** are captured by specifying a class method starting
    with ``on_`` followed by the signal name. To capture the frontend signal
    ``run`` that is emitted with the JavaScript code

    .. code-block:: js

        // on the JavaScript frontend
        databench.emit('run', {my_param: 'helloworld'});

    use

    .. code-block:: python

        # here in Python
        def on_run(self, my_param):

    here. The entries of a dictionary will be used as keyword arguments in the
    function call; as in this example. If the emitted message is an array,
    the entries will be used as positional arguments in the function call.
    If the message is neither of type ``list`` nor ``dict`` (for example a
    plain ``string`` or ``float``), the function will be called with that
    message as its first parameter.

    **Outgoing messages** are sent using ``emit(signal_name, message)``.
    For example, use

    .. code-block:: python

        self.emit('result', {'msg': 'done'})

    to send the signal ``result`` with the message ``{'msg': 'done'}`` to
    the frontend.

    """

    def __init__(self):
        pass

    def set_emit_fn(self, emit_fn):
        """Sets what the emit function for this analysis will be."""
        self.emit = emit_fn

    def set_emit_action_fn(self, emit_action_fn):
        """Sets what the emit action function for this analysis will be."""
        self.emit_action = emit_action_fn

    """Events."""

    def onall(self, message_data):
        logging.debug('onall called.')

    def on_connect(self):
        logging.debug('on_connect called.')

    def on_disconnect(self):
        logging.debug('on_disconnect called.')

class Meta(object):
    """
    Args:
        name (str): Name of this analysis. If ``signals`` is not specified,
            this also becomes the namespace for the WebSocket connection and
            has to match the frontend's :js:class:`Databench` ``name``.
        import_name (str): Usually the file name ``__name__`` where this
            analysis is instantiated.
        description (str): Usually the ``__doc__`` string of the analysis.
        analysis_class (:class:`databench.Analysis`): Object
            that should be instantiated for every new websocket connection.

    For standard use cases, you don't have to modify this class. However,
    If you want to serve more than the ``index.html`` page, say a
    ``details.html`` page, you can derive from this class and add this
    to the constructor

    .. code-block:: python

        self.blueprint.add_url_rule('/details.html', 'render_details',
                                    self.render_details)

    and add a new method to the class

    .. code-block:: python

        def render_details(self):
            return render_template(
                self.name+'/details.html',
                analysis_description=self.description
            )

    and create the file ``details.html`` similar to ``index.html``.

    """

    all_instances = []

    def __init__(
            self,
            name,
            import_name,
            description,
            analysis_class,
    ):
        Meta.all_instances.append(self)
        self.show_in_index = True

        self.name = name
        self.import_name = import_name
        self.description = description
        self.analysis_class = analysis_class

        analyses_path = os.getcwd()+'/'+'analyses'
        if not os.path.exists(analyses_path):
            analyses_path = os.getcwd()+'/'+'analyses_packaged'
        self.analyses_path = analyses_path

        # detect whether thumbnail.png is present
        if os.path.isfile(analyses_path+'/'+self.name+'/thumbnail.png'):
            self.thumbnail = 'thumbnail.png'

        self.blueprint = Blueprint(
            name,
            import_name,
            template_folder=analyses_path,
            static_folder=analyses_path+'/'+self.name,
            static_url_path='/static',
        )
        self.blueprint.add_url_rule('/', 'render_index', self.render_index)
        self.blueprint.add_url_rule('/'+name+'.zip', 'zip_analysis',
                                    self.zip_analysis, methods=['GET'])

        self.sockets = None

    def render_index(self):
        """Renders the main analysis frontend template."""
        return render_template(
            self.name+'/index.html',
            analysis_name=self.name,
            analysis_description=self.description,
        )

    def zip_analysis(self):
        def generator():
            z = zipstream.ZipFile(mode='w',
                                  compression=zipstream.ZIP_DEFLATED)

            # find all analysis files
            folder = self.analyses_path+'/'+self.name
            for root, dirnames, filenames in os.walk(folder):
                invisible_dirs = [d for d in dirnames if d[0] == '.']
                for d in invisible_dirs:
                    dirnames.remove(d)
                for filename in filenames:
                    if filename[0] == '.':
                        continue
                    if filename[-4:] == '.pyc':
                        continue

                    # add the file to zipstream
                    fullname = os.path.join(root, filename)
                    arcname = fullname.replace(self.analyses_path+'/', '')
                    z.write(fullname, arcname=arcname)

            # add requirements.txt if present
            if os.path.isfile(self.analyses_path+'/requirements.txt'):
                z.write(self.analyses_path+'/requirements.txt')

            for chunk in z:
                yield chunk

        response = Response(generator(), mimetype='application/zip')
        response.headers['Content-Disposition'] = \
            'attachment; filename={}'.format(self.name+'.zip')
        return response

    def wire_sockets(self, sockets, url_prefix=''):
        self.sockets = sockets
        self.sockets.add_url_rule(url_prefix+'/ws', 'ws_serve', self.ws_serve)

    def instantiate_analysis_class(self):
        return self.analysis_class()

    def ws_serve(self, ws):
        """Handle a new websocket connection."""
        logging.debug('ws_serve()')

        def emit(signal, message):
            try:
                ws.send(json.dumps({'signal': signal, 'message': message}))
            except geventwebsocket.WebSocketError, e:
                logging.info('websocket closed. could not send: '+signal +
                             ' -- '+str(message))

        def emit_action(action_id, status):
            try:
                ws.send(json.dumps({'action_id': action_id, 'status': status}))
            except geventwebsocket.WebSocketError, e:
                logging.info('websocket could not send action: '+action_id +
                             ' -- '+str(status))

        analysis_instance = self.instantiate_analysis_class()
        logging.debug("analysis instantiated")
        analysis_instance.set_emit_fn(emit)
        analysis_instance.set_emit_action_fn(emit_action)
        greenlets = []
        greenlets.append(gevent.Greenlet.spawn(
            analysis_instance.on_connect
        ))

        def process_message(message):
            if message is None:
                logging.debug('empty message received.')
                return

            message_data = json.loads(message)
            analysis_instance.onall(message_data)
            if 'signal' in message_data and 'message' in message_data:
                fn_name = 'on_'+message_data['signal']
                if hasattr(self.analysis_class, fn_name):
                    logging.debug('calling '+fn_name)

                    # detect whether an action_id should be registered
                    action_id = None
                    if '__action_id' in message_data['message']:
                        action_id = message_data['message']['__action_id']
                        del message_data['message']['__action_id']

                    # wrap the action
                    def spawn_action():
                        action_id_local = action_id
                        if action_id_local:
                            emit_action(action_id_local, 'start')

                        # Check whether this is a list (positional arguments)
                        # or a dictionary (keyword arguments).
                        if isinstance(message_data['message'], list):
                            getattr(analysis_instance, fn_name)(
                                *message_data['message']
                            )
                        elif isinstance(message_data['message'], dict):
                            getattr(analysis_instance, fn_name)(
                                **message_data['message']
                            )
                        else:
                            getattr(analysis_instance, fn_name)(
                                message_data['message']
                            )

                        if action_id_local:
                            emit_action(action_id_local, 'end')

                    # every 'on_' is processed in a separate greenlet
                    greenlets.append(gevent.Greenlet.spawn(spawn_action))
                else:
                    logging.warning('frontend wants to call '
                                    'on_'+message_data['signal']+' which is '
                                    'not in the Analysis class.')
            else:
                logging.info('message not processed: '+message)

        while True:
            try:
                message = ws.receive()
                logging.debug('received message: '+str(message))
                process_message(message)
            except geventwebsocket.WebSocketError, e:
                break

        # disconnected
        logging.debug("disconnecting analysis instance")
        gevent.killall(greenlets)
        analysis_instance.on_disconnect()

class AnalysisZMQ(Analysis):

    def __init__(self, namespace, analysis_id, zmq_publish):
        self.namespace = namespace
        self.analysis_id = analysis_id
        self.zmq_publish = zmq_publish

    def onall(self, message_data):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': message_data,
        }
        self.zmq_publish.send_json(msg)
        logging.debug('onall called with: '+str(msg))

    def on_connect(self):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': {'signal': 'connect', 'message': {}},
        }
        self.zmq_publish.send_json(msg)
        logging.debug('on_connect called')

    def on_disconnect(self):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': {'signal': 'disconnect', 'message': {}},
        }
        self.zmq_publish.send_json(msg)
        logging.debug('on_disconnect called')

class MetaZMQ(Meta):
    """A Meta class that pipes all messages to ZMQ and back.

    The entire ZMQ interface of Databench is defined here and in
    :class`AnalysisZMQ`.

    """

    def __init__(
            self,
            name,
            import_name,
            description,

            executable,
            zmq_publish,
            port_subscribe=None,
    ):
        Meta.__init__(self, name, import_name, description, AnalysisZMQ)

        self.zmq_publish = zmq_publish

        self.zmq_analysis_id = 0
        self.zmq_analyses = {}
        self.zmq_confirmed = False

        # check whether we have to determine port_subscribe ourselves first
        if port_subscribe is None:
            socket = zmq.Context().socket(zmq.PUB)
            port_subscribe = socket.bind_to_random_port(
                'tcp://127.0.0.1',
                min_port=3000, max_port=9000,
            )
            socket.unbind('tcp://127.0.0.1:'+str(port_subscribe))
            logging.debug('determined: port_subscribe='+str(port_subscribe))

        # zmq subscription to listen for messages from backend
        logging.debug('main listening on port: '+str(port_subscribe))
        self.zmq_sub = zmq.Context().socket(zmq.SUB)
        self.zmq_sub.connect('tcp://127.0.0.1:'+str(port_subscribe))
        self.zmq_sub.setsockopt(zmq.SUBSCRIBE, '')

        # @copy_current_request_context
        def zmq_listener():
            while True:
                msg = self.zmq_sub.recv_json()
                self.zmq_confirmed = True
                logging.debug('main ('+') received '
                              'msg: '+str(msg))

                if 'description' in msg:
                    self.description = msg['description']

                if '__analysis_id' in msg and \
                   msg['__analysis_id'] in self.zmq_analyses:
                    analysis = self.zmq_analyses[msg['__analysis_id']]
                    del msg['__analysis_id']

                    if 'message' in msg and \
                       'signal' in msg['message'] and \
                       'message' in msg['message']:
                        analysis.emit(msg['message']['signal'],
                                      msg['message']['message'])
                    elif 'action_id' in msg and \
                         'status' in msg:
                        analysis.emit_action(msg['action_id'],
                                             msg['status'])
                    else:
                        logging.debug('dont understand this message: ' +
                                      str(msg))
                else:
                    logging.debug('__analysis_id not in message or '
                                  'AnalysisZMQ with that id not found.')
        self.zmq_listener = gevent.Greenlet.spawn(zmq_listener)

        # launch the language kernel process
        self.kernel_process = subprocess.Popen(executable, shell=False)

        # init language kernel
        def sending_init():
            while not self.zmq_confirmed:
                logging.debug('init kernel '+self.name+' to publish on '
                              'port '+str(port_subscribe))
                self.zmq_publish.send_json({
                    '__databench_namespace': self.name,
                    'publish_on_port': port_subscribe,
                })
                time.sleep(0.1)
        gevent.Greenlet.spawn(sending_init)

    def __del__(self):
        self.kernel_process.terminate()
        self.kernel_process.kill()

    def instantiate_analysis_class(self):
        self.zmq_analysis_id += 1
        i = self.analysis_class(self.name,
                                self.zmq_analysis_id,
                                self.zmq_publish)
        self.zmq_analyses[self.zmq_analysis_id] = i
        return i
<commit_msg>minor doc<commit_after>"""Analysis module for Databench."""

import os
import json
import time
import gevent
import logging
import zipstream
import subprocess
import geventwebsocket
import zmq.green as zmq
from flask import Blueprint, render_template, Response

class Analysis(object):
    """Databench's analysis class.

    This contains the analysis code. Every browser connection corresponds to
    and instance of this class.

    **Incoming messages** are captured by specifying a class method starting
    with ``on_`` followed by the signal name. To capture the frontend signal
    ``run`` that is emitted with the JavaScript code

    .. code-block:: js

        // on the JavaScript frontend
        databench.emit('run', {my_param: 'helloworld'});

    use

    .. code-block:: python

        # here in Python
        def on_run(self, my_param):

    here. The entries of a dictionary will be used as keyword arguments in the
    function call; as in this example. If the emitted message is an array,
    the entries will be used as positional arguments in the function call.
    If the message is neither of type ``list`` nor ``dict`` (for example a
    plain ``string`` or ``float``), the function will be called with that
    as its first parameter.

    **Outgoing messages** are sent using ``emit(signal_name, message)``.
    For example, use

    .. code-block:: python

        self.emit('result', {'msg': 'done'})

    to send the signal ``result`` with the message ``{'msg': 'done'}`` to
    the frontend.

    """

    def __init__(self):
        pass

    def set_emit_fn(self, emit_fn):
        """Sets what the emit function for this analysis will be."""
        self.emit = emit_fn

    def set_emit_action_fn(self, emit_action_fn):
        """Sets what the emit action function for this analysis will be."""
        self.emit_action = emit_action_fn

    """Events."""

    def onall(self, message_data):
        logging.debug('onall called.')

    def on_connect(self):
        logging.debug('on_connect called.')

    def on_disconnect(self):
        logging.debug('on_disconnect called.')

class Meta(object):
    """
    Args:
        name (str): Name of this analysis. If ``signals`` is not specified,
            this also becomes the namespace for the WebSocket connection and
            has to match the frontend's :js:class:`Databench` ``name``.
        import_name (str): Usually the file name ``__name__`` where this
            analysis is instantiated.
        description (str): Usually the ``__doc__`` string of the analysis.
        analysis_class (:class:`databench.Analysis`): Object
            that should be instantiated for every new websocket connection.

    For standard use cases, you don't have to modify this class. However,
    If you want to serve more than the ``index.html`` page, say a
    ``details.html`` page, you can derive from this class and add this
    to the constructor

    .. code-block:: python

        self.blueprint.add_url_rule('/details.html', 'render_details',
                                    self.render_details)

    and add a new method to the class

    .. code-block:: python

        def render_details(self):
            return render_template(
                self.name+'/details.html',
                analysis_description=self.description
            )

    and create the file ``details.html`` similar to ``index.html``.

    """

    all_instances = []

    def __init__(
            self,
            name,
            import_name,
            description,
            analysis_class,
    ):
        Meta.all_instances.append(self)
        self.show_in_index = True

        self.name = name
        self.import_name = import_name
        self.description = description
        self.analysis_class = analysis_class

        analyses_path = os.getcwd()+'/'+'analyses'
        if not os.path.exists(analyses_path):
            analyses_path = os.getcwd()+'/'+'analyses_packaged'
        self.analyses_path = analyses_path

        # detect whether thumbnail.png is present
        if os.path.isfile(analyses_path+'/'+self.name+'/thumbnail.png'):
            self.thumbnail = 'thumbnail.png'

        self.blueprint = Blueprint(
            name,
            import_name,
            template_folder=analyses_path,
            static_folder=analyses_path+'/'+self.name,
            static_url_path='/static',
        )
        self.blueprint.add_url_rule('/', 'render_index', self.render_index)
        self.blueprint.add_url_rule('/'+name+'.zip', 'zip_analysis',
                                    self.zip_analysis, methods=['GET'])

        self.sockets = None

    def render_index(self):
        """Renders the main analysis frontend template."""
        return render_template(
            self.name+'/index.html',
            analysis_name=self.name,
            analysis_description=self.description,
        )

    def zip_analysis(self):
        def generator():
            z = zipstream.ZipFile(mode='w',
                                  compression=zipstream.ZIP_DEFLATED)

            # find all analysis files
            folder = self.analyses_path+'/'+self.name
            for root, dirnames, filenames in os.walk(folder):
                invisible_dirs = [d for d in dirnames if d[0] == '.']
                for d in invisible_dirs:
                    dirnames.remove(d)
                for filename in filenames:
                    if filename[0] == '.':
                        continue
                    if filename[-4:] == '.pyc':
                        continue

                    # add the file to zipstream
                    fullname = os.path.join(root, filename)
                    arcname = fullname.replace(self.analyses_path+'/', '')
                    z.write(fullname, arcname=arcname)

            # add requirements.txt if present
            if os.path.isfile(self.analyses_path+'/requirements.txt'):
                z.write(self.analyses_path+'/requirements.txt')

            for chunk in z:
                yield chunk

        response = Response(generator(), mimetype='application/zip')
        response.headers['Content-Disposition'] = \
            'attachment; filename={}'.format(self.name+'.zip')
        return response

    def wire_sockets(self, sockets, url_prefix=''):
        self.sockets = sockets
        self.sockets.add_url_rule(url_prefix+'/ws', 'ws_serve', self.ws_serve)

    def instantiate_analysis_class(self):
        return self.analysis_class()

    def ws_serve(self, ws):
        """Handle a new websocket connection."""
        logging.debug('ws_serve()')

        def emit(signal, message):
            try:
                ws.send(json.dumps({'signal': signal, 'message': message}))
            except geventwebsocket.WebSocketError, e:
                logging.info('websocket closed. could not send: '+signal +
                             ' -- '+str(message))

        def emit_action(action_id, status):
            try:
                ws.send(json.dumps({'action_id': action_id, 'status': status}))
            except geventwebsocket.WebSocketError, e:
                logging.info('websocket could not send action: '+action_id +
                             ' -- '+str(status))

        analysis_instance = self.instantiate_analysis_class()
        logging.debug("analysis instantiated")
        analysis_instance.set_emit_fn(emit)
        analysis_instance.set_emit_action_fn(emit_action)
        greenlets = []
        greenlets.append(gevent.Greenlet.spawn(
            analysis_instance.on_connect
        ))

        def process_message(message):
            if message is None:
                logging.debug('empty message received.')
                return

            message_data = json.loads(message)
            analysis_instance.onall(message_data)
            if 'signal' in message_data and 'message' in message_data:
                fn_name = 'on_'+message_data['signal']
                if hasattr(self.analysis_class, fn_name):
                    logging.debug('calling '+fn_name)

                    # detect whether an action_id should be registered
                    action_id = None
                    if '__action_id' in message_data['message']:
                        action_id = message_data['message']['__action_id']
                        del message_data['message']['__action_id']

                    # wrap the action
                    def spawn_action():
                        action_id_local = action_id
                        if action_id_local:
                            emit_action(action_id_local, 'start')

                        # Check whether this is a list (positional arguments)
                        # or a dictionary (keyword arguments).
                        if isinstance(message_data['message'], list):
                            getattr(analysis_instance, fn_name)(
                                *message_data['message']
                            )
                        elif isinstance(message_data['message'], dict):
                            getattr(analysis_instance, fn_name)(
                                **message_data['message']
                            )
                        else:
                            getattr(analysis_instance, fn_name)(
                                message_data['message']
                            )

                        if action_id_local:
                            emit_action(action_id_local, 'end')

                    # every 'on_' is processed in a separate greenlet
                    greenlets.append(gevent.Greenlet.spawn(spawn_action))
                else:
                    logging.warning('frontend wants to call '
                                    'on_'+message_data['signal']+' which is '
                                    'not in the Analysis class.')
            else:
                logging.info('message not processed: '+message)

        while True:
            try:
                message = ws.receive()
                logging.debug('received message: '+str(message))
                process_message(message)
            except geventwebsocket.WebSocketError, e:
                break

        # disconnected
        logging.debug("disconnecting analysis instance")
        gevent.killall(greenlets)
        analysis_instance.on_disconnect()

class AnalysisZMQ(Analysis):

    def __init__(self, namespace, analysis_id, zmq_publish):
        self.namespace = namespace
        self.analysis_id = analysis_id
        self.zmq_publish = zmq_publish

    def onall(self, message_data):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': message_data,
        }
        self.zmq_publish.send_json(msg)
        logging.debug('onall called with: '+str(msg))

    def on_connect(self):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': {'signal': 'connect', 'message': {}},
        }
        self.zmq_publish.send_json(msg)
        logging.debug('on_connect called')

    def on_disconnect(self):
        msg = {
            '__analysis_id': self.analysis_id,
            '__databench_namespace': self.namespace,
            'message': {'signal': 'disconnect', 'message': {}},
        }
        self.zmq_publish.send_json(msg)
        logging.debug('on_disconnect called')

class MetaZMQ(Meta):
    """A Meta class that pipes all messages to ZMQ and back.

    The entire ZMQ interface of Databench is defined here and in
    :class`AnalysisZMQ`.

    """

    def __init__(
            self,
            name,
            import_name,
            description,

            executable,
            zmq_publish,
            port_subscribe=None,
    ):
        Meta.__init__(self, name, import_name, description, AnalysisZMQ)

        self.zmq_publish = zmq_publish

        self.zmq_analysis_id = 0
        self.zmq_analyses = {}
        self.zmq_confirmed = False

        # check whether we have to determine port_subscribe ourselves first
        if port_subscribe is None:
            socket = zmq.Context().socket(zmq.PUB)
            port_subscribe = socket.bind_to_random_port(
                'tcp://127.0.0.1',
                min_port=3000, max_port=9000,
            )
            socket.unbind('tcp://127.0.0.1:'+str(port_subscribe))
            logging.debug('determined: port_subscribe='+str(port_subscribe))

        # zmq subscription to listen for messages from backend
        logging.debug('main listening on port: '+str(port_subscribe))
        self.zmq_sub = zmq.Context().socket(zmq.SUB)
        self.zmq_sub.connect('tcp://127.0.0.1:'+str(port_subscribe))
        self.zmq_sub.setsockopt(zmq.SUBSCRIBE, '')

        # @copy_current_request_context
        def zmq_listener():
            while True:
                msg = self.zmq_sub.recv_json()
                self.zmq_confirmed = True
                logging.debug('main ('+') received '
                              'msg: '+str(msg))

                if 'description' in msg:
                    self.description = msg['description']

                if '__analysis_id' in msg and \
                   msg['__analysis_id'] in self.zmq_analyses:
                    analysis = self.zmq_analyses[msg['__analysis_id']]
                    del msg['__analysis_id']

                    if 'message' in msg and \
                       'signal' in msg['message'] and \
                       'message' in msg['message']:
                        analysis.emit(msg['message']['signal'],
                                      msg['message']['message'])
                    elif 'action_id' in msg and \
                         'status' in msg:
                        analysis.emit_action(msg['action_id'],
                                             msg['status'])
                    else:
                        logging.debug('dont understand this message: ' +
                                      str(msg))
                else:
                    logging.debug('__analysis_id not in message or '
                                  'AnalysisZMQ with that id not found.')
        self.zmq_listener = gevent.Greenlet.spawn(zmq_listener)

        # launch the language kernel process
        self.kernel_process = subprocess.Popen(executable, shell=False)

        # init language kernel
        def sending_init():
            while not self.zmq_confirmed:
                logging.debug('init kernel '+self.name+' to publish on '
                              'port '+str(port_subscribe))
                self.zmq_publish.send_json({
                    '__databench_namespace': self.name,
                    'publish_on_port': port_subscribe,
                })
                time.sleep(0.1)
        gevent.Greenlet.spawn(sending_init)

    def __del__(self):
        self.kernel_process.terminate()
        self.kernel_process.kill()

    def instantiate_analysis_class(self):
        self.zmq_analysis_id += 1
        i = self.analysis_class(self.name,
                                self.zmq_analysis_id,
                                self.zmq_publish)
        self.zmq_analyses[self.zmq_analysis_id] = i
        return i
<|endoftext|>


<commit_before>import json
import os

from databuild.facets import sum_facets
from databuild.loader import load_classpath

from databuild import settings

class Operator(object):
    operations = []

    def __init__(self, workbook):
        self.workbook = workbook
        self.languages = self.build_languages()

        super(Operator, self).__init__()

    def build_languages(self):
        languages = {}
        for name, runtime in settings.LANGUAGES.items():
            try:
                RuntimeClass = load_classpath(runtime)
            except ImportError:
                continue
            else:
                languages[name] = RuntimeClass(self.workbook)
        return languages

    def apply_operations(self, build_file, echo=False):
        self.build_file = os.path.abspath(build_file)

        with open(build_file, 'rb') as fh:
            operations = json.load(fh)
            [self.apply_operation(op, echo) for op in operations]

    def apply_operation(self, operation, echo=False):
        if echo and operation['description']:
            print(operation['description'])

        fn = load_classpath(operation['path'])

        if 'expression' in operation['params']:
            operation['params']['expression'] = self.parse_expression(operation['params']['expression'])

        if 'facets' in operation['params']:
            facets = [self.parse_expression(facet['expression']) for facet in operation['params']['facets']]
            operation['params']['facets'] = sum_facets(facets)

        kwargs = operation['params']
        fn(self.workbook, **kwargs)
        self.operations.append(operation)

    def parse_expression(self, expression):
        language, exp = expression['language'], expression['content']
        runtime = self.languages[language]
        return runtime.eval(exp)
<commit_msg>specified encoding in json loads<commit_after>import json
import os

from databuild.facets import sum_facets
from databuild.loader import load_classpath

from databuild import settings

class Operator(object):
    operations = []

    def __init__(self, workbook):
        self.workbook = workbook
        self.languages = self.build_languages()

        super(Operator, self).__init__()

    def build_languages(self):
        languages = {}
        for name, runtime in settings.LANGUAGES.items():
            try:
                RuntimeClass = load_classpath(runtime)
            except ImportError:
                continue
            else:
                languages[name] = RuntimeClass(self.workbook)
        return languages

    def apply_operations(self, build_file, echo=False):
        self.build_file = os.path.abspath(build_file)

        with open(build_file, 'rb') as fh:
            operations = json.load(fh, encoding='utf-8')
            [self.apply_operation(op, echo) for op in operations]

    def apply_operation(self, operation, echo=False):
        if echo and operation['description']:
            print(operation['description'])

        fn = load_classpath(operation['path'])

        if 'expression' in operation['params']:
            operation['params']['expression'] = self.parse_expression(operation['params']['expression'])

        if 'facets' in operation['params']:
            facets = [self.parse_expression(facet['expression']) for facet in operation['params']['facets']]
            operation['params']['facets'] = sum_facets(facets)

        kwargs = operation['params']
        fn(self.workbook, **kwargs)
        self.operations.append(operation)

    def parse_expression(self, expression):
        language, exp = expression['language'], expression['content']
        runtime = self.languages[language]
        return runtime.eval(exp)
<|endoftext|>


<commit_before>SHORTCUTS = {
    'core.import_data': "databuild.operations.core.import_data",
    'core.export_data': "databuild.operations.core.export_data",
    'columns.update_column': "databuild.operations.columns.update_column",
    'lua': "databuild.parsers.lua.lua",
}

FUNCTIONS = (
    'databuild.functions.data',
)

ADAPTER = 'databuild.adapters.locmem.LocMemBook'
<commit_msg>added `columns.*` shortcut<commit_after>SHORTCUTS = {
    'core.import_data': "databuild.operations.core.import_data",
    'core.export_data': "databuild.operations.core.export_data",
    'columns.add_column': "databuild.operations.columns.add_column",
    'columns.update_column': "databuild.operations.columns.update_column",
    'columns.remove_column': "databuild.operations.columns.remove_column",
    'columns.rename_column': "databuild.operations.columns.rename_column",
    'lua': "databuild.parsers.lua.lua",
}

FUNCTIONS = (
    'databuild.functions.data',
)

ADAPTER = 'databuild.adapters.locmem.LocMemBook'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# 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 os

import fixtures
import mock
from oslo_log import log
from oslotest import base
import testscenarios

from datahub.common import context as dh_context
import datahub.conf
from datahub.tests import conf_fixture

CONF = datahub.conf.CONF
try:
    log.register_options(CONF)
except CONF.ArgsAlreadyParsedError:
    pass
CONF.set_override('use_stderr', False)

class BaseTestCase(testscenarios.WithScenarios, base.BaseTestCase):
    """Test case base class for all unit tests."""

    def setUp(self):
        super(BaseTestCase, self).setUp()
        self.addCleanup(CONF.reset)

class TestCase(base.BaseTestCase):
    """Test case base class for all unit tests."""

    def setUp(self):
        super(TestCase, self).setUp()
        self.context = dh_context.RequestContext(user='fake_user')

        def make_context(*args, **kwargs):
            if not kwargs.get('user_id'):
                kwargs['user_id'] = 'fake_user'

            context = dh_context.RequestContext(*args, **kwargs)
            return dh_context.RequestContext.from_dict(context.to_dict())

        p = mock.patch.object(dh_context, 'make_context',
                              side_effect=make_context)
        self.mock_make_context = p.start()
        self.addCleanup(p.stop)

        self.useFixture(conf_fixture.ConfFixture())
        self.useFixture(fixtures.NestedTempfile())

    def config(self, **kw):
        """Override config options for a test."""
        group = kw.pop('group', None)
        for k, v in kw.items():
            CONF.set_override(k, v, group)

    def get_path(self, project_file=None):
        """Get the absolute path to a file. Used for testing the API.

        :param project_file: File whose path to return. Default: None.
        :returns: path to the specified file, or path to project root.
        """
        root = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                            '..',
                                            '..',
                                            )
                               )
        if project_file:
            return os.path.join(root, project_file)
        else:
            return root
<commit_msg>capture stderr test<commit_after># -*- coding: utf-8 -*-
#
# 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 os

import fixtures
import mock
from oslo_log import log
from oslotest import base
import testscenarios

from datahub.common import context as dh_context
import datahub.conf
from datahub.tests import conf_fixture

CONF = datahub.conf.CONF
try:
    log.register_options(CONF)
except CONF.ArgsAlreadyParsedError:
    pass

class BaseTestCase(testscenarios.WithScenarios, base.BaseTestCase):
    """Test case base class for all unit tests."""

    def setUp(self):
        super(BaseTestCase, self).setUp()
        self.addCleanup(CONF.reset)

class TestCase(base.BaseTestCase):
    """Test case base class for all unit tests."""

    def setUp(self):
        super(TestCase, self).setUp()
        self.context = dh_context.RequestContext(user='fake_user')

        def make_context(*args, **kwargs):
            if not kwargs.get('user_id'):
                kwargs['user_id'] = 'fake_user'

            context = dh_context.RequestContext(*args, **kwargs)
            return dh_context.RequestContext.from_dict(context.to_dict())

        p = mock.patch.object(dh_context, 'make_context',
                              side_effect=make_context)
        self.mock_make_context = p.start()
        self.addCleanup(p.stop)

        self.useFixture(conf_fixture.ConfFixture())
        self.useFixture(fixtures.NestedTempfile())

    def config(self, **kw):
        """Override config options for a test."""
        group = kw.pop('group', None)
        for k, v in kw.items():
            CONF.set_override(k, v, group)

    def get_path(self, project_file=None):
        """Get the absolute path to a file. Used for testing the API.

        :param project_file: File whose path to return. Default: None.
        :returns: path to the specified file, or path to project root.
        """
        root = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                            '..',
                                            '..',
                                            )
                               )
        if project_file:
            return os.path.join(root, project_file)
        else:
            return root
<|endoftext|>


<commit_before>"""
Drudge, a symbolic system for non-commutative and tensor algebra
================================================================

"""

from .canonpy import Perm, Group
from .term import Range, Vec, Term
from .canon import IDENT, NEG, CONJ
from .drudge import Tensor, Drudge
from .wick import WickDrudge
from .fock import (
    CR, AN, FERMI, BOSE, FockDrudge, GenMBDrudge, PartHoleDrudge,
    UP, DOWN, SpinOneHalfGenDrudge, SpinOneHalfPartHoleDrudge
)
from .utils import sum_, prod_

__all__ = [
    # Canonpy.
    'Perm',
    'Group',

    # Vec.
    'Vec',

    # Term.
    'Range',
    'Term',

    # Canon.
    'IDENT',
    'NEG',
    'CONJ',

    # Drudge.
    'Tensor',
    'Drudge',

    # Different problem-specific drudges.
    #
    # Base Wick algebra.
    'WickDrudge',

    # Many-body theories.
    'CR', 'AN', 'FERMI', 'BOSE',
    'FockDrudge',
    'GenMBDrudge', 'PartHoleDrudge',
    'UP', 'DOWN',
    'SpinOneHalfGenDrudge', 'SpinOneHalfPartHoleDrudge',

    # Small user utilities.
    'sum_',
    'prod_'
]
<commit_msg>Forward the time stamper to the base package<commit_after>"""
Drudge, a symbolic system for non-commutative and tensor algebra
================================================================

"""

from .canonpy import Perm, Group
from .term import Range, Vec, Term
from .canon import IDENT, NEG, CONJ
from .drudge import Tensor, Drudge
from .wick import WickDrudge
from .fock import (
    CR, AN, FERMI, BOSE, FockDrudge, GenMBDrudge, PartHoleDrudge,
    UP, DOWN, SpinOneHalfGenDrudge, SpinOneHalfPartHoleDrudge
)
from .utils import sum_, prod_, TimeStamper

__all__ = [
    # Canonpy.
    'Perm',
    'Group',

    # Vec.
    'Vec',

    # Term.
    'Range',
    'Term',

    # Canon.
    'IDENT',
    'NEG',
    'CONJ',

    # Drudge.
    'Tensor',
    'Drudge',

    # Different problem-specific drudges.
    #
    # Base Wick algebra.
    'WickDrudge',

    # Many-body theories.
    'CR', 'AN', 'FERMI', 'BOSE',
    'FockDrudge',
    'GenMBDrudge', 'PartHoleDrudge',
    'UP', 'DOWN',
    'SpinOneHalfGenDrudge', 'SpinOneHalfPartHoleDrudge',

    # Small user utilities.
    'sum_',
    'prod_',
    'TimeStamper'
]
<|endoftext|>