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

import numpy as np
import simsurvey.cadence as simul

def load_ztf_fields(filename='data/ZTF_Fields.txt', mwebv=False):
    """Load the ZTF fields propose by Eran from the included file.

    Parameters
    ----------
    filename: [str]
        File name of the ASCII file containing the field definitions

    mwebv: [bool]
        Include the Milky Way E(B-V) from the file in the output

    Return
    ------
    Dictionary of np.arrays with the field coordinates, IDs (and exinction)
    """
    fields = np.genfromtxt(filename, comments='%')

    out = {'field_id': np.array(fields[:,0], dtype=int),
           'ra': fields[:,1],
           'dec': fields[:,2]}

    if mwebv:
        out['mwebv'] = fields[:,3]

    return out

def get_survey_plan_single_field(mjd_range=(58000, 58020),
                                 width=10.,
                                 height=10.,
                                 dt=3600.,
                                 bands=['r_r_gggg', 'r_r_ggii'],
                                 band_str='des%s',
                                 skynoise={'g': 800., 'r': 800., 'i': 1260.}):
    """Generates a very simple plan with just one field. This can be used to
    simulate some lightcurves without specifying a full survey.
    See `SNIa_early_lcs.py/ipynb` for examples of use

    Parameters
    ----------
    mjd_range: [tuple of floats]
        Start and end date (in units of days) for the plan

    width, height: [floats]
        Width and height of the single survey field

    dt: [float]
        Time interval between two observations a night

    bands: [list of strs]
        List of filters that should be used for the observations;
        The plan will alternate between the strings each night and
        each string specifies the bands (as single letters) to be used
        with separation `dt`. Use underscore ("_") to skip an observation.

    band_str: [str]
        Formatting string used to create the full name of the bandpass
        as registered in `sncosmo`.

    skynoise: [dict]
        Dictionary of skynoise values to be used for each band.

    Return
    ------
    SurveyPlan object
    """
    obs_days = np.arange(mjd_range[0], mjd_range[1]+1, 1)

    obs = {'time': [], 'field': [], 'band': [], 'skynoise': []}

    for l, d in enumerate(obs_days):
        for k, band in enumerate(bands[l % len(bands)]):
            if band != '_':
                obs['field'].append(0)
                obs['time'].append(d + k * dt / 86400.)
                obs['band'].append(band_str%band)
                obs['skynoise'].append(skynoise[band])

    return simul.SurveyPlan(time=obs['time'], band=obs['band'],
                            obs_field=obs['field'],
                            skynoise=obs['skynoise'],
                            fields=dict(ra=[0.], dec=[0.]),
                            width=width, height=height)

def get_survey_plan_simple(mjd_range=(58000, 59095),
                           dec_range=(-30, 90),
                           ra_range=(0, 120),
                           dec_good=30.,
                           mwebv_max=0.4,
                           t_obs=45.,
                           t_night=28801.,
                           n_repeat=80,
                           bands=['r_r_gggg', 'r_r_ggii'],
                           band_str='des%s',
                           bands_alt=['rrrrrrrr', 'rrrrrrrr'],
                           fields=None,
                           skynoise={'g': 800., 'r': 800., 'i': 1260.}):
    """Generates a more sophisticated plan using a list of field.
    This can be used to simulate a full survey until Eric's code .

    See `SNIa_90d.ipynb` for an example of use.

    Parameters
    ----------
    mjd_range: [tuple of floats]
        Start and end date (in units of days) for the plan

    dec_range: [tuple of floats]
        Allowed declination range for the centers of the fields.

    ra_range: [tuple of floats]
        Allowed right ascension range for the centers of the fields *during the
        first night*. This is shifted by a little bit each night to account for
        Earth orbit the Sun.

    dec_good: [float]
        "Good" declination value. The fields for each night will be sorted
        according to how close they are to this values. Only `2 * n_repeat`
        fields will be used.  The first `n_repeat` fields are considered the
        "good" fields (e.g. near zenith) while the second set is considered
        the "alternate" fields.

    mwebv_max: [float]
        Maximum MW E(B-V) value for field to be used.

    t_obs: [float]
        Average time (in seconds) needed per observation (including slewing
        and other overhead)

    t_night: [float]
        Total available time (in seconds) per night.

    n_repeat: [int]
        Number of repeats for one band before moving to the next (or repeating
        the same one).

    bands: [list of strs]
        List of filters that should be used for the observations;
        The plan will alternate between the strings each night and
        each string specifies the bands (as single letters) to be used
        with separation `n_repeat` times with seperation `t_obs`.
        Use underscore ("_") to use `bands_alt` and the seconds list of
        `n_repeat` fields.

    band_str: [str]
        Formatting string used to create the full name of the bandpass
        as registered in `sncosmo`.

    bands_alt: [list of strs]
        Similar to `bands` but only use if.

    fields: [dict similar to output of `load_ztf_fields()`]
        Dictionary of field definitions; should include MW E(B-V).

    skynoise: [dict]
        Dictionary of skynoise values to be used for each band.

    Return
    ------
    SurveyPlan object
    """
    if fields is None:
        fields = load_ztf_fields(mwebv=True)

    ra_shift = 360. / 365.25
    obs_days = np.arange(mjd_range[0], mjd_range[1]+1, 1)

    obs = {'time': [], 'field': [], 'band': [], 'skynoise': []}

    for l, d in enumerate(obs_days):
        # Find fields that have their center within dec_range and in ra_range shifted
        # by ra_shift for each day
        field_mask = (((fields['ra'] - ra_shift*(d-obs_days[0]) + 180)%360 - 180
                       > ra_range[0]) &
                      ((fields['ra'] - ra_shift*(d-obs_days[0]) + 180)%360 - 180
                       < ra_range[1]) &
                      (fields['dec'] > dec_range[0]) &
                      (fields['dec'] < dec_range[1]) &
                      (fields['mwebv'] < 0.4))
        field_idx = np.where(field_mask)[0]

        # Sort first by ra then by dec
        field_idx = field_idx[np.argsort(fields['ra'][field_idx])]
        field_idx = field_idx[np.argsort(np.abs(dec_good - fields['dec'][field_idx]))]
        fields_good = fields['field_id'][field_idx[:n_repeat]]
        fields_alt = fields['field_id'][field_idx[n_repeat:2*n_repeat]]

        for k, t in enumerate(np.arange(0, t_night, t_obs)):
            b0, b1 = (l%len(bands), (k / len(fields_good)) % len(bands[l%len(bands)]))
            band = bands[b0][b1]
            if band != '_':
                obs['field'].append(fields_good[k % len(fields_good)])
            else:
                obs['field'].append(fields_alt[k % len(fields_alt)])
                band = bands_alt[b0][b1]

            obs['time'].append(d + t / 86400.)
            obs['band'].append(band_str%band)
            obs['skynoise'].append(skynoise[band])

    return simul.SurveyPlan(time=obs['time'], band=obs['band'], obs_field=obs['field'],
                            skynoise=obs['skynoise'],
                            fields={k: v for k, v in fields.items()
                                    if k in ['ra', 'dec', 'field_id',
                                             'width', 'height']})

def get_survey_lcs_single_field(tr, progress_bar=True, notebook=False,
                                instprop=None, **kw):
    """
    """
    plan = get_survey_plan_single_field(**kw)
    if instprop is None:
        instprop = {"desg":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desr":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desi":{"gain":1.,"zp":30,"zpsys":'ab'}}

    survey = simul.SimulSurvey(
        generator=tr,
        plan=plan,
        instprop=instprop
    )

    lcs = survey.get_lightcurves(progress_bar=progress_bar, notebook=notebook)

    return survey, lcs

def get_survey_lcs_simple(tr, progress_bar=True, notebook=False, instprop=None, **kw):
    """
    """
    plan = get_survey_plan_simple(**kw)
    if instprop is None:
        instprop = {"desg":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desr":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desi":{"gain":1.,"zp":30,"zpsys":'ab'}}

    survey = simul.SimulSurvey(
        generator=tr,
        plan=plan,
        instprop=instprop
    )

    lcs = survey.get_lightcurves(progress_bar=progress_bar, notebook=notebook)

    return survey, lcs
<commit_msg>fixed field selection<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import simsurvey.cadence as simul

def load_ztf_fields(filename='data/ZTF_Fields.txt', mwebv=False):
    """Load the ZTF fields propose by Eran from the included file.

    Parameters
    ----------
    filename: [str]
        File name of the ASCII file containing the field definitions

    mwebv: [bool]
        Include the Milky Way E(B-V) from the file in the output

    Return
    ------
    Dictionary of np.arrays with the field coordinates, IDs (and exinction)
    """
    fields = np.genfromtxt(filename, comments='%')

    out = {'field_id': np.array(fields[:,0], dtype=int),
           'ra': fields[:,1],
           'dec': fields[:,2]}

    if mwebv:
        out['mwebv'] = fields[:,3]

    return out

def get_survey_plan_single_field(mjd_range=(58000, 58020),
                                 width=10.,
                                 height=10.,
                                 dt=3600.,
                                 bands=['r_r_gggg', 'r_r_ggii'],
                                 band_str='des%s',
                                 skynoise={'g': 800., 'r': 800., 'i': 1260.}):
    """Generates a very simple plan with just one field. This can be used to
    simulate some lightcurves without specifying a full survey.
    See `SNIa_early_lcs.py/ipynb` for examples of use

    Parameters
    ----------
    mjd_range: [tuple of floats]
        Start and end date (in units of days) for the plan

    width, height: [floats]
        Width and height of the single survey field

    dt: [float]
        Time interval between two observations a night

    bands: [list of strs]
        List of filters that should be used for the observations;
        The plan will alternate between the strings each night and
        each string specifies the bands (as single letters) to be used
        with separation `dt`. Use underscore ("_") to skip an observation.

    band_str: [str]
        Formatting string used to create the full name of the bandpass
        as registered in `sncosmo`.

    skynoise: [dict]
        Dictionary of skynoise values to be used for each band.

    Return
    ------
    SurveyPlan object
    """
    obs_days = np.arange(mjd_range[0], mjd_range[1]+1, 1)

    obs = {'time': [], 'field': [], 'band': [], 'skynoise': []}

    for l, d in enumerate(obs_days):
        for k, band in enumerate(bands[l % len(bands)]):
            if band != '_':
                obs['field'].append(0)
                obs['time'].append(d + k * dt / 86400.)
                obs['band'].append(band_str%band)
                obs['skynoise'].append(skynoise[band])

    return simul.SurveyPlan(time=obs['time'], band=obs['band'],
                            obs_field=obs['field'],
                            skynoise=obs['skynoise'],
                            fields=dict(ra=[0.], dec=[0.]),
                            width=width, height=height)

def get_survey_plan_simple(mjd_range=(58000, 59095),
                           dec_range=(-30, 90),
                           ra_range=(0, 120),
                           dec_good=30.,
                           mwebv_max=0.4,
                           t_obs=45.,
                           t_night=28801.,
                           n_repeat=80,
                           bands=['r_r_gggg', 'r_r_ggii'],
                           band_str='des%s',
                           bands_alt=['rrrrrrrr', 'rrrrrrrr'],
                           fields=None,
                           skynoise={'g': 800., 'r': 800., 'i': 1260.}):
    """Generates a more sophisticated plan using a list of field.
    This can be used to simulate a full survey until Eric's code .

    See `SNIa_90d.ipynb` for an example of use.

    Parameters
    ----------
    mjd_range: [tuple of floats]
        Start and end date (in units of days) for the plan

    dec_range: [tuple of floats]
        Allowed declination range for the centers of the fields.

    ra_range: [tuple of floats]
        Allowed right ascension range for the centers of the fields *during the
        first night*. This is shifted by a little bit each night to account for
        Earth orbit the Sun.

    dec_good: [float]
        "Good" declination value. The fields for each night will be sorted
        according to how close they are to this values. Only `2 * n_repeat`
        fields will be used.  The first `n_repeat` fields are considered the
        "good" fields (e.g. near zenith) while the second set is considered
        the "alternate" fields.

    mwebv_max: [float]
        Maximum MW E(B-V) value for field to be used.

    t_obs: [float]
        Average time (in seconds) needed per observation (including slewing
        and other overhead)

    t_night: [float]
        Total available time (in seconds) per night.

    n_repeat: [int]
        Number of repeats for one band before moving to the next (or repeating
        the same one).

    bands: [list of strs]
        List of filters that should be used for the observations;
        The plan will alternate between the strings each night and
        each string specifies the bands (as single letters) to be used
        with separation `n_repeat` times with seperation `t_obs`.
        Use underscore ("_") to use `bands_alt` and the seconds list of
        `n_repeat` fields.

    band_str: [str]
        Formatting string used to create the full name of the bandpass
        as registered in `sncosmo`.

    bands_alt: [list of strs]
        Similar to `bands` but only use if.

    fields: [dict similar to output of `load_ztf_fields()`]
        Dictionary of field definitions; should include MW E(B-V).

    skynoise: [dict]
        Dictionary of skynoise values to be used for each band.

    Return
    ------
    SurveyPlan object
    """
    if fields is None:
        fields = load_ztf_fields(mwebv=True)
        for k in fields.keys():
            fields[k] = fields[k][:906]

    ra_shift = 360. / 365.25
    obs_days = np.arange(mjd_range[0], mjd_range[1]+1, 1)

    obs = {'time': [], 'field': [], 'band': [], 'skynoise': []}

    for l, d in enumerate(obs_days):
        # Find fields that have their center within dec_range and in ra_range shifted
        # by ra_shift for each day
        field_mask = (((fields['ra'] - ra_shift*(d-obs_days[0]) + 180)%360 - 180
                       > ra_range[0]) &
                      ((fields['ra'] - ra_shift*(d-obs_days[0]) + 180)%360 - 180
                       < ra_range[1]) &
                      (fields['dec'] > dec_range[0]) &
                      (fields['dec'] < dec_range[1]) &
                      (fields['mwebv'] < 0.4))
        field_idx = np.where(field_mask)[0]

        # Sort first by ra then by dec
        field_idx = field_idx[np.argsort(fields['ra'][field_idx])]
        field_idx = field_idx[np.argsort(np.abs(dec_good - fields['dec'][field_idx]))]
        fields_good = fields['field_id'][field_idx[:n_repeat]]
        fields_alt = fields['field_id'][field_idx[n_repeat:2*n_repeat]]

        for k, t in enumerate(np.arange(0, t_night, t_obs)):
            b0, b1 = (l%len(bands), (k / len(fields_good)) % len(bands[l%len(bands)]))
            band = bands[b0][b1]
            if band != '_':
                obs['field'].append(fields_good[k % len(fields_good)])
            else:
                obs['field'].append(fields_alt[k % len(fields_alt)])
                band = bands_alt[b0][b1]

            obs['time'].append(d + t / 86400.)
            obs['band'].append(band_str%band)
            obs['skynoise'].append(skynoise[band])

    return simul.SurveyPlan(time=obs['time'], band=obs['band'], obs_field=obs['field'],
                            skynoise=obs['skynoise'],
                            fields={k: v for k, v in fields.items()
                                    if k in ['ra', 'dec', 'field_id',
                                             'width', 'height']})

def get_survey_lcs_single_field(tr, progress_bar=True, notebook=False,
                                instprop=None, **kw):
    """
    """
    plan = get_survey_plan_single_field(**kw)
    if instprop is None:
        instprop = {"desg":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desr":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desi":{"gain":1.,"zp":30,"zpsys":'ab'}}

    survey = simul.SimulSurvey(
        generator=tr,
        plan=plan,
        instprop=instprop
    )

    lcs = survey.get_lightcurves(progress_bar=progress_bar, notebook=notebook)

    return survey, lcs

def get_survey_lcs_simple(tr, progress_bar=True, notebook=False, instprop=None, **kw):
    """
    """
    plan = get_survey_plan_simple(**kw)
    if instprop is None:
        instprop = {"desg":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desr":{"gain":1.,"zp":30,"zpsys":'ab'},
                    "desi":{"gain":1.,"zp":30,"zpsys":'ab'}}

    survey = simul.SimulSurvey(
        generator=tr,
        plan=plan,
        instprop=instprop
    )

    lcs = survey.get_lightcurves(progress_bar=progress_bar, notebook=notebook)

    return survey, lcs
<|endoftext|>


<commit_before>import logging
from collections import defaultdict
from ethereum.opcodes import opcodes
from typing import List, Tuple, Union, Callable, Dict
from mythril.analysis.security import get_detection_modules
from mythril.disassembler.disassembly import Disassembly
from mythril.laser.ethereum.state.account import Account
from mythril.laser.ethereum.state.world_state import WorldState
from mythril.laser.ethereum.state.global_state import GlobalState
from mythril.laser.ethereum.transaction import (
    TransactionStartSignal,
    TransactionEndSignal,
    ContractCreationTransaction,
)
from mythril.laser.ethereum.evm_exceptions import StackUnderflowException
from mythril.laser.ethereum.instructions import Instruction
from mythril.laser.ethereum.cfg import NodeFlags, Node, Edge, JumpType
from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy
from datetime import datetime, timedelta
from copy import copy
from mythril.laser.ethereum.transaction import (
    execute_contract_creation,
    execute_message_call,
)
from functools import reduce
from mythril.laser.ethereum.evm_exceptions import VmException

class SVMError(Exception):
    pass

"""
Main symbolic execution engine.
"""

OPCODE_LIST = [c[0] for _, c in opcodes.items()]

class LaserEVM:
    """
    Laser EVM class
    """

    def __init__(
        self,
        accounts: Dict[str, Account],
        dynamic_loader=None,
        max_depth=float("inf"),
        execution_timeout=60,
        create_timeout=10,
        strategy=DepthFirstSearchStrategy,
        max_transaction_count=3,
    ):
        world_state = WorldState()
        world_state.accounts = accounts
        # this sets the initial world state
        self.world_state = world_state
        self.open_states = [world_state]

        self.nodes = {}
        self.edges = []
        self.coverage = {}

        self.total_states = 0
        self.dynamic_loader = dynamic_loader

        self.work_list = []
        self.strategy = strategy(self.work_list, max_depth)
        self.max_depth = max_depth
        self.max_transaction_count = max_transaction_count

        self.execution_timeout = execution_timeout
        self.create_timeout = create_timeout

        self.time = None

        self.pre_hooks = defaultdict(list)
        self.post_hooks = defaultdict(list)

        self.register_detection_modules()

        logging.info(
            "LASER EVM initialized with dynamic loader: " + str(dynamic_loader)
        )

    def register_detection_modules(self):
        modules = get_detection_modules(entrypoint="callback")
        for module in modules:
            for hook in module.detector.hooks:
                hook = hook.upper()
                if hook in OPCODE_LIST:
                    self.post_hooks[hook].append(module.detector.execute)
                elif hook.endswith("*"):
                    to_register = filter(lambda x: x.startswith(hook), OPCODE_LIST)
                    for actual_hook in to_register:
                        self.post_hooks[actual_hook].append(module.detector.execute)
                else:
                    logging.error(
                        "Encountered invalid hook opcode %s in module %s",
                        hook,
                        module.detector.name,
                    )

    @property
    def accounts(self) -> Dict[str, Account]:
        return self.world_state.accounts

    def sym_exec(
        self,
        main_address=None,
        creation_code=None,
        contract_name=None,
        max_transactions=3,
    ) -> None:
        logging.debug("Starting LASER execution")
        self.time = datetime.now()

        if main_address:
            logging.info("Starting message call transaction to {}".format(main_address))
            self._execute_transactions(main_address)

        elif creation_code:
            logging.info("Starting contract creation transaction")
            created_account = execute_contract_creation(
                self, creation_code, contract_name
            )
            logging.info(
                "Finished contract creation, found {} open states".format(
                    len(self.open_states)
                )
            )
            if len(self.open_states) == 0:
                logging.warning(
                    "No contract was created during the execution of contract creation "
                    "Increase the resources for creation execution (--max-depth or --create-timeout)"
                )

            self._execute_transactions(created_account.address)

        logging.info("Finished symbolic execution")
        logging.info(
            "%d nodes, %d edges, %d total states",
            len(self.nodes),
            len(self.edges),
            self.total_states,
        )
        for code, coverage in self.coverage.items():
            cov = (
                reduce(lambda sum_, val: sum_ + 1 if val else sum_, coverage[1])
                / float(coverage[0])
                * 100
            )
            logging.info("Achieved {:.2f}% coverage for code: {}".format(cov, code))

    def _execute_transactions(self, address):
        """
        This function executes multiple transactions on the address based on the coverage
        :param address: Address of the contract
        :return:
        """
        self.coverage = {}
        for i in range(self.max_transaction_count):
            initial_coverage = self._get_covered_instructions()

            self.time = datetime.now()
            logging.info("Starting message call transaction, iteration: {}".format(i))
            execute_message_call(self, address)

            end_coverage = self._get_covered_instructions()
            if end_coverage == initial_coverage:
                break

    def _get_covered_instructions(self) -> int:
        """ Gets the total number of covered instructions for all accounts in the svm"""
        total_covered_instructions = 0
        for _, cv in self.coverage.items():
            total_covered_instructions += reduce(
                lambda sum_, val: sum_ + 1 if val else sum_, cv[1]
            )
        return total_covered_instructions

    def exec(self, create=False, track_gas=False) -> Union[List[GlobalState], None]:
        final_states = []
        for global_state in self.strategy:
            if self.execution_timeout and not create:
                if (
                    self.time + timedelta(seconds=self.execution_timeout)
                    <= datetime.now()
                ):
                    return final_states + [global_state] if track_gas else None
            elif self.create_timeout and create:
                if self.time + timedelta(seconds=self.create_timeout) <= datetime.now():
                    return final_states + [global_state] if track_gas else None

            try:
                new_states, op_code = self.execute_state(global_state)
            except NotImplementedError:
                logging.debug("Encountered unimplemented instruction")
                continue
            self.manage_cfg(op_code, new_states)

            if new_states:
                self.work_list += new_states
            elif track_gas:
                final_states.append(global_state)
            self.total_states += len(new_states)
        return final_states if track_gas else None

    def execute_state(
        self, global_state: GlobalState
    ) -> Tuple[List[GlobalState], Union[str, None]]:
        instructions = global_state.environment.code.instruction_list
        try:
            op_code = instructions[global_state.mstate.pc]["opcode"]
        except IndexError:
            self.open_states.append(global_state.world_state)
            return [], None

        self._execute_pre_hook(op_code, global_state)
        try:
            self._measure_coverage(global_state)
            new_global_states = Instruction(op_code, self.dynamic_loader).evaluate(
                global_state
            )

        except VmException as e:
            transaction, return_global_state = global_state.transaction_stack.pop()

            if return_global_state is None:
                # In this case we don't put an unmodified world state in the open_states list Since in the case of an
                #  exceptional halt all changes should be discarded, and this world state would not provide us with a
                #  previously unseen world state
                logging.debug(
                    "Encountered a VmException, ending path: `{}`".format(str(e))
                )
                new_global_states = []
            else:
                # First execute the post hook for the transaction ending instruction
                self._execute_post_hook(op_code, [global_state])
                new_global_states = self._end_message_call(
                    return_global_state,
                    global_state,
                    revert_changes=True,
                    return_data=None,
                )

        except TransactionStartSignal as start_signal:
            # Setup new global state
            new_global_state = start_signal.transaction.initial_global_state()

            new_global_state.transaction_stack = copy(
                global_state.transaction_stack
            ) + [(start_signal.transaction, global_state)]
            new_global_state.node = global_state.node
            new_global_state.mstate.constraints = global_state.mstate.constraints

            return [new_global_state], op_code

        except TransactionEndSignal as end_signal:
            transaction, return_global_state = (
                end_signal.global_state.transaction_stack.pop()
            )

            if return_global_state is None:
                if (
                    not isinstance(transaction, ContractCreationTransaction)
                    or transaction.return_data
                ) and not end_signal.revert:
                    end_signal.global_state.world_state.node = global_state.node
                    self.open_states.append(end_signal.global_state.world_state)
                new_global_states = []
            else:
                # First execute the post hook for the transaction ending instruction
                self._execute_post_hook(op_code, [end_signal.global_state])

                new_global_states = self._end_message_call(
                    return_global_state,
                    global_state,
                    revert_changes=False or end_signal.revert,
                    return_data=transaction.return_data,
                )

        self._execute_post_hook(op_code, new_global_states)

        return new_global_states, op_code

    def _end_message_call(
        self,
        return_global_state: GlobalState,
        global_state: GlobalState,
        revert_changes=False,
        return_data=None,
    ) -> List[GlobalState]:
        # Resume execution of the transaction initializing instruction
        op_code = return_global_state.environment.code.instruction_list[
            return_global_state.mstate.pc
        ]["opcode"]

        # Set execution result in the return_state
        return_global_state.last_return_data = return_data
        if not revert_changes:
            return_global_state.world_state = copy(global_state.world_state)
            return_global_state.environment.active_account = global_state.accounts[
                return_global_state.environment.active_account.address
            ]

        # Execute the post instruction handler
        new_global_states = Instruction(op_code, self.dynamic_loader).evaluate(
            return_global_state, True
        )

        # In order to get a nice call graph we need to set the nodes here
        for state in new_global_states:
            state.node = global_state.node

        return new_global_states

    def _measure_coverage(self, global_state: GlobalState) -> None:
        code = global_state.environment.code.bytecode
        number_of_instructions = len(global_state.environment.code.instruction_list)
        instruction_index = global_state.mstate.pc

        if code not in self.coverage.keys():
            self.coverage[code] = [
                number_of_instructions,
                [False] * number_of_instructions,
            ]

        self.coverage[code][1][instruction_index] = True

    def manage_cfg(self, opcode: str, new_states: List[GlobalState]) -> None:
        if opcode == "JUMP":
            assert len(new_states) <= 1
            for state in new_states:
                self._new_node_state(state)
        elif opcode == "JUMPI":
            for state in new_states:
                self._new_node_state(
                    state, JumpType.CONDITIONAL, state.mstate.constraints[-1]
                )
        elif opcode in ("SLOAD", "SSTORE") and len(new_states) > 1:
            for state in new_states:
                self._new_node_state(
                    state, JumpType.CONDITIONAL, state.mstate.constraints[-1]
                )

        elif opcode in ("CALL", "CALLCODE", "DELEGATECALL", "STATICCALL"):
            assert len(new_states) <= 1
            for state in new_states:
                self._new_node_state(state, JumpType.CALL)
                # Keep track of added contracts so the graph can be generated properly
                if (
                    state.environment.active_account.contract_name
                    not in self.world_state.accounts.keys()
                ):
                    self.world_state.accounts[
                        state.environment.active_account.address
                    ] = state.environment.active_account
        elif opcode == "RETURN":
            for state in new_states:
                self._new_node_state(state, JumpType.RETURN)

        for state in new_states:
            state.node.states.append(state)

    def _new_node_state(
        self, state: GlobalState, edge_type=JumpType.UNCONDITIONAL, condition=None
    ) -> None:
        new_node = Node(state.environment.active_account.contract_name)
        old_node = state.node
        state.node = new_node
        new_node.constraints = state.mstate.constraints
        self.nodes[new_node.uid] = new_node
        self.edges.append(
            Edge(old_node.uid, new_node.uid, edge_type=edge_type, condition=condition)
        )

        if edge_type == JumpType.RETURN:
            new_node.flags |= NodeFlags.CALL_RETURN
        elif edge_type == JumpType.CALL:
            try:
                if "retval" in str(state.mstate.stack[-1]):
                    new_node.flags |= NodeFlags.CALL_RETURN
                else:
                    new_node.flags |= NodeFlags.FUNC_ENTRY
            except StackUnderflowException:
                new_node.flags |= NodeFlags.FUNC_ENTRY
        address = state.environment.code.instruction_list[state.mstate.pc]["address"]

        environment = state.environment
        disassembly = environment.code
        if address in disassembly.address_to_function_name:
            # Enter a new function
            environment.active_function_name = disassembly.address_to_function_name[
                address
            ]
            new_node.flags |= NodeFlags.FUNC_ENTRY

            logging.debug(
                "- Entering function "
                + environment.active_account.contract_name
                + ":"
                + new_node.function_name
            )
        elif address == 0:
            environment.active_function_name = "fallback"

        new_node.function_name = environment.active_function_name

    def _execute_pre_hook(self, op_code: str, global_state: GlobalState) -> None:
        if op_code not in self.pre_hooks.keys():
            return
        for hook in self.pre_hooks[op_code]:
            hook(global_state)

    def _execute_post_hook(
        self, op_code: str, global_states: List[GlobalState]
    ) -> None:
        if op_code not in self.post_hooks.keys():
            return

        for hook in self.post_hooks[op_code]:
            for global_state in global_states:
                hook(global_state)

    def pre_hook(self, op_code: str) -> Callable:
        def hook_decorator(func: Callable):
            if op_code not in self.pre_hooks.keys():
                self.pre_hooks[op_code] = []
            self.pre_hooks[op_code].append(func)
            return func

        return hook_decorator

    def post_hook(self, op_code: str) -> Callable:
        def hook_decorator(func: Callable):
            if op_code not in self.post_hooks.keys():
                self.post_hooks[op_code] = []
            self.post_hooks[op_code].append(func)
            return func

        return hook_decorator
<commit_msg>Fix OPCODE* prefix<commit_after>import logging
from collections import defaultdict
from ethereum.opcodes import opcodes
from typing import List, Tuple, Union, Callable, Dict
from mythril.analysis.security import get_detection_modules
from mythril.disassembler.disassembly import Disassembly
from mythril.laser.ethereum.state.account import Account
from mythril.laser.ethereum.state.world_state import WorldState
from mythril.laser.ethereum.state.global_state import GlobalState
from mythril.laser.ethereum.transaction import (
    TransactionStartSignal,
    TransactionEndSignal,
    ContractCreationTransaction,
)
from mythril.laser.ethereum.evm_exceptions import StackUnderflowException
from mythril.laser.ethereum.instructions import Instruction
from mythril.laser.ethereum.cfg import NodeFlags, Node, Edge, JumpType
from mythril.laser.ethereum.strategy.basic import DepthFirstSearchStrategy
from datetime import datetime, timedelta
from copy import copy
from mythril.laser.ethereum.transaction import (
    execute_contract_creation,
    execute_message_call,
)
from functools import reduce
from mythril.laser.ethereum.evm_exceptions import VmException

class SVMError(Exception):
    pass

"""
Main symbolic execution engine.
"""

OPCODE_LIST = [c[0] for _, c in opcodes.items()]

class LaserEVM:
    """
    Laser EVM class
    """

    def __init__(
        self,
        accounts: Dict[str, Account],
        dynamic_loader=None,
        max_depth=float("inf"),
        execution_timeout=60,
        create_timeout=10,
        strategy=DepthFirstSearchStrategy,
        max_transaction_count=3,
    ):
        world_state = WorldState()
        world_state.accounts = accounts
        # this sets the initial world state
        self.world_state = world_state
        self.open_states = [world_state]

        self.nodes = {}
        self.edges = []
        self.coverage = {}

        self.total_states = 0
        self.dynamic_loader = dynamic_loader

        self.work_list = []
        self.strategy = strategy(self.work_list, max_depth)
        self.max_depth = max_depth
        self.max_transaction_count = max_transaction_count

        self.execution_timeout = execution_timeout
        self.create_timeout = create_timeout

        self.time = None

        self.pre_hooks = defaultdict(list)
        self.post_hooks = defaultdict(list)

        self.register_detection_modules()

        logging.info(
            "LASER EVM initialized with dynamic loader: " + str(dynamic_loader)
        )

    def register_detection_modules(self):
        modules = get_detection_modules(entrypoint="callback")
        for module in modules:
            for hook in module.detector.hooks:
                hook = hook.upper()
                if hook in OPCODE_LIST:
                    self.post_hooks[hook].append(module.detector.execute)
                elif hook.endswith("*"):
                    to_register = filter(lambda x: x.startswith(hook[:-1]), OPCODE_LIST)
                    for actual_hook in to_register:
                        self.post_hooks[actual_hook].append(module.detector.execute)
                else:
                    logging.error(
                        "Encountered invalid hook opcode %s in module %s",
                        hook,
                        module.detector.name,
                    )

    @property
    def accounts(self) -> Dict[str, Account]:
        return self.world_state.accounts

    def sym_exec(
        self,
        main_address=None,
        creation_code=None,
        contract_name=None,
        max_transactions=3,
    ) -> None:
        logging.debug("Starting LASER execution")
        self.time = datetime.now()

        if main_address:
            logging.info("Starting message call transaction to {}".format(main_address))
            self._execute_transactions(main_address)

        elif creation_code:
            logging.info("Starting contract creation transaction")
            created_account = execute_contract_creation(
                self, creation_code, contract_name
            )
            logging.info(
                "Finished contract creation, found {} open states".format(
                    len(self.open_states)
                )
            )
            if len(self.open_states) == 0:
                logging.warning(
                    "No contract was created during the execution of contract creation "
                    "Increase the resources for creation execution (--max-depth or --create-timeout)"
                )

            self._execute_transactions(created_account.address)

        logging.info("Finished symbolic execution")
        logging.info(
            "%d nodes, %d edges, %d total states",
            len(self.nodes),
            len(self.edges),
            self.total_states,
        )
        for code, coverage in self.coverage.items():
            cov = (
                reduce(lambda sum_, val: sum_ + 1 if val else sum_, coverage[1])
                / float(coverage[0])
                * 100
            )
            logging.info("Achieved {:.2f}% coverage for code: {}".format(cov, code))

    def _execute_transactions(self, address):
        """
        This function executes multiple transactions on the address based on the coverage
        :param address: Address of the contract
        :return:
        """
        self.coverage = {}
        for i in range(self.max_transaction_count):
            initial_coverage = self._get_covered_instructions()

            self.time = datetime.now()
            logging.info("Starting message call transaction, iteration: {}".format(i))
            execute_message_call(self, address)

            end_coverage = self._get_covered_instructions()
            if end_coverage == initial_coverage:
                break

    def _get_covered_instructions(self) -> int:
        """ Gets the total number of covered instructions for all accounts in the svm"""
        total_covered_instructions = 0
        for _, cv in self.coverage.items():
            total_covered_instructions += reduce(
                lambda sum_, val: sum_ + 1 if val else sum_, cv[1]
            )
        return total_covered_instructions

    def exec(self, create=False, track_gas=False) -> Union[List[GlobalState], None]:
        final_states = []
        for global_state in self.strategy:
            if self.execution_timeout and not create:
                if (
                    self.time + timedelta(seconds=self.execution_timeout)
                    <= datetime.now()
                ):
                    return final_states + [global_state] if track_gas else None
            elif self.create_timeout and create:
                if self.time + timedelta(seconds=self.create_timeout) <= datetime.now():
                    return final_states + [global_state] if track_gas else None

            try:
                new_states, op_code = self.execute_state(global_state)
            except NotImplementedError:
                logging.debug("Encountered unimplemented instruction")
                continue
            self.manage_cfg(op_code, new_states)

            if new_states:
                self.work_list += new_states
            elif track_gas:
                final_states.append(global_state)
            self.total_states += len(new_states)
        return final_states if track_gas else None

    def execute_state(
        self, global_state: GlobalState
    ) -> Tuple[List[GlobalState], Union[str, None]]:
        instructions = global_state.environment.code.instruction_list
        try:
            op_code = instructions[global_state.mstate.pc]["opcode"]
        except IndexError:
            self.open_states.append(global_state.world_state)
            return [], None

        self._execute_pre_hook(op_code, global_state)
        try:
            self._measure_coverage(global_state)
            new_global_states = Instruction(op_code, self.dynamic_loader).evaluate(
                global_state
            )

        except VmException as e:
            transaction, return_global_state = global_state.transaction_stack.pop()

            if return_global_state is None:
                # In this case we don't put an unmodified world state in the open_states list Since in the case of an
                #  exceptional halt all changes should be discarded, and this world state would not provide us with a
                #  previously unseen world state
                logging.debug(
                    "Encountered a VmException, ending path: `{}`".format(str(e))
                )
                new_global_states = []
            else:
                # First execute the post hook for the transaction ending instruction
                self._execute_post_hook(op_code, [global_state])
                new_global_states = self._end_message_call(
                    return_global_state,
                    global_state,
                    revert_changes=True,
                    return_data=None,
                )

        except TransactionStartSignal as start_signal:
            # Setup new global state
            new_global_state = start_signal.transaction.initial_global_state()

            new_global_state.transaction_stack = copy(
                global_state.transaction_stack
            ) + [(start_signal.transaction, global_state)]
            new_global_state.node = global_state.node
            new_global_state.mstate.constraints = global_state.mstate.constraints

            return [new_global_state], op_code

        except TransactionEndSignal as end_signal:
            transaction, return_global_state = (
                end_signal.global_state.transaction_stack.pop()
            )

            if return_global_state is None:
                if (
                    not isinstance(transaction, ContractCreationTransaction)
                    or transaction.return_data
                ) and not end_signal.revert:
                    end_signal.global_state.world_state.node = global_state.node
                    self.open_states.append(end_signal.global_state.world_state)
                new_global_states = []
            else:
                # First execute the post hook for the transaction ending instruction
                self._execute_post_hook(op_code, [end_signal.global_state])

                new_global_states = self._end_message_call(
                    return_global_state,
                    global_state,
                    revert_changes=False or end_signal.revert,
                    return_data=transaction.return_data,
                )

        self._execute_post_hook(op_code, new_global_states)

        return new_global_states, op_code

    def _end_message_call(
        self,
        return_global_state: GlobalState,
        global_state: GlobalState,
        revert_changes=False,
        return_data=None,
    ) -> List[GlobalState]:
        # Resume execution of the transaction initializing instruction
        op_code = return_global_state.environment.code.instruction_list[
            return_global_state.mstate.pc
        ]["opcode"]

        # Set execution result in the return_state
        return_global_state.last_return_data = return_data
        if not revert_changes:
            return_global_state.world_state = copy(global_state.world_state)
            return_global_state.environment.active_account = global_state.accounts[
                return_global_state.environment.active_account.address
            ]

        # Execute the post instruction handler
        new_global_states = Instruction(op_code, self.dynamic_loader).evaluate(
            return_global_state, True
        )

        # In order to get a nice call graph we need to set the nodes here
        for state in new_global_states:
            state.node = global_state.node

        return new_global_states

    def _measure_coverage(self, global_state: GlobalState) -> None:
        code = global_state.environment.code.bytecode
        number_of_instructions = len(global_state.environment.code.instruction_list)
        instruction_index = global_state.mstate.pc

        if code not in self.coverage.keys():
            self.coverage[code] = [
                number_of_instructions,
                [False] * number_of_instructions,
            ]

        self.coverage[code][1][instruction_index] = True

    def manage_cfg(self, opcode: str, new_states: List[GlobalState]) -> None:
        if opcode == "JUMP":
            assert len(new_states) <= 1
            for state in new_states:
                self._new_node_state(state)
        elif opcode == "JUMPI":
            for state in new_states:
                self._new_node_state(
                    state, JumpType.CONDITIONAL, state.mstate.constraints[-1]
                )
        elif opcode in ("SLOAD", "SSTORE") and len(new_states) > 1:
            for state in new_states:
                self._new_node_state(
                    state, JumpType.CONDITIONAL, state.mstate.constraints[-1]
                )

        elif opcode in ("CALL", "CALLCODE", "DELEGATECALL", "STATICCALL"):
            assert len(new_states) <= 1
            for state in new_states:
                self._new_node_state(state, JumpType.CALL)
                # Keep track of added contracts so the graph can be generated properly
                if (
                    state.environment.active_account.contract_name
                    not in self.world_state.accounts.keys()
                ):
                    self.world_state.accounts[
                        state.environment.active_account.address
                    ] = state.environment.active_account
        elif opcode == "RETURN":
            for state in new_states:
                self._new_node_state(state, JumpType.RETURN)

        for state in new_states:
            state.node.states.append(state)

    def _new_node_state(
        self, state: GlobalState, edge_type=JumpType.UNCONDITIONAL, condition=None
    ) -> None:
        new_node = Node(state.environment.active_account.contract_name)
        old_node = state.node
        state.node = new_node
        new_node.constraints = state.mstate.constraints
        self.nodes[new_node.uid] = new_node
        self.edges.append(
            Edge(old_node.uid, new_node.uid, edge_type=edge_type, condition=condition)
        )

        if edge_type == JumpType.RETURN:
            new_node.flags |= NodeFlags.CALL_RETURN
        elif edge_type == JumpType.CALL:
            try:
                if "retval" in str(state.mstate.stack[-1]):
                    new_node.flags |= NodeFlags.CALL_RETURN
                else:
                    new_node.flags |= NodeFlags.FUNC_ENTRY
            except StackUnderflowException:
                new_node.flags |= NodeFlags.FUNC_ENTRY
        address = state.environment.code.instruction_list[state.mstate.pc]["address"]

        environment = state.environment
        disassembly = environment.code
        if address in disassembly.address_to_function_name:
            # Enter a new function
            environment.active_function_name = disassembly.address_to_function_name[
                address
            ]
            new_node.flags |= NodeFlags.FUNC_ENTRY

            logging.debug(
                "- Entering function "
                + environment.active_account.contract_name
                + ":"
                + new_node.function_name
            )
        elif address == 0:
            environment.active_function_name = "fallback"

        new_node.function_name = environment.active_function_name

    def _execute_pre_hook(self, op_code: str, global_state: GlobalState) -> None:
        if op_code not in self.pre_hooks.keys():
            return
        for hook in self.pre_hooks[op_code]:
            hook(global_state)

    def _execute_post_hook(
        self, op_code: str, global_states: List[GlobalState]
    ) -> None:
        if op_code not in self.post_hooks.keys():
            return

        for hook in self.post_hooks[op_code]:
            for global_state in global_states:
                hook(global_state)

    def pre_hook(self, op_code: str) -> Callable:
        def hook_decorator(func: Callable):
            if op_code not in self.pre_hooks.keys():
                self.pre_hooks[op_code] = []
            self.pre_hooks[op_code].append(func)
            return func

        return hook_decorator

    def post_hook(self, op_code: str) -> Callable:
        def hook_decorator(func: Callable):
            if op_code not in self.post_hooks.keys():
                self.post_hooks[op_code] = []
            self.post_hooks[op_code].append(func)
            return func

        return hook_decorator
<|endoftext|>


<commit_before>import pytest
from django.contrib.sites.models import Site

from controlled_vocabularies.models import Vocabulary, Term, Property

pytestmark = pytest.mark.django_db

def create_vocab(name='languages', label='Language', order='name',
                 maintainer='Buddy Noone', maintainerEmail='noone@nowhere.com',
                 definition='Lorem ipsum dolor sit amet'):
    """Creates a saved instance of the Vocabulary model."""
    vocab_instance = Vocabulary(name=name,
                                label=label,
                                order=order,
                                maintainer=maintainer,
                                maintainerEmail=maintainer,
                                definition=definition)
    vocab_instance.save()
    return(vocab_instance)

def create_term(name='spanish', label='Spanish', order=2):
    """Creates a saved instance of the Term model."""
    vocab_instance = create_vocab()
    term_instance = Term(vocab_list=vocab_instance,
                         name=name,
                         label=label,
                         order=order)
    term_instance.save()
    return(term_instance)

def create_property(property_name='dialect', label='Catalan'):
    """Creates a saved instance of the Property model."""
    term_instance = create_term()
    property_instance = Property(term_key=term_instance,
                                 property_name=property_name,
                                 label=label)
    property_instance.save()
    return(property_instance)

class TestVocabulary:
    def test_save(self):
        assert Vocabulary.objects.all().count() == 0
        create_vocab()
        assert Vocabulary.objects.all().count() == 1

    def test_unicode(self):
        vocab_instance = create_vocab()
        assert unicode(vocab_instance) == 'languages'

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        create_vocab(
            name=' spaces  here ',
            label=' Spaces here   too  ',
            maintainer=' Please   Fix Me ',
            maintainerEmail='  here@there.com ',
            definition='  Strip leading   and trailing, condense   others. '
        )
        vocab_from_db = Vocabulary.objects.get(pk=1)

        assert vocab_from_db.name == 'spaces here'
        assert vocab_from_db.label == 'Spaces here too'
        assert vocab_from_db.maintainer == 'Please Fix Me'
        assert vocab_from_db.maintainerEmail == 'here@there.com'
        assert vocab_from_db.definition == 'Strip leading and trailing, condense others.'

class TestTerm:
    def test_unicode(self):
        term_instance = create_term()
        assert unicode(term_instance) == 'spanish'

    def test_save(self):
        assert Term.objects.all().count() == 0
        create_term()
        assert Term.objects.all().count() == 1

    def test_get_vocab(self):
        term_instance = create_term()
        expected = "{}".format(
            Site.objects.get_current().domain,
            term_instance.vocab_list.id,
            term_instance.vocab_list
        )
        assert term_instance.get_vocab() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        create_term(
            name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        term_from_db = Term.objects.get(pk=1)

        assert term_from_db.name == 'spaces here'
        assert term_from_db.label == 'Spaces here too'

class TestProperty:
    def test_save(self):
        assert Property.objects.all().count() == 0
        create_property()
        assert Property.objects.all().count() == 1

    def test_get_vocab(self):
        property_instance = create_property()
        expected = "{}".format(
            Site.objects.get_current().domain,
            property_instance.term_key.vocab_list.id,
            property_instance.term_key.vocab_list
        )
        assert property_instance.get_vocab() == expected

    def test_get_term(self):
        property_instance = create_property()
        expected = "{}".format(
            Site.objects.get_current().domain,
            property_instance.term_key.id,
            property_instance.term_key
        )
        assert property_instance.get_term() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        create_property(
            property_name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        property_from_db = Property.objects.get(pk=1)

        assert property_from_db.property_name == 'spaces here'
        assert property_from_db.label == 'Spaces here too'
<commit_msg>Switched to using factories and made the re-used link a constant.<commit_after>import pytest
from django.contrib.sites.models import Site

from controlled_vocabularies.models import Vocabulary, Term, Property
from . import factories

pytestmark = pytest.mark.django_db

LINK= "{}"

class TestVocabulary:
    def test_save(self):
        vocab =  factories.VocabularyFactory.build()
        assert Vocabulary.objects.all().count() == 0
        vocab.save()
        assert Vocabulary.objects.all().count() == 1

    def test_unicode(self):
        vocab = factories.VocabularyFactory()
        assert unicode(vocab) == vocab.name

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.VocabularyFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
            maintainer=' Please   Fix Me ',
            maintainerEmail='  here@there.com ',
            definition='  Strip leading   and trailing, condense   others. '
        )
        vocab = Vocabulary.objects.get(pk=1)

        assert vocab.name == 'spaces here'
        assert vocab.label == 'Spaces here too'
        assert vocab.maintainer == 'Please Fix Me'
        assert vocab.maintainerEmail == 'here@there.com'
        assert vocab.definition == 'Strip leading and trailing, condense others.'

class TestTerm:
    def test_unicode(self):
        term = factories.TermFactory()
        assert unicode(term) == term.name

    def test_save(self):
        term = factories.TermFactory.build(vocab_list=factories.VocabularyFactory())
        assert Term.objects.all().count() == 0
        term.save()
        assert Term.objects.all().count() == 1

    def test_get_vocab(self):
        term = factories.TermFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            term.vocab_list.id,
            term.vocab_list
        )
        assert term.get_vocab() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.TermFactory(
            name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        term = Term.objects.get(pk=1)

        assert term.name == 'spaces here'
        assert term.label == 'Spaces here too'

class TestProperty:
    def test_save(self):
        prop = factories.PropertyFactory.build(term_key=factories.TermFactory())
        assert Property.objects.all().count() == 0
        prop.save()
        assert Property.objects.all().count() == 1

    def test_get_vocab(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'vocabulary',
            prop.term_key.vocab_list.id,
            prop.term_key.vocab_list
        )
        assert prop.get_vocab() == expected

    def test_get_term(self):
        prop = factories.PropertyFactory()
        expected = LINK.format(
            Site.objects.get_current().domain,
            'term',
            prop.term_key.id,
            prop.term_key
        )
        assert prop.get_term() == expected

    @pytest.mark.xfail(reason='Input is not being cleaned before being saved.')
    def test_whitespace_stripped(self):
        factories.PropertyFactory(
            property_name=' spaces  here ',
            label=' Spaces here   too  ',
        )
        prop = Property.objects.get(pk=1)

        assert prop.property_name == 'spaces here'
        assert prop.label == 'Spaces here too'
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
import itertools

class Hyperparameters(object):
    """
    Define hyperparameters with a default value
    and optionaly a list of possible values.

    Parameters
    ----------
    name : `string`
        The name of the hyperparameter.
        expected hyperparameters: {'batch_size', 'n_epochs', 'learning_rate', 'l1_reg', 'l2_reg', 'patience'}
    value : `float`
        The default value of the hyperparameter.
    range : `list` of `float`
        A list of values itterated over during the gris search

    Examples
    --------
    # Create an Hyperparameters instance
    >>>hp = Hyperparameters()
    # Define an hyperparameter
    >>>>hp('batch_size', 500)
    # Define an hyperparameter with its default value
    # and the possible values for the grid search
    >>>hp('n_epochs', 1000, [10, 100, 1000, 1000])

    """
    def __init__(self):
        self.hp_value = dict()
        self.hp_default = dict()
        self.hp_range = dict()
        self.iteration = 0

    def __call__(self, name, value, range=None):
        self.__setattr__(name, value)
        self.hp_value[name] = value
        self.hp_default[name] = value
        self.hp_range[name] = range
        if not range:
            self.hp_range[name] = [value]
        product = [x for x in apply(itertools.product, self.hp_range.values())]
        self.hp_product = [dict(zip(self.hp_value.keys(), p)) for p in product]

    def __str__(self):
        return str(self.hp_value)

    def __iter__(self):
        return self

    def next(self):
        if self.iteration > len(self.hp_product) - 1:
            raise StopIteration
        self.hp_value = self.hp_product[self.iteration]
        for name, value in self.hp_value.iteritems():
            self.__setattr__(name, value)
        self.iteration += 1
        return self

    def reset(self):
        for name, value in self.hp_default.iteritems():
            self.__setattr__(name, value)
        self.iteration = 0

<commit_msg>dl first commit<commit_after># -*- coding: UTF-8 -*-
import itertools

class Hyperparameters(object):
    """
    Define hyperparameters with a default value
    and optionaly a list of possible values.

    Parameters
    ----------
    name : `string`
        The name of the hyperparameter.
        expected hyperparameters: {'batch_size', 'n_epochs', 'learning_rate', 'l1_reg', 'l2_reg', 'patience'}
    value : `float`
        The default value of the hyperparameter.
    range : `list` of `float`
        A list of values itterated over during the gris search

    Examples
    --------
    >>>hp = Hyperparameters()   # Create an Hyperparameters instance
    >>>hp('batch_size', 500)    # Define an hyperparameter with its default value
    >>>hp('n_epochs', 1000, [10, 100, 1000, 1000])    # default value and range for the grid search

    """
    def __init__(self):
        self.hp_value = dict()
        self.hp_default = dict()
        self.hp_range = dict()
        self.iteration = 0

    def __call__(self, name, value, range=None):
        self.__setattr__(name, value)
        self.hp_value[name] = value
        self.hp_default[name] = value
        self.hp_range[name] = range
        if not range:
            self.hp_range[name] = [value]
        product = [x for x in apply(itertools.product, self.hp_range.values())]
        self.hp_product = [dict(zip(self.hp_value.keys(), p)) for p in product]

    def __str__(self):
        return str(self.hp_value)

    def __iter__(self):
        return self

    def next(self):
        if self.iteration > len(self.hp_product) - 1:
            raise StopIteration
        self.hp_value = self.hp_product[self.iteration]
        for name, value in self.hp_value.iteritems():
            self.__setattr__(name, value)
        self.iteration += 1
        return self

    def reset(self):
        for name, value in self.hp_default.iteritems():
            self.__setattr__(name, value)
        self.iteration = 0

<|endoftext|>


<commit_before>from datetime import datetime

from discord.ext import commands
import discord

from db.query_wrappers import get_event_role, get_event_delete_role
from cogs.utils.messages import delete_all, MessageManager
from cogs.utils.checks import is_event, is_int
from cogs.utils import constants
from cogs.utils.format import format_role_name

class Events:

    def __init__(self, bot):
        self.bot = bot

    @commands.command()
    @commands.guild_only()
    @commands.cooldown(rate=2, per=5, type=commands.BucketType.user)
    async def event(self, ctx):
        """
        Create an event in the events channel

        After invoking the event command, the bot will ask
        you to enter the event details. Once the event
        is created, it will appear in the upcoming-events
        channel. The upcoming-events channel is designed with
        the assumption that it isn't used for anything but
        displaying events; non event messages will be deleted.

        Users will be able to accept and decline the
        event by adding reactions. If a maximum number
        of attendees is set and the event is full,
        additional attendees will be placed in a standby
        section. If a spot opens up, the user at the
        top of the standby section will be automatically
        moved into the event.

        By default, everyone can make events. However, a minimum role
        requirement to create events can be defined in the settings.
        See `help settings seteventrole` for more information.

        The event creator and those with the Manage Sever permission
        can delete events by reacting to the event message with \U0001f480.
        """
        manager = MessageManager(self.bot, ctx.author, ctx.channel, ctx.prefix, [ctx.message])
        event_role = get_event_role(self.bot, ctx.guild)
        member_permissions = ctx.author.permissions_in(ctx.channel)

        if event_role:
            if ctx.author.top_role < event_role:
                event_role_str = format_role_name(event_role)
                await manager.say("You must be of role **{}** or higher to do that.".format(event_role))
                return await manager.clear()

        await manager.say('Event creation instructions have been messaged to you')

        res = await manager.say_and_wait("Enter event title:", dm=True)
        if not res:
            return await manager.clear()
        title = res.content

        description = ""
        res = await manager.say_and_wait("Enter event description (type 'none' for no description):", dm=True)
        if not res:
            return await manager.clear()
        if res.content.upper() != 'NONE':
            description = res.content

        max_members = 0
        while not max_members:
            res = await manager.say_and_wait("Enter the maximum numbers of attendees (type 'none' for no maximum):", dm=True)
            if not res:
                return await manager.clear()
            if res.content.upper() == 'NONE':
                break
            elif is_int(res.content) and int(res.content) in range(1,10000):
                max_members = int(res.content)
            else:
                await manager.say("Invalid entry. Must be a number between 1 and 9999.", dm=True)

        start_time = None
        while not start_time:
            res = await manager.say_and_wait("Enter event time (YYYY-MM-DD HH:MM AM/PM):", dm=True)
            if not res:
                return await manager.clear()
            start_time_format = '%Y-%m-%d %I:%M %p'
            try:
                start_time = datetime.strptime(res.content, start_time_format)
            except ValueError:
                await manager.say("Invalid event time!", dm=True)

        time_zone = None
        while not time_zone:
            res = await manager.say_and_wait("Enter the time zone (PST, EST, etc):", dm=True)
            if not res:
                return await manager.clear()
            user_timezone = "".join(res.content.upper().split())
            if user_timezone not in constants.TIME_ZONES:
                await manager.say("Unsupported time zone", dm=True)
            else:
                time_zone = user_timezone

        affected_rows = self.bot.db.create_event(title, start_time, time_zone, ctx.guild.id, description, max_members, ctx.author.id)
        if affected_rows == 0:
            await manager.say("An event with that name already exists!", dm=True)
            return await manager.clear()

        event_channel = await self.get_events_channel(ctx.guild)
        await manager.say("Event created! The " + event_channel.mention + " channel will be updated momentarily.", dm=True)
        await self.list_events(ctx.guild)
        await manager.clear()

    async def list_events(self, guild):
        """Clear the event channel and display all upcoming events"""
        events_channel = await self.get_events_channel(guild)
        await events_channel.purge(limit=999, check=delete_all)
        events = self.bot.db.get_events(guild.id)

        if len(events) > 0:
            for event in events:
                event_embed = self.create_event_embed(guild, event)
                msg = await events_channel.send(embed=event_embed)
                await msg.add_reaction("\N{WHITE HEAVY CHECK MARK}")
                await msg.add_reaction("\N{CROSS MARK}")
        else:
            await events_channel.send("There are no upcoming events.")

    async def on_raw_reaction_add(self, emoji, message_id, channel_id, user_id):
        """If a reaction represents a user RSVP, update the DB and event message"""
        channel = self.bot.get_channel(channel_id)
        if isinstance(channel, discord.abc.PrivateChannel):
            return
        try:
            message = await channel.get_message(message_id)
        except:
            return

        guild = channel.guild
        member = guild.get_member(user_id)
        deleted = None

        # We check that the user is not the message author as to not count
        # the initial reactions added by the bot as being indicative of attendance
        if is_event(message) and member != message.author:
            title = message.embeds[0].title
            if emoji.name == "\N{WHITE HEAVY CHECK MARK}":
                await self.set_attendance(member, guild, 1, title, message)
            elif emoji.name == "\N{CROSS MARK}":
                await self.set_attendance(member, guild, 0, title, message)
            elif emoji.name == "\N{SKULL}":
                deleted =  await self.delete_event(guild, title, member, channel)

            if not deleted:
                try:
                    await message.remove_reaction(emoji, member)
                except:
                    pass

    async def set_attendance(self, member, guild, attending, title, message):
        """Send updated event attendance info to db and update the event"""
        self.bot.db.add_user(member.id)
        self.bot.db.update_attendance(member.id, guild.id, attending, title, datetime.now())

        # Update event message in place for a more seamless user experience
        event = self.bot.db.get_event(guild.id, title)
        if event:
            event_embed = self.create_event_embed(guild, event)
            await message.edit(embed=event_embed)
        else:
            raise ValueError("Could not retrieve event")
            return

    async def delete_event(self, guild, title, member, channel):
        """Delete an event and update the events channel on success"""
        event_delete_role = get_event_delete_role(self.bot, guild)
        result = self.bot.db.get_event_creator(guild.id, title)
        creator_id = result.get('user_id') if result else None

        if member.permissions_in(channel).manage_guild or (member.id == creator_id) or (event_delete_role and member.top_role >= event_delete_role):
            deleted = self.bot.db.delete_event(guild.id, title)
            if deleted:
                await self.list_events(guild)
                return True
        else:
            try:
                await member.send("You don't have permission to delete that event.")
            except:
                pass

    async def get_events_channel(self, guild):
        """Return the events channel if it exists, otherwise create one and return it"""
        for channel in guild.channels:
            if channel.name == "upcoming-events":
                return channel

        # Need to make sure the bot can still send messages in the events channel
        overwrites = {
            guild.default_role: discord.PermissionOverwrite(send_messages=False, add_reactions=True),
            guild.me: discord.PermissionOverwrite(send_messages=True, add_reactions=True)
        }
        return await guild.create_text_channel("upcoming-events", overwrites=overwrites)

    def create_event_embed(self, guild, event):
        """Create and return a Discord Embed object that represents an upcoming event"""
        title = event.get('event_title')
        description = event.get('description')
        time = event.get('start_time')
        timezone = event.get('timezone')
        creator_id = event.get('user_id')
        accepted = event.get('accepted')
        declined = event.get('declined')
        max_members = event.get('max_members')

        embed_msg = discord.Embed(color=constants.BLUE)
        embed_msg.title = title

        creator = guild.get_member(creator_id)
        if creator:
            embed_msg.set_footer(text="Created by {} | React with {} to remove this event".format(creator.display_name, '\U0001f480'))
        else:
            embed_msg.set_footer(text="React with {} to remove this event".format('\U0001f480'))

        if description:
            embed_msg.description = description
        time_str = time.strftime("%A %b %-d, %Y @ %-I:%M %p")
        embed_msg.add_field(name="Time", value=time_str + " " + timezone, inline=False)

        if accepted:
            accepted_list = None
            if max_members:
                accepted_list = accepted.split(',')[:max_members]
            else:
                accepted_list = accepted.split(',')
            text = ""
            for user_id in accepted_list:
                member = guild.get_member(int(user_id))
                if member:
                    text += "{}\n".format(member.display_name)
            if max_members:
                embed_msg.add_field(name="Accepted ({}/{})".format(len(accepted_list), max_members), value=text)
            else:
                embed_msg.add_field(name="Accepted", value=text)
        else:
            if max_members:
                embed_msg.add_field(name="Accepted (0/{})".format(max_members), value="-")
            else:
                embed_msg.add_field(name="Accepted", value="-")

        if declined:
            declined_list = declined.split(',')
            text = ""
            for user_id in declined_list:
                member = guild.get_member(int(user_id))
                if member:
                    text += "{}\n".format(member.display_name)
            embed_msg.add_field(name="Declined", value=text)
        else:
            embed_msg.add_field(name="Declined", value="-")

        if accepted and max_members:
            standby_list = accepted.split(',')[max_members:]
            if standby_list:
                text = ""
                for user_id in standby_list:
                    member = guild.get_member(int(user_id))
                    if member:
                        text += "{}\n".format(member.display_name)
                embed_msg.add_field(name="Standby", value=text, inline=False)

        return embed_msg
<commit_msg>add debug code<commit_after>from datetime import datetime

from discord.ext import commands
import discord

from db.query_wrappers import get_event_role, get_event_delete_role
from cogs.utils.messages import delete_all, MessageManager
from cogs.utils.checks import is_event, is_int
from cogs.utils import constants
from cogs.utils.format import format_role_name

class Events:

    def __init__(self, bot):
        self.bot = bot

    @commands.command()
    @commands.guild_only()
    @commands.cooldown(rate=2, per=5, type=commands.BucketType.user)
    async def event(self, ctx):
        """
        Create an event in the events channel

        After invoking the event command, the bot will ask
        you to enter the event details. Once the event
        is created, it will appear in the upcoming-events
        channel. The upcoming-events channel is designed with
        the assumption that it isn't used for anything but
        displaying events; non event messages will be deleted.

        Users will be able to accept and decline the
        event by adding reactions. If a maximum number
        of attendees is set and the event is full,
        additional attendees will be placed in a standby
        section. If a spot opens up, the user at the
        top of the standby section will be automatically
        moved into the event.

        By default, everyone can make events. However, a minimum role
        requirement to create events can be defined in the settings.
        See `help settings seteventrole` for more information.

        The event creator and those with the Manage Sever permission
        can delete events by reacting to the event message with \U0001f480.
        """
        manager = MessageManager(self.bot, ctx.author, ctx.channel, ctx.prefix, [ctx.message])
        event_role = get_event_role(self.bot, ctx.guild)
        member_permissions = ctx.author.permissions_in(ctx.channel)

        if event_role:
            if ctx.author.top_role < event_role:
                event_role_str = format_role_name(event_role)
                await manager.say("You must be of role **{}** or higher to do that.".format(event_role))
                return await manager.clear()

        await manager.say('Event creation instructions have been messaged to you')

        res = await manager.say_and_wait("Enter event title:", dm=True)
        if not res:
            return await manager.clear()
        title = res.content

        description = ""
        res = await manager.say_and_wait("Enter event description (type 'none' for no description):", dm=True)
        if not res:
            return await manager.clear()
        if res.content.upper() != 'NONE':
            description = res.content

        max_members = 0
        while not max_members:
            res = await manager.say_and_wait("Enter the maximum numbers of attendees (type 'none' for no maximum):", dm=True)
            if not res:
                return await manager.clear()
            if res.content.upper() == 'NONE':
                break
            elif is_int(res.content) and int(res.content) in range(1,10000):
                max_members = int(res.content)
            else:
                await manager.say("Invalid entry. Must be a number between 1 and 9999.", dm=True)

        start_time = None
        while not start_time:
            res = await manager.say_and_wait("Enter event time (YYYY-MM-DD HH:MM AM/PM):", dm=True)
            if not res:
                return await manager.clear()
            start_time_format = '%Y-%m-%d %I:%M %p'
            try:
                start_time = datetime.strptime(res.content, start_time_format)
            except ValueError:
                await manager.say("Invalid event time!", dm=True)

        time_zone = None
        while not time_zone:
            res = await manager.say_and_wait("Enter the time zone (PST, EST, etc):", dm=True)
            if not res:
                return await manager.clear()
            user_timezone = "".join(res.content.upper().split())
            if user_timezone not in constants.TIME_ZONES:
                await manager.say("Unsupported time zone", dm=True)
            else:
                time_zone = user_timezone

        affected_rows = self.bot.db.create_event(title, start_time, time_zone, ctx.guild.id, description, max_members, ctx.author.id)
        if affected_rows == 0:
            await manager.say("An event with that name already exists!", dm=True)
            return await manager.clear()

        event_channel = await self.get_events_channel(ctx.guild)
        await manager.say("Event created! The " + event_channel.mention + " channel will be updated momentarily.", dm=True)
        await self.list_events(ctx.guild)
        await manager.clear()

    async def list_events(self, guild):
        """Clear the event channel and display all upcoming events"""
        events_channel = await self.get_events_channel(guild)
        await events_channel.purge(limit=999, check=delete_all)
        events = self.bot.db.get_events(guild.id)

        if len(events) > 0:
            for event in events:
                event_embed = self.create_event_embed(guild, event)
                msg = await events_channel.send(embed=event_embed)
                await msg.add_reaction("\N{WHITE HEAVY CHECK MARK}")
                await msg.add_reaction("\N{CROSS MARK}")
        else:
            await events_channel.send("There are no upcoming events.")

    async def on_raw_reaction_add(self, emoji, message_id, channel_id, user_id):
        """If a reaction represents a user RSVP, update the DB and event message"""
        channel = self.bot.get_channel(channel_id)
        if isinstance(channel, discord.abc.PrivateChannel):
            return
        try:
            message = await channel.get_message(message_id)
        except:
            return

        guild = channel.guild
        member = guild.get_member(user_id)
        deleted = None

        # We check that the user is not the message author as to not count
        # the initial reactions added by the bot as being indicative of attendance
        if is_event(message) and member != message.author:
            title = message.embeds[0].title
            if emoji.name == "\N{WHITE HEAVY CHECK MARK}":
                await self.set_attendance(member, guild, 1, title, message)
            elif emoji.name == "\N{CROSS MARK}":
                await self.set_attendance(member, guild, 0, title, message)
            elif emoji.name == "\N{SKULL}":
                deleted = await self.delete_event(guild, title, member, channel)

            if not deleted:
                try:
                    await message.remove_reaction(emoji, member)
                except:
                    pass

    async def set_attendance(self, member, guild, attending, title, message):
        """Send updated event attendance info to db and update the event"""
        self.bot.db.add_user(member.id)
        self.bot.db.update_attendance(member.id, guild.id, attending, title, datetime.now())

        # Update event message in place for a more seamless user experience
        event = self.bot.db.get_event(guild.id, title)
        if event:
            event_embed = self.create_event_embed(guild, event)
            await message.edit(embed=event_embed)
        else:
            raise ValueError("Could not retrieve event")
            return

    async def delete_event(self, guild, title, member, channel):
        """Delete an event and update the events channel on success"""
        event_delete_role = get_event_delete_role(self.bot, guild)
        result = self.bot.db.get_event_creator(guild.id, title)
        creator_id = result.get('user_id') if result else None

        if member.permissions_in(channel).manage_guild or (member.id == creator_id) or (event_delete_role and member.top_role >= event_delete_role):
            deleted = self.bot.db.delete_event(guild.id, title)
            if deleted:

                if guild.id == 351073221241470976:
                    asal = self.bot.get_user(118926942404608003)
                    await asal.send("Event **{}** deleted by **{}**".format(title, str(member)))

                await self.list_events(guild)
                return True
        else:
            try:
                await member.send("You don't have permission to delete that event.")
            except:
                pass

    async def get_events_channel(self, guild):
        """Return the events channel if it exists, otherwise create one and return it"""
        for channel in guild.channels:
            if channel.name == "upcoming-events":
                return channel

        # Need to make sure the bot can still send messages in the events channel
        overwrites = {
            guild.default_role: discord.PermissionOverwrite(send_messages=False, add_reactions=True),
            guild.me: discord.PermissionOverwrite(send_messages=True, add_reactions=True)
        }
        return await guild.create_text_channel("upcoming-events", overwrites=overwrites)

    def create_event_embed(self, guild, event):
        """Create and return a Discord Embed object that represents an upcoming event"""
        title = event.get('event_title')
        description = event.get('description')
        time = event.get('start_time')
        timezone = event.get('timezone')
        creator_id = event.get('user_id')
        accepted = event.get('accepted')
        declined = event.get('declined')
        max_members = event.get('max_members')

        embed_msg = discord.Embed(color=constants.BLUE)
        embed_msg.title = title

        creator = guild.get_member(creator_id)
        if creator:
            embed_msg.set_footer(text="Created by {} | React with {} to remove this event".format(creator.display_name, '\U0001f480'))
        else:
            embed_msg.set_footer(text="React with {} to remove this event".format('\U0001f480'))

        if description:
            embed_msg.description = description
        time_str = time.strftime("%A %b %-d, %Y @ %-I:%M %p")
        embed_msg.add_field(name="Time", value=time_str + " " + timezone, inline=False)

        if accepted:
            accepted_list = None
            if max_members:
                accepted_list = accepted.split(',')[:max_members]
            else:
                accepted_list = accepted.split(',')
            text = ""
            for user_id in accepted_list:
                member = guild.get_member(int(user_id))
                if member:
                    text += "{}\n".format(member.display_name)
            if max_members:
                embed_msg.add_field(name="Accepted ({}/{})".format(len(accepted_list), max_members), value=text)
            else:
                embed_msg.add_field(name="Accepted", value=text)
        else:
            if max_members:
                embed_msg.add_field(name="Accepted (0/{})".format(max_members), value="-")
            else:
                embed_msg.add_field(name="Accepted", value="-")

        if declined:
            declined_list = declined.split(',')
            text = ""
            for user_id in declined_list:
                member = guild.get_member(int(user_id))
                if member:
                    text += "{}\n".format(member.display_name)
            embed_msg.add_field(name="Declined", value=text)
        else:
            embed_msg.add_field(name="Declined", value="-")

        if accepted and max_members:
            standby_list = accepted.split(',')[max_members:]
            if standby_list:
                text = ""
                for user_id in standby_list:
                    member = guild.get_member(int(user_id))
                    if member:
                        text += "{}\n".format(member.display_name)
                embed_msg.add_field(name="Standby", value=text, inline=False)

        return embed_msg
<|endoftext|>