<commit_before>#--output_dir=/Users/piotr.milos/Downloads/18 --alsologtostderr --loop_hparams_set=rl_modelrl_base_quick
#--output_dir=/Users/piotr.milos/t2t/rl_v1 --alsologtostderr --loop_hparams_set=rl_modelrl_tiny

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import copy
import os
from gym.spaces import Box
import numpy as np

import six
from gym.spaces import Discrete
from gym.utils.play import PlayPlot

from tensor2tensor.bin import t2t_trainer
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import gym_problems_specs
from tensor2tensor.layers import discretization
from tensor2tensor.rl import rl_trainer_lib
from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
from tensor2tensor.rl.envs.utils import get_policy
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib
from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
from gym.utils import play
import tensorflow as tf

HP_SCOPES = ["loop", "model", "ppo"]

_font = None
FONT_SIZE = 20

def _get_font():
  global _font
  if _font is None:
    #weirdness due to various working dirs
    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "deepsense_experiments/Xerox Serif Narrow.ttf"]

    for path in FONT_PATHS:
      try:
        _font = ImageFont.truetype(path, FONT_SIZE)
        return _font
      except:
        pass

def _assert_image(img):
  if isinstance(img, np.ndarray):
    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
  return img

def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
  img = _assert_image(img)
  if text=="":
    return img
  draw = ImageDraw.Draw(img)
  font = _get_font()
  draw.text(position, text, color, font=font)

  return img

def concatenate_images(*imgs, axis=1):
  imgs = [_assert_image(img) for img in imgs]
  imgs_np = [np.array(img) for img in imgs]
  concatenated_im_np = np.concatenate(imgs_np, axis=axis)

  return _assert_image(concatenated_im_np)

def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams):
  """Train the PPO agent in the simulated environment."""
  gym_problem = registry.problem(problem_name)
  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)

  # Adding model hparams for model specific adjustments
  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
  ppo_hparams.add_hparam("model_hparams", model_hparams)

  environment_spec = copy.copy(gym_problem.environment_spec)
  environment_spec.simulation_random_starts = hparams.simulation_random_starts
  environment_spec.simulation_flip_first_random_for_beginning = False
  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale

  ppo_hparams.add_hparam("environment_spec", environment_spec)
  ppo_hparams.num_agents = 1

  with temporary_flags({
      "problem": problem_name,
      "model": hparams.generative_model,
      "hparams_set": hparams.generative_model_params,
      "output_dir": world_model_dir,
      "data_dir": epoch_data_dir,
  }):

    sess = tf.Session()
    env = DebugBatchEnv(ppo_hparams, sess)
    sess.run(tf.global_variables_initializer())
    env.initialize()

    env_model_loader = tf.train.Saver(
      tf.global_variables("next_frame*"))
    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
      must_restore=True)

    model_saver = tf.train.Saver(
      tf.global_variables(".*network_parameters.*"))
    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)

    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                   (ord('r'),):101,
                   (ord('p'),):102}

    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)

from gym.core import Env

class DebugBatchEnv(Env):

  def __init__(self, hparams, sess = None):
    if sess == None:
      self.sess = tf.Session()
    else:
      self.sess = sess

    self.action_space = Discrete(6)
    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
    self._tmp = 1
    self.res = None
    self._prepare_networks(hparams, self.sess)

  def _prepare_networks(self, hparams, sess):
    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
    batch_env = batch_env_factory(hparams)
    self.reward, self.done = batch_env.simulate(self.action)
    self.observation = batch_env.observ
    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))

    environment_wrappers = hparams.environment_spec.wrappers
    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []

    to_initialize = [batch_env]
    for w in wrappers:
      batch_env = w[0](batch_env, **w[1])
      to_initialize.append(batch_env)

    def initialization_lambda():
      for batch_env in to_initialize:
        batch_env.initialize(sess)

    self.initialize = initialization_lambda

    obs_copy = batch_env.observ + 0

    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
    self.policy_probs = actor_critic.policy.probs[0, 0, :]
    self.value = actor_critic.value[0, :]

  def render(self, mode='human'):
    raise NotImplemented()

  def _fake_reset(self):
    self._tmp = 0
    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def _reset_env(self):
    _observ = self.sess.run(self.reset_op)[0, ...]
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    #TODO:(put correct numbers)
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def reset(self):
    self._reset_env()
    observ = self._augment_observation()
    return observ

  def _step_fake(self, action):
    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
    observ[0, 0, 0] = 0
    observ[0, 0, 1] = 255

    self._tmp += 1
    if self._tmp>20:
      self._tmp = 0

    rew = 1
    done = False
    probs = np.ones(shape=(6,), dtype=np.float32)/6
    vf = 0.0

    return observ, rew, done, probs, vf

  def _step_env(self, action):
    observ, rew, done, probs, vf = self.sess.\
      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
          feed_dict={self.action: [action]})

    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf

  def _augment_observation(self):
    _observ, rew, done, probs, vf = self.res
    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
    probs_str = ""
    for p in probs:
      probs_str += "%.2f" % p +", "

    action = np.argmax(probs)

    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                           vf, rew)
    info_pane = write_on_image(info_pane, info_str)

    augmented_observ = concatenate_images(_observ, info_pane)
    augmented_observ = np.array(augmented_observ)
    return augmented_observ

  def step(self, action):
    #Special codes
    if action==100:
      #Skip action
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 101:
      #reset
      self.reset()
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 102:
      #play
      raise NotImplemented()

    #standard codes
    _observ, rew, done, probs, vf = self._step_env(action)
    self.res = (_observ, rew, done, probs, vf)

    observ = self._augment_observation()
    return observ, rew, done, {"probs": probs, "vf": vf}

def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
  """Run the main training loop."""

  # Directories
  subdirectories = ["data", "tmp", "world_model", "ppo"]
  using_autoencoder = hparams.autoencoder_train_steps > 0
  if using_autoencoder:
    subdirectories.append("autoencoder")
  directories = setup_directories(output_dir, subdirectories)

  if hparams.game in gym_problems_specs.ATARI_GAMES:
    game_with_mode = hparams.game + "_deterministic-v4"
  else:
    game_with_mode = hparams.game
  # Problems
  if using_autoencoder:
    simulated_problem_name = (
        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
        % game_with_mode)
  else:
    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                              % game_with_mode)

  epoch = hparams.epochs-1
  epoch_data_dir = os.path.join(directories["data"], str(epoch))

  ppo_event_dir = os.path.join(directories["world_model"],
                               "ppo_summaries", str(epoch))
  ppo_model_dir = directories["ppo"]
  if not hparams.ppo_continue_training:
    ppo_model_dir = ppo_event_dir
  show_agent(simulated_problem_name, ppo_model_dir, directories["world_model"], epoch_data_dir, hparams)

def create_loop_hparams():
  hparams = registry.hparams(FLAGS.loop_hparams_set)
  hparams.parse(FLAGS.loop_hparams)
  return hparams

def main(_):
  hp = create_loop_hparams()
  output_dir = FLAGS.output_dir
  training_loop(hp, output_dir)

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run()
<commit_msg>more cleanups<commit_after>#--output_dir=/Users/piotr.milos/Downloads/18 --alsologtostderr --loop_hparams_set=rl_modelrl_base_quick
#--output_dir=/Users/piotr.milos/t2t/rl_v1 --alsologtostderr --loop_hparams_set=rl_modelrl_tiny

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import copy
import os
from gym.spaces import Box
import numpy as np

import six
from gym.spaces import Discrete
from gym.utils.play import PlayPlot

from tensor2tensor.bin import t2t_trainer
from tensor2tensor.data_generators import generator_utils
from tensor2tensor.data_generators import gym_problems_specs
from tensor2tensor.layers import discretization
from tensor2tensor.rl import rl_trainer_lib
from tensor2tensor.rl.envs.batch_env_factory import batch_env_factory
from tensor2tensor.rl.envs.utils import get_policy
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib
from tensor2tensor.rl.trainer_model_based import FLAGS, setup_directories, temporary_flags
from gym.utils import play
import tensorflow as tf

HP_SCOPES = ["loop", "model", "ppo"]

_font = None
FONT_SIZE = 20

def _get_font():
  global _font
  if _font is None:
    #weirdness due to various working dirs
    FONT_PATHS = ["tensor-2-tensor-with-mrunner/tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "tensor-2-tensor-with-mrunner/deepsense_experiments/Xerox Serif Narrow.ttf",
                  "deepsense_experiments/Xerox Serif Narrow.ttf"]

    for path in FONT_PATHS:
      try:
        _font = ImageFont.truetype(path, FONT_SIZE)
        return _font
      except:
        pass

def _assert_image(img):
  if isinstance(img, np.ndarray):
    img = Image.fromarray(np.ndarray.astype(img, np.uint8))
  return img

def write_on_image(img, text="", position=(0, 0), color=(255, 255, 255)):
  img = _assert_image(img)
  if text=="":
    return img
  draw = ImageDraw.Draw(img)
  font = _get_font()
  draw.text(position, text, color, font=font)

  return img

def concatenate_images(*imgs, axis=1):
  imgs = [_assert_image(img) for img in imgs]
  imgs_np = [np.array(img) for img in imgs]
  concatenated_im_np = np.concatenate(imgs_np, axis=axis)

  return _assert_image(concatenated_im_np)

def show_agent(problem_name, agent_model_dir, world_model_dir, epoch_data_dir, hparams):
  """Train the PPO agent in the simulated environment."""
  gym_problem = registry.problem(problem_name)
  ppo_hparams = trainer_lib.create_hparams(hparams.ppo_params)

  # Adding model hparams for model specific adjustments
  model_hparams = trainer_lib.create_hparams(hparams.generative_model_params)
  ppo_hparams.add_hparam("model_hparams", model_hparams)

  environment_spec = copy.copy(gym_problem.environment_spec)
  environment_spec.simulation_random_starts = hparams.simulation_random_starts
  environment_spec.simulation_flip_first_random_for_beginning = False
  environment_spec.intrinsic_reward_scale = hparams.intrinsic_reward_scale

  ppo_hparams.add_hparam("environment_spec", environment_spec)
  ppo_hparams.num_agents = 1

  with temporary_flags({
      "problem": problem_name,
      "model": hparams.generative_model,
      "hparams_set": hparams.generative_model_params,
      "output_dir": world_model_dir,
      "data_dir": epoch_data_dir,
  }):

    sess = tf.Session()
    env = DebugBatchEnv(ppo_hparams, sess)
    sess.run(tf.global_variables_initializer())
    env.initialize()

    env_model_loader = tf.train.Saver(
      tf.global_variables("next_frame*"))
    trainer_lib.restore_checkpoint(world_model_dir, env_model_loader, sess,
      must_restore=True)

    model_saver = tf.train.Saver(
      tf.global_variables(".*network_parameters.*"))
    trainer_lib.restore_checkpoint(agent_model_dir, model_saver, sess)

    key_mapping = {(): 100, (ord('q'),):1, (ord('a'),):2,
                   (ord('r'),):101,
                   (ord('p'),):102}

    play.play(env, zoom=2, fps=10, keys_to_action=key_mapping)

from gym.core import Env

class DebugBatchEnv(Env):

  def __init__(self, hparams, sess = None):
    if sess == None:
      self.sess = tf.Session()
    else:
      self.sess = sess

    self.action_space = Discrete(6)
    self.observation_space = Box(low=0, high=255, shape=(210, 160+250, 3), dtype=np.uint8)
    self._tmp = 1
    self.res = None
    self._prepare_networks(hparams, self.sess)

  def _prepare_networks(self, hparams, sess):
    self.action = tf.placeholder(shape=(1,), dtype=tf.int32)
    batch_env = batch_env_factory(hparams)
    self.reward, self.done = batch_env.simulate(self.action)
    self.observation = batch_env.observ
    self.reset_op = batch_env.reset(tf.constant([0], dtype=tf.int32))

    environment_wrappers = hparams.environment_spec.wrappers
    wrappers = copy.copy(environment_wrappers) if environment_wrappers else []

    to_initialize = [batch_env]
    for w in wrappers:
      batch_env = w[0](batch_env, **w[1])
      to_initialize.append(batch_env)

    def initialization_lambda():
      for batch_env in to_initialize:
        batch_env.initialize(sess)

    self.initialize = initialization_lambda

    obs_copy = batch_env.observ + 0

    actor_critic = get_policy(tf.expand_dims(obs_copy, 0), hparams)
    self.policy_probs = actor_critic.policy.probs[0, 0, :]
    self.value = actor_critic.value[0, :]

  def render(self, mode='human'):
    raise NotImplemented()

  def _fake_reset(self):
    self._tmp = 0
    _observ = np.ones(shape=(210, 160, 3), dtype=np.uint8) * 10 * self._tmp
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def _reset_env(self):
    _observ = self.sess.run(self.reset_op)[0, ...]
    _observ[0, 0, 0] = 0
    _observ[0, 0, 1] = 255
    #TODO:(put correct numbers)
    self.res = (_observ, 0, False, [0.1, 0.5, 0.5], 1.1)

  def reset(self):
    self._reset_env()
    observ = self._augment_observation()
    return observ

  def _step_fake(self, action):
    observ = np.ones(shape=(210, 160, 3), dtype=np.uint8)*10*self._tmp
    observ[0, 0, 0] = 0
    observ[0, 0, 1] = 255

    self._tmp += 1
    if self._tmp>20:
      self._tmp = 0

    rew = 1
    done = False
    probs = np.ones(shape=(6,), dtype=np.float32)/6
    vf = 0.0

    return observ, rew, done, probs, vf

  def _step_env(self, action):
    observ, rew, done, probs, vf = self.sess.\
      run([self.observation, self.reward, self.done, self.policy_probs, self.value],
          feed_dict={self.action: [action]})

    return observ[0, ...], rew[0, ...], done[0, ...], probs, vf

  def _augment_observation(self):
    _observ, rew, done, probs, vf = self.res
    info_pane = np.zeros(shape=(210, 250, 3), dtype=np.uint8)
    probs_str = ""
    for p in probs:
      probs_str += "%.2f" % p +", "

    action = np.argmax(probs)

    info_str = " Policy:{}\n Action:{}\n Value function:{}\n Reward:{}".format(probs_str, action,
                                                                           vf, rew)
    info_pane = write_on_image(info_pane, info_str)

    augmented_observ = concatenate_images(_observ, info_pane)
    augmented_observ = np.array(augmented_observ)
    return augmented_observ

  def step(self, action):
    #Special codes
    if action==100:
      #Skip action
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 101:
      #reset
      self.reset()
      _, rew, done, _, _ = self.res
      observ = self._augment_observation()
      return observ, rew, done, {}

    if action == 102:
      #play
      raise NotImplemented()

    #standard codes
    _observ, rew, done, probs, vf = self._step_env(action)
    self.res = (_observ, rew, done, probs, vf)

    observ = self._augment_observation()
    return observ, rew, done, {"probs": probs, "vf": vf}

def setup_show_metadata(hparams, output_dir):
  """Setup metadata"""

  # Directories
  subdirectories = ["data", "tmp", "world_model", "ppo"]
  using_autoencoder = hparams.autoencoder_train_steps > 0
  if using_autoencoder:
    subdirectories.append("autoencoder")
  directories = setup_directories(output_dir, subdirectories)

  if hparams.game in gym_problems_specs.ATARI_GAMES:
    game_with_mode = hparams.game + "_deterministic-v4"
  else:
    game_with_mode = hparams.game
  # Problems
  if using_autoencoder:
    simulated_problem_name = (
        "gym_simulated_discrete_problem_with_agent_on_%s_autoencoded"
        % game_with_mode)
  else:
    simulated_problem_name = ("gym_simulated_discrete_problem_with_agent_on_%s"
                              % game_with_mode)
  epoch = hparams.epochs-1
  epoch_data_dir = os.path.join(directories["data"], str(epoch))
  ppo_model_dir = directories["ppo"]

  show_agent(simulated_problem_name, ppo_model_dir,
             directories["world_model"], epoch_data_dir, hparams)

def main(_):
  hparams = registry.hparams(FLAGS.loop_hparams_set)
  hparams.parse(FLAGS.loop_hparams)
  output_dir = FLAGS.output_dir
  setup_show_metadata(hparams, output_dir)

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
  tf.app.run()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# PySPED - Python libraries to deal with Brazil's SPED Project
#
# Copyright (C) 2010-2012
# Copyright (C) Aristides Caldeira <aristides.caldeira at tauga.com.br>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# PySPED - Bibliotecas Python para o
#          SPED - Sistema Público de Escrituração Digital
#
# Copyright (C) 2010-2012
# Copyright (C) Aristides Caldeira <aristides.caldeira arroba tauga.com.br>
#
# Este programa é um software livre: você pode redistribuir e/ou modificar
# este programa sob os termos da licença GNU Affero General Public License,
# publicada pela Free Software Foundation, em sua versão 3 ou, de acordo
# com sua opção, qualquer versão posterior.
#
# Este programa é distribuido na esperança de que venha a ser útil,
# porém SEM QUAISQUER GARANTIAS, nem mesmo a garantia implícita de
# COMERCIABILIDADE ou ADEQUAÇÃO A UMA FINALIDADE ESPECÍFICA. Veja a
# GNU Affero General Public License para mais detalhes.
#
# Você deve ter recebido uma cópia da GNU Affero General Public License
# juntamente com este programa. Caso esse não seja o caso, acesse:
# <http://www.gnu.org/licenses/>
#

from __future__ import division, print_function, unicode_literals

from pysped.xml_sped import *
from pysped.nfe.leiaute import ESQUEMA_ATUAL_VERSAO_3 as ESQUEMA_ATUAL
from pysped.nfe.leiaute import consrecinfe_200
import os
from nfe_310 import NFe

DIRNAME = os.path.dirname(__file__)

class ConsReciNFe(consrecinfe_200.ConsReciNFe):
    def __init__(self):
        super(ConsReciNFe, self).__init__()
        self.versao  = TagDecimal(nome='consReciNFe', codigo='BP02', propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'consReciNFe_v3.10.xsd'

class InfProt(consrecinfe_200.InfProt):
    def __init__(self):
        super(InfProt, self).__init__()
        self.dhRecbto = TagDataHoraUTC(nome='dhRecbto', codigo='PR08', raiz='//infProt')

class ProtNFe(consrecinfe_200.ProtNFe):
    def __init__(self):
        super(ProtNFe, self).__init__()
        self.versao  = TagDecimal(nome='protNFe', codigo='PR02' , propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.infProt = InfProt()

class RetConsReciNFe(consrecinfe_200.RetConsReciNFe):
    def __init__(self):
        super(RetConsReciNFe, self).__init__()
        self.versao   = TagDecimal(nome='retConsReciNFe', codigo='BR02' , propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'retConsReciNFe_v3.10.xsd'

    def get_xml(self):
        xml = XMLNFe.get_xml(self)
        xml += self.versao.xml
        xml += self.tpAmb.xml
        xml += self.verAplic.xml
        xml += self.nRec.xml
        xml += self.cStat.xml
        xml += self.xMotivo.xml
        xml += self.cUF.xml
        xml += self.cMsg.xml
        xml += self.xMsg.xml

        for pn in self.protNFe:
            xml += pn.xml

        xml += '</retConsReciNFe>'
        return xml

    def set_xml(self, arquivo):
        if self._le_xml(arquivo):
            self.versao.xml   = arquivo
            self.tpAmb.xml    = arquivo
            self.verAplic.xml = arquivo
            self.nRec.xml     = arquivo
            self.cStat.xml    = arquivo
            self.xMotivo.xml  = arquivo
            self.cUF.xml      = arquivo
            self.cMsg.xml     = arquivo
            self.xMsg.xml     = arquivo
            self.protNFe      = self.le_grupo('//retConsReciNFe/protNFe', ProtNFe)

            #
            # Monta o dicionário dos protocolos
            #
            for pn in self.protNFe:
                self.dic_protNFe[pn.infProt.chNFe.valor] = pn

    xml = property(get_xml, set_xml)

class ProcNFe(consrecinfe_200.ProcNFe):
    def __init__(self):
        super(ProcNFe, self).__init__()
        self.versao  = TagDecimal(nome='nfeProc', propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.NFe     = NFe()
        self.protNFe = ProtNFe()
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'procNFe_v3.10.xsd'
<commit_msg>Corrigindo tratamento do fuso horário no retorno da recibo da NF-e<commit_after># -*- coding: utf-8 -*-
#
# PySPED - Python libraries to deal with Brazil's SPED Project
#
# Copyright (C) 2010-2012
# Copyright (C) Aristides Caldeira <aristides.caldeira at tauga.com.br>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# PySPED - Bibliotecas Python para o
#          SPED - Sistema Público de Escrituração Digital
#
# Copyright (C) 2010-2012
# Copyright (C) Aristides Caldeira <aristides.caldeira arroba tauga.com.br>
#
# Este programa é um software livre: você pode redistribuir e/ou modificar
# este programa sob os termos da licença GNU Affero General Public License,
# publicada pela Free Software Foundation, em sua versão 3 ou, de acordo
# com sua opção, qualquer versão posterior.
#
# Este programa é distribuido na esperança de que venha a ser útil,
# porém SEM QUAISQUER GARANTIAS, nem mesmo a garantia implícita de
# COMERCIABILIDADE ou ADEQUAÇÃO A UMA FINALIDADE ESPECÍFICA. Veja a
# GNU Affero General Public License para mais detalhes.
#
# Você deve ter recebido uma cópia da GNU Affero General Public License
# juntamente com este programa. Caso esse não seja o caso, acesse:
# <http://www.gnu.org/licenses/>
#

from __future__ import division, print_function, unicode_literals

from pysped.xml_sped import *
from pysped.nfe.leiaute import ESQUEMA_ATUAL_VERSAO_3 as ESQUEMA_ATUAL
from pysped.nfe.leiaute import consrecinfe_200
import os
from nfe_310 import NFe

DIRNAME = os.path.dirname(__file__)

class ConsReciNFe(consrecinfe_200.ConsReciNFe):
    def __init__(self):
        super(ConsReciNFe, self).__init__()
        self.versao  = TagDecimal(nome='consReciNFe', codigo='BP02', propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'consReciNFe_v3.10.xsd'

class InfProt(consrecinfe_200.InfProt):
    def __init__(self):
        super(InfProt, self).__init__()
        self.dhRecbto = TagDataHoraUTC(nome='dhRecbto', codigo='PR08', raiz='//infProt')

    def get_xml(self):
        xml = XMLNFe.get_xml(self)

        if self.Id.valor:
            xml += self.Id.xml
        else:
            xml += '<infProt>'

        xml += self.tpAmb.xml
        xml += self.verAplic.xml
        xml += self.chNFe.xml
        xml += self.dhRecbto.xml
        xml += self.nProt.xml
        xml += self.digVal.xml
        xml += self.cStat.xml
        xml += self.xMotivo.xml
        xml += '</infProt>'
        return xml

    def set_xml(self, arquivo):
        if self._le_xml(arquivo):
            self.Id.xml        = arquivo
            self.tpAmb.xml     = arquivo
            self.verAplic.xml  = arquivo
            self.chNFe.xml     = arquivo
            self.dhRecbto.xml  = arquivo
            self.nProt.xml     = arquivo
            self.digVal.xml    = arquivo
            self.cStat.xml     = arquivo
            self.xMotivo.xml   = arquivo

    xml = property(get_xml, set_xml)

class ProtNFe(consrecinfe_200.ProtNFe):
    def __init__(self):
        super(ProtNFe, self).__init__()
        self.versao  = TagDecimal(nome='protNFe', codigo='PR02' , propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.infProt = InfProt()

class RetConsReciNFe(consrecinfe_200.RetConsReciNFe):
    def __init__(self):
        super(RetConsReciNFe, self).__init__()
        self.versao   = TagDecimal(nome='retConsReciNFe', codigo='BR02' , propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'retConsReciNFe_v3.10.xsd'

    def get_xml(self):
        xml = XMLNFe.get_xml(self)
        xml += self.versao.xml
        xml += self.tpAmb.xml
        xml += self.verAplic.xml
        xml += self.nRec.xml
        xml += self.cStat.xml
        xml += self.xMotivo.xml
        xml += self.cUF.xml
        xml += self.cMsg.xml
        xml += self.xMsg.xml

        for pn in self.protNFe:
            xml += pn.xml

        xml += '</retConsReciNFe>'
        return xml

    def set_xml(self, arquivo):
        if self._le_xml(arquivo):
            self.versao.xml   = arquivo
            self.tpAmb.xml    = arquivo
            self.verAplic.xml = arquivo
            self.nRec.xml     = arquivo
            self.cStat.xml    = arquivo
            self.xMotivo.xml  = arquivo
            self.cUF.xml      = arquivo
            self.cMsg.xml     = arquivo
            self.xMsg.xml     = arquivo
            self.protNFe      = self.le_grupo('//retConsReciNFe/protNFe', ProtNFe)

            #
            # Monta o dicionário dos protocolos
            #
            for pn in self.protNFe:
                self.dic_protNFe[pn.infProt.chNFe.valor] = pn

    xml = property(get_xml, set_xml)

class ProcNFe(consrecinfe_200.ProcNFe):
    def __init__(self):
        super(ProcNFe, self).__init__()
        self.versao  = TagDecimal(nome='nfeProc', propriedade='versao', namespace=NAMESPACE_NFE, valor='3.10', raiz='/')
        self.NFe     = NFe()
        self.protNFe = ProtNFe()
        self.caminho_esquema = os.path.join(DIRNAME, 'schema/', ESQUEMA_ATUAL + '/')
        self.arquivo_esquema = 'procNFe_v3.10.xsd'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Script to test the plugin support.

https://github.com/pySTEPS/cookiecutter-pysteps-plugin
"""
import os
import subprocess
import sys
import tempfile
from cookiecutter.main import cookiecutter

PLUGIN_TEMPLATE_URL = "https://github.com/pysteps/cookiecutter-pysteps-plugin"

from contextlib import contextmanager
import pysteps

def _check_installed_plugin(import_func_name):
    # reload the pysteps module to detect the installed plugin
    pysteps.io.discover_importers()
    assert hasattr(pysteps.io.importers, import_func_name)
    assert import_func_name in pysteps.io.interface._importer_methods
    importer = getattr(pysteps.io.importers, import_func_name)
    importer("filename")

@contextmanager
def _create_and_install_plugin(project_name, importer_name):
    with tempfile.TemporaryDirectory() as tmpdirname:
        print(f"Installing plugin {project_name} providing the {importer_name} module")
        cookiecutter(
            PLUGIN_TEMPLATE_URL,
            no_input=True,
            overwrite_if_exists=True,
            extra_context={
                "project_name": project_name,
                "importer_name": importer_name,
            },
            output_dir=tmpdirname,
        )
        # Install the plugin
        subprocess.check_call(
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--force-reinstall",
                os.path.join(tmpdirname, project_name),
            ]
        )

        # The block below, together with the decorator used in this function are used
        # to create a context manager that uninstall the plugin packages after the
        # tests finish (even if they fail).
        # https://docs.pytest.org/en/stable/fixture.html?highlight=context#fixture-finalization-executing-teardown-code
        try:
            yield project_name
        finally:
            _uninstall_plugin(project_name)

def _uninstall_plugin(project_name):
    # Install the plugin
    subprocess.check_call(
        [sys.executable, "-m", "pip", "uninstall", "-y", project_name]
    )

def test_importers_plugins():
    with _create_and_install_plugin("test_importer_aaa", "importer_aaa"):
        # The default plugin template appends an _xxx to the importer function.
        _check_installed_plugin("importer_aaa_xxx")
        with _create_and_install_plugin("test_importer_bbb", "importer_bbb"):
            _check_installed_plugin("importer_aaa_xxx")
            _check_installed_plugin("importer_bbb_xxx")
<commit_msg>Skip test if cookiecutter is not installed<commit_after># -*- coding: utf-8 -*-
"""
Script to test the plugin support.

https://github.com/pySTEPS/cookiecutter-pysteps-plugin
"""
import os
import pytest
import subprocess
import sys
import tempfile

__ = pytest.importorskip("cookiecutter")
from cookiecutter.main import cookiecutter

PLUGIN_TEMPLATE_URL = "https://github.com/pysteps/cookiecutter-pysteps-plugin"

from contextlib import contextmanager
import pysteps

def _check_installed_plugin(import_func_name):
    # reload the pysteps module to detect the installed plugin
    pysteps.io.discover_importers()
    assert hasattr(pysteps.io.importers, import_func_name)
    assert import_func_name in pysteps.io.interface._importer_methods
    importer = getattr(pysteps.io.importers, import_func_name)
    importer("filename")

@contextmanager
def _create_and_install_plugin(project_name, importer_name):
    with tempfile.TemporaryDirectory() as tmpdirname:
        print(f"Installing plugin {project_name} providing the {importer_name} module")
        cookiecutter(
            PLUGIN_TEMPLATE_URL,
            no_input=True,
            overwrite_if_exists=True,
            extra_context={
                "project_name": project_name,
                "importer_name": importer_name,
            },
            output_dir=tmpdirname,
        )
        # Install the plugin
        subprocess.check_call(
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--force-reinstall",
                os.path.join(tmpdirname, project_name),
            ]
        )

        # The block below, together with the decorator used in this function are used
        # to create a context manager that uninstall the plugin packages after the
        # tests finish (even if they fail).
        # https://docs.pytest.org/en/stable/fixture.html?highlight=context#fixture-finalization-executing-teardown-code
        try:
            yield project_name
        finally:
            _uninstall_plugin(project_name)

def _uninstall_plugin(project_name):
    # Install the plugin
    subprocess.check_call(
        [sys.executable, "-m", "pip", "uninstall", "-y", project_name]
    )

def test_importers_plugins():
    with _create_and_install_plugin("test_importer_aaa", "importer_aaa"):
        # The default plugin template appends an _xxx to the importer function.
        _check_installed_plugin("importer_aaa_xxx")
        with _create_and_install_plugin("test_importer_bbb", "importer_bbb"):
            _check_installed_plugin("importer_aaa_xxx")
            _check_installed_plugin("importer_bbb_xxx")
<|endoftext|>


<commit_before>##########################################################################
#
#  Copyright (c) 2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of Image Engine Design nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import unittest

import IECore
import Gaffer
import GafferImage
import os

class SamplerTest( unittest.TestCase ) :

	fileName = os.path.expandvars( "$GAFFER_ROOT/python/GafferTest/images/checker.exr" )

	def testOutOfBoundsSampleModeBlack( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		c = Gaffer.Context()
		c["image:channelName"] = 'R'
		c["image:tileOrigin"] = IECore.V2i( 0 )

		bounds = r["out"]["dataWindow"].getValue();
		f = GafferImage.Filter.create( "Box" )

		testCases = [
			( bounds.min.x-1, bounds.min.y ),
			( bounds.min.x, bounds.min.y-1 ),
			( bounds.max.x, bounds.max.y+1 ),
			( bounds.max.x+1, bounds.max.y ),
			( bounds.min.x-1, bounds.max.y ),
			( bounds.min.x, bounds.max.y+1 ),
			( bounds.max.x+1, bounds.min.y ),
			( bounds.max.x, bounds.min.y-1 )
		]

		self.assertTrue( "Box" in GafferImage.Filter.filters() )
		f = GafferImage.Filter.create("Box")

		with c :

			self.assertTrue( "R" in r["out"]["channelNames"].getValue() )
			s = GafferImage.Sampler( r["out"], "R", bounds, GafferImage.BoundingMode.Black )

			# Check that the bounding pixels are non zero.
			self.assertNotEqual( s.sample( bounds.min.x+.5, bounds.min.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.max.x+.5, bounds.max.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.min.x+.5, bounds.max.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.max.x+.5, bounds.min.y+.5, f ), 0. )

			# Sample out of bounds and assert that a zero is returned.
			for x, y in testCases :
				self.assertEqual( s.sample( x+.5, y+.5, f ), 0. )

	def testOutOfBoundsSampleModeClamp( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		c = Gaffer.Context()
		c["image:channelName"] = 'R'
		c["image:tileOrigin"] = IECore.V2i( 0 )

		bounds = r["out"]["dataWindow"].getValue();
		f = GafferImage.Filter.create( "Box" )

		with c :

			self.assertTrue( "R" in r["out"]["channelNames"].getValue() )
			s = GafferImage.Sampler( r["out"], "R", bounds, GafferImage.BoundingMode.Clamp )

			# Get the values of the corner pixels.
			bl = s.sample( bounds.min.x+.5, bounds.min.y+.5, f )
			br = s.sample( bounds.max.x+.5, bounds.min.y+.5, f )
			tr = s.sample( bounds.max.x+.5, bounds.max.y+.5, f )
			tl = s.sample( bounds.min.x+.5, bounds.max.y+.5, f )

			# Sample out of bounds and assert that the same value as the nearest pixel is returned.
			self.assertEqual( s.sample( bounds.min.x-1, bounds.min.y, f ), bl )
			self.assertEqual( s.sample( bounds.min.x, bounds.min.y-1, f ), bl )
			self.assertEqual( s.sample( bounds.max.x, bounds.max.y+1, f ), tr )
			self.assertEqual( s.sample( bounds.max.x+1, bounds.max.y, f ), tr )
			self.assertEqual( s.sample( bounds.min.x-1, bounds.max.y, f ), tl )
			self.assertEqual( s.sample( bounds.min.x, bounds.max.y+1, f ), tl )
			self.assertEqual( s.sample( bounds.max.x+1, bounds.min.y, f ), br )
			self.assertEqual( s.sample( bounds.max.x, bounds.min.y-1, f ), br )

	# Test that the hash() method accumulates all of the hashes of the tiles within the sample area
	# for a large number of different sample areas.
	def testSampleHash( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		tileSize = GafferImage.ImagePlug.tileSize()
		# Test one tile first.

		for w in range(2, 6) :
			for x in range(-3, 3) :
				for y in range(-3, 3) :
					sampleBox = IECore.Box2i( IECore.V2i( x*tileSize, y*tileSize ), IECore.V2i( x*tileSize+w*tileSize/2, y*tileSize+w*tileSize/2 ) )
					self.__testHashOfBounds( sampleBox, "R", r["out"] )

	# A private method that acumulates the hashes of the tiles within
	# a box and compares them to the hash returned by the sampler.
	def __testHashOfBounds( self, box, channel, plug ) :

		tileOrigin = GafferImage.ImagePlug.tileOrigin( IECore.V2i( box.min ) )
		self.assertTrue( channel in plug["channelNames"].getValue() )

		c = Gaffer.Context()
		c["image:channelName"] = channel
		c["image:tileOrigin"] = tileOrigin

		h = IECore.MurmurHash()
		h2 = h

		# Get the hash from the sampler.
		with c :
			s = GafferImage.Sampler( plug, channel, box, GafferImage.BoundingMode.Clamp )
			s.hash( h )

		# Get the hash from the tiles within our desired sample area.
		with c :
			y = box.min.y
			while y <= box.max.y :
				x = box.min.x
				while x <= box.max.x :
					tileOrigin = GafferImage.ImagePlug.tileOrigin( IECore.V2i( x, y ) )
					tile = IECore.Box2i( tileOrigin, tileOrigin + IECore.V2i( GafferImage.ImagePlug.tileSize()-1 ) )
					h2.append( plug.channelDataHash( channel, tileOrigin ) )
					x += GafferImage.ImagePlug.tileSize()
				y += GafferImage.ImagePlug.tileSize()

		self.assertEqual( h, h2 )

<commit_msg>Rebasing.<commit_after>##########################################################################
#
#  Copyright (c) 2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of Image Engine Design nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import unittest

import IECore
import Gaffer
import GafferImage
import os

class SamplerTest( unittest.TestCase ) :

	fileName = os.path.expandvars( "$GAFFER_ROOT/python/GafferTest/images/checker.exr" )

	def testOutOfBoundsSampleModeBlack( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		c = Gaffer.Context()
		c["image:channelName"] = 'R'
		c["image:tileOrigin"] = IECore.V2i( 0 )

		bounds = r["out"]["dataWindow"].getValue();

		testCases = [
			( bounds.min.x-1, bounds.min.y ),
			( bounds.min.x, bounds.min.y-1 ),
			( bounds.max.x, bounds.max.y+1 ),
			( bounds.max.x+1, bounds.max.y ),
			( bounds.min.x-1, bounds.max.y ),
			( bounds.min.x, bounds.max.y+1 ),
			( bounds.max.x+1, bounds.min.y ),
			( bounds.max.x, bounds.min.y-1 )
		]

		self.assertTrue( "Box" in GafferImage.Filter.filters() )
		f = GafferImage.Filter.create("Box")

		with c :

			self.assertTrue( "R" in r["out"]["channelNames"].getValue() )
			s = GafferImage.Sampler( r["out"], "R", bounds, GafferImage.BoundingMode.Black )

			# Check that the bounding pixels are non zero.
			self.assertNotEqual( s.sample( bounds.min.x+.5, bounds.min.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.max.x+.5, bounds.max.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.min.x+.5, bounds.max.y+.5, f ), 0. )
			self.assertNotEqual( s.sample( bounds.max.x+.5, bounds.min.y+.5, f ), 0. )

			# Sample out of bounds and assert that a zero is returned.
			for x, y in testCases :
				self.assertEqual( s.sample( x+.5, y+.5, f ), 0. )

	def testOutOfBoundsSampleModeClamp( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		c = Gaffer.Context()
		c["image:channelName"] = 'R'
		c["image:tileOrigin"] = IECore.V2i( 0 )

		bounds = r["out"]["dataWindow"].getValue();
		f = GafferImage.Filter.create( "Box" )

		with c :

			self.assertTrue( "R" in r["out"]["channelNames"].getValue() )
			s = GafferImage.Sampler( r["out"], "R", bounds, GafferImage.BoundingMode.Clamp )

			# Get the values of the corner pixels.
			bl = s.sample( bounds.min.x+.5, bounds.min.y+.5, f )
			br = s.sample( bounds.max.x+.5, bounds.min.y+.5, f )
			tr = s.sample( bounds.max.x+.5, bounds.max.y+.5, f )
			tl = s.sample( bounds.min.x+.5, bounds.max.y+.5, f )

			# Sample out of bounds and assert that the same value as the nearest pixel is returned.
			self.assertEqual( s.sample( bounds.min.x-1, bounds.min.y, f ), bl )
			self.assertEqual( s.sample( bounds.min.x, bounds.min.y-1, f ), bl )
			self.assertEqual( s.sample( bounds.max.x, bounds.max.y+1, f ), tr )
			self.assertEqual( s.sample( bounds.max.x+1, bounds.max.y, f ), tr )
			self.assertEqual( s.sample( bounds.min.x-1, bounds.max.y, f ), tl )
			self.assertEqual( s.sample( bounds.min.x, bounds.max.y+1, f ), tl )
			self.assertEqual( s.sample( bounds.max.x+1, bounds.min.y, f ), br )
			self.assertEqual( s.sample( bounds.max.x, bounds.min.y-1, f ), br )

	# Test that the hash() method accumulates all of the hashes of the tiles within the sample area
	# for a large number of different sample areas.
	def testSampleHash( self ) :

		s = Gaffer.ScriptNode()
		r = GafferImage.ImageReader()
		r["fileName"].setValue( self.fileName )
		s.addChild( r )

		tileSize = GafferImage.ImagePlug.tileSize()
		# Test one tile first.

		for w in range(2, 6) :
			for x in range(-3, 3) :
				for y in range(-3, 3) :
					sampleBox = IECore.Box2i( IECore.V2i( x*tileSize, y*tileSize ), IECore.V2i( x*tileSize+w*tileSize/2, y*tileSize+w*tileSize/2 ) )
					self.__testHashOfBounds( sampleBox, "R", r["out"] )

	# A private method that acumulates the hashes of the tiles within
	# a box and compares them to the hash returned by the sampler.
	def __testHashOfBounds( self, box, channel, plug ) :

		tileOrigin = GafferImage.ImagePlug.tileOrigin( IECore.V2i( box.min ) )
		self.assertTrue( channel in plug["channelNames"].getValue() )

		c = Gaffer.Context()
		c["image:channelName"] = channel
		c["image:tileOrigin"] = tileOrigin

		h = IECore.MurmurHash()
		h2 = h

		# Get the hash from the sampler.
		with c :
			s = GafferImage.Sampler( plug, channel, box, GafferImage.BoundingMode.Clamp )
			s.hash( h )

		# Get the hash from the tiles within our desired sample area.
		with c :
			y = box.min.y
			while y <= box.max.y :
				x = box.min.x
				while x <= box.max.x :
					tileOrigin = GafferImage.ImagePlug.tileOrigin( IECore.V2i( x, y ) )
					tile = IECore.Box2i( tileOrigin, tileOrigin + IECore.V2i( GafferImage.ImagePlug.tileSize()-1 ) )
					h2.append( plug.channelDataHash( channel, tileOrigin ) )
					x += GafferImage.ImagePlug.tileSize()
				y += GafferImage.ImagePlug.tileSize()

		self.assertEqual( h, h2 )

<|endoftext|>


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

"""
Half sample images from a rosbag and saved to another rosbag.
The data from other topics are copied directly to the output rosbag.
The image topics are preset to be /cam0/image_raw and /cam1/image_raw.
"""

import os
import argparse

import cv2

import rosbag
from cv_bridge import CvBridge
import play_images_in_rosbag

def decide_output_encoding(cv_img):
    coding = 'passthrough'
    if cv_img.dtype == 'uint16':
        coding = 'mono16'
    return coding

def main():
    """Extract a folder of images from a rosbag.
    """
    parser = argparse.ArgumentParser(description="Downsample images in a ROS bag.")
    parser.add_argument("bag_file", help="Input ROS bag.")
    parser.add_argument("--out_bag_file", help="Output ROS bag file.", default=None)

    args = parser.parse_args()
    out_bag_file = args.out_bag_file
    if args.out_bag_file is None:
        out_bag_file = os.path.join(os.path.splitext(args.bag_file)[0] + '_512_16.bag')

    in_bag = rosbag.Bag(args.bag_file, "r")
    topic_list = in_bag.get_type_and_topic_info()[1].keys()
    print('Topics {} are in {}'.format(topic_list, args.bag_file))

    out_bag = rosbag.Bag(out_bag_file, 'w')
    bridge = CvBridge()
    for k in range(2):
        count = 0
        image_topic = '/cam{}/image_raw'.format(k)
        encoding = ''
        for _, msg, t in in_bag.read_messages(topics=[image_topic]):
            cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
            h, w = cv_img.shape[:2]
            cv_half_img = cv2.pyrDown(cv_img, dstsize=(w / 2, h / 2))
            if count == 0:
                print('Image info before and after half sampling:')
                play_images_in_rosbag.print_image_info(cv_img)
                play_images_in_rosbag.print_image_info(cv_half_img)
                encoding = decide_output_encoding(cv_img)

            cv2.imshow('Downsampled frame', cv_half_img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            count += 1

            # For 16UC1 input image type, we may need to use mono16 as output encoding option.
            # see http://library.isr.ist.utl.pt/docs/roswiki/cv_bridge(2f)Tutorials(2f)UsingCvBridgeToConvertBetweenROSImagesAndOpenCVImages.html
            rosimage = bridge.cv2_to_imgmsg(cv_half_img, encoding=encoding)
            rosimage.header.stamp = t  # This is needed because header.stamp could have been 0.0.
            out_bag.write(image_topic, rosimage, t)
        print('Saved {} images on topic {}'.format(count, image_topic))

    # copy the other messages
    for topic in topic_list:
        if 'image_raw' in topic:
            continue
        count = 0
        for returned_topic, msg, t in in_bag.read_messages(topics=[topic]):
            msg.header.stamp = t  # This is needed because header.stamp could have been very different from t.
            out_bag.write(topic, msg, t)
            count += 1
        print('Saved {} messages on topic {}'.format(count, topic))

    cv2.destroyAllWindows()
    out_bag.close()
    in_bag.close()

    print("Downsampled images of {} on topics /cam0/image_raw /cam1/image_raw into {}".
          format(args.bag_file, out_bag_file))

if __name__ == '__main__':
    main()
<commit_msg>use header.stamp instead of message time for TUM VI<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Half sample images from a rosbag and saved to another rosbag.
The data from other topics are copied directly to the output rosbag.
The image topics are preset to be /cam0/image_raw and /cam1/image_raw.
"""
from __future__ import print_function
import os
import argparse

import rosbag
from cv_bridge import CvBridge
import cv2

import play_images_in_rosbag

def decide_output_encoding(cv_img):
    """
    For 16UC1 input image, we need to use mono16 as output encoding option.
    see http://library.isr.ist.utl.pt/docs/roswiki/cv_bridge(2f)Tutorials(2f)
    UsingCvBridgeToConvertBetweenROSImagesAndOpenCVImages.html
    :param cv_img:
    :return:
    """

    coding = 'passthrough'
    if cv_img.dtype == 'uint16':
        coding = 'mono16'
    return coding

def main():
    """Extract a folder of images from a rosbag.
    """
    parser = argparse.ArgumentParser(
        description="Downsample images in a ROS bag.")
    parser.add_argument("bag_file", help="Input ROS bag.")
    parser.add_argument("--out_bag_file",
                        help="Output ROS bag file.",
                        default=None)

    args = parser.parse_args()
    out_bag_file = args.out_bag_file
    if args.out_bag_file is None:
        out_bag_file = os.path.join(
            os.path.splitext(args.bag_file)[0] + '_512_16.bag')

    in_bag = rosbag.Bag(args.bag_file, "r")
    topic_list = in_bag.get_type_and_topic_info()[1].keys()
    print('Topics {} are in {}'.format(topic_list, args.bag_file))

    out_bag = rosbag.Bag(out_bag_file, 'w')
    bridge = CvBridge()
    time_shift = None
    for k in range(2):
        count = 0
        image_topic = '/cam{}/image_raw'.format(k)
        encoding = ''
        for _, msg, t in in_bag.read_messages(topics=[image_topic]):
            cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
            h, w = cv_img.shape[:2]
            cv_half_img = cv2.pyrDown(cv_img, dstsize=(w / 2, h / 2))
            if count == 0:
                print('Image info before and after half sampling:')
                play_images_in_rosbag.print_image_info(cv_img)
                play_images_in_rosbag.print_image_info(cv_half_img)
                encoding = decide_output_encoding(cv_img)

            cv2.imshow('Downsampled frame', cv_half_img)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            count += 1

            rosimage = bridge.cv2_to_imgmsg(cv_half_img, encoding=encoding)
            if time_shift is None:
                time_shift = t - msg.header.stamp
            # assign because rosimage.header.stamp is zero.
            rosimage.header.stamp = time_shift + msg.header.stamp
            out_bag.write(image_topic, rosimage, rosimage.header.stamp)
        print('Saved {} images on topic {}'.format(count, image_topic))

    # copy the other messages
    for topic in topic_list:
        if 'image_raw' in topic:
            continue
        count = 0
        for _, msg, t in in_bag.read_messages(topics=[topic]):
            msg.header.stamp = time_shift + msg.header.stamp
            out_bag.write(topic, msg, msg.header.stamp)
            count += 1
        print('Saved {} messages on topic {}'.format(count, topic))

    cv2.destroyAllWindows()
    out_bag.close()
    in_bag.close()

    print("Downsampled images of {} on topics /cam0/image_raw "
          "/cam1/image_raw into {}".format(args.bag_file, out_bag_file))

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


<commit_before>from cybox.core import Object, Observable, ObservableComposition
from cybox.objects.file_object import File
from cybox.objects.address_object import Address
from cybox.objects.hostname_object import Hostname
from cybox.objects.uri_object import URI
from cybox.objects.pipe_object import Pipe
from cybox.objects.mutex_object import Mutex
from cybox.objects.artifact_object import Artifact
from cybox.objects.memory_object import Memory
from cybox.objects.email_message_object import EmailMessage, EmailHeader, Attachments
from cybox.objects.domain_name_object import DomainName
from cybox.objects.win_registry_key_object import *
from cybox.common import Hash, ByteRun, ByteRuns
from cybox.objects.http_session_object import *
from cybox.objects.as_object import AutonomousSystem
from stix.extensions.test_mechanism.snort_test_mechanism import *
import ntpath, socket, sys
from stix.indicator import Indicator

this_module = sys.modules[__name__]

simple_type_to_method = {}
simple_type_to_method.update(dict.fromkeys(["md5", "sha1", "sha256", "filename", "filename|md5", "filename|sha1", "filename|sha256", "malware-sample", "attachment"], "resolveFileObservable"))
simple_type_to_method.update(dict.fromkeys(["ip-src", "ip-dst"], "generateIPObservable"))
simple_type_to_method.update(dict.fromkeys(["regkey", "regkey|value"], "generateRegkeyObservable"))
simple_type_to_method.update(dict.fromkeys(["hostname", "domain", "url", "AS", "mutex", "named pipe", "link"], "generateSimpleObservable"))
simple_type_to_method.update(dict.fromkeys(["email-src", "email-dst", "email-subject"], "resolveEmailObservable"))
simple_type_to_method.update(dict.fromkeys(["http-method", "user-agent"], "resolveHTTPObservable"))
simple_type_to_method.update(dict.fromkeys(["pattern-in-file", "pattern-in-traffic", "pattern-in-memory"], "resolvePatternObservable"))

# mapping for the attributes that can go through the simpleobservable script
misp_cybox_name = {"domain" : "DomainName", "hostname" : "Hostname", "url" : "URI", "AS" : "AutonomousSystem", "mutex" : "Mutex", "named pipe" : "Pipe", "link" : "URI"}
cybox_name_attribute = {"DomainName" : "value", "Hostname" : "hostname_value", "URI" : "value", "AutonomousSystem" : "number", "Pipe" : "name", "Mutex" : "name"}
misp_indicator_type = {"domain" : "Domain Watchlist", "hostname" : "Domain Watchlist", "url" : "URL Watchlist", "AS" : "", "mutex" : "Host Characteristics", "named pipe" : "Host Characteristics", "link" : ""}

# mapping Windows Registry Hives and their abbreviations
# see https://cybox.mitre.org/language/version2.1/xsddocs/objects/Win_Registry_Key_Object_xsd.html#RegistryHiveEnum
# the dict keys must be UPPER CASE and end with \\
misp_reghive = {
    "HKEY_CLASSES_ROOT\\"                : "HKEY_CLASSES_ROOT",
    "HKCR\\"                             : "HKEY_CLASSES_ROOT",
    "HKEY_CURRENT_CONFIG\\"              : "HKEY_CURRENT_CONFIG",
    "HKCC\\"                             : "HKEY_CURRENT_CONFIG",
    "HKEY_CURRENT_USER\\"                : "HKEY_CURRENT_USER",
    "HKCU\\"                             : "HKEY_CURRENT_USER",
    "HKEY_LOCAL_MACHINE\\"               : "HKEY_LOCAL_MACHINE",
    "HKLM\\"                             : "HKEY_LOCAL_MACHINE",
    "HKEY_USERS\\"                       : "HKEY_USERS",
    "HKU\\"                              : "HKEY_USERS",
    "HKEY_CURRENT_USER_LOCAL_SETTINGS\\" : "HKEY_CURRENT_USER_LOCAL_SETTINGS",
    "HKCULS\\"                           : "HKEY_CURRENT_USER_LOCAL_SETTINGS",
    "HKEY_PERFORMANCE_DATA\\"            : "HKEY_PERFORMANCE_DATA",
    "HKPD\\"                             : "HKEY_PERFORMANCE_DATA",
    "HKEY_PERFORMANCE_NLSTEXT\\"         : "HKEY_PERFORMANCE_NLSTEXT",
    "HKPN\\"                             : "HKEY_PERFORMANCE_NLSTEXT",
    "HKEY_PERFORMANCE_TEXT\\"            : "HKEY_PERFORMANCE_TEXT",
    "HKPT\\"                             : "HKEY_PERFORMANCE_TEXT",
}

def generateObservable(indicator, attribute):
    if (attribute["type"] in ("snort", "yara")):
        generateTM(indicator, attribute)
    else:
        observable = None;
        if (attribute["type"] in simple_type_to_method.keys()):
            action = getattr(this_module, simple_type_to_method[attribute["type"]], None)
            if (action != None):
                property = action(indicator, attribute)
		property.condition = "Equals"
                object = Object(property)
                object.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":" + property.__class__.__name__ + "-" + attribute["uuid"]
                observable = Observable(object)
                observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
                indicator.add_observable(observable)

def resolveFileObservable(indicator, attribute):
    hashValue = ""
    filenameValue = ""
    if (attribute["type"] in ("filename|md5", "filename|sha1", "filename|sha256", "malware-sample")):
        values = attribute["value"].split('|')
        filenameValue = values[0]
        hashValue = values[1]
        indicator.add_indicator_type("File Hash Watchlist")
    else:
        if (attribute["type"] in ("filename", "attachment")):
            filenameValue = attribute["value"]
        else:
            hashValue = attribute["value"]
            indicator.add_indicator_type("File Hash Watchlist")
    observable = generateFileObservable(filenameValue, hashValue)
    return observable

def generateFileObservable(filenameValue, hashValue):
    file_object = File()
    if (filenameValue != ""):
        if (("/" in filenameValue) or ("\\" in filenameValue)):
            file_object.file_path = ntpath.dirname(filenameValue)
            file_object.file_name = ntpath.basename(filenameValue)
        else:
            file_object.file_name = filenameValue
    if (hashValue != ""):
        file_object.add_hash(Hash(hash_value=hashValue, exact=True))
    return file_object

def generateIPObservable(indicator, attribute):
    indicator.add_indicator_type("IP Watchlist")
    address_object = Address()
    cidr = False
    if ("/" in attribute["value"]):
        ip = attribute["value"].split('/')[0]
        cidr = True
    else:
        ip = attribute["value"]
    try:
        socket.inet_aton(ip)
        ipv4 = True
    except socket.error:
        ipv4 = False
    if (cidr == True):
        address_object.category = "cidr"
    elif (ipv4 == True):
        address_object.category = "ipv4-addr"
    else:
        address_object.category = "ipv6-addr"
    if (attribute["type"] == "ip-src"):
        address_object.is_source = True
    else:
        address_object.is_source = False
    address_object.address_value = attribute["value"]
    return address_object

def generateRegkeyObservable(indicator, attribute):
    indicator.add_indicator_type("Host Characteristics")
    regkey = ""
    regvalue = ""
    if (attribute["type"] == "regkey|value"):
        regkey = attribute["value"].split('|')[0]
        regvalue = attribute["value"].split('|')[1]
    else:
        regkey = attribute["value"]
    reghive, regkey = resolveRegHive(regkey)
    reg_object = WinRegistryKey()
    reg_object.key = regkey
    reg_object.key.condition = "Equals"
    if (reghive != None):
        reg_object.hive = reghive
        reg_object.hive.condition = "Equals"
    if (regvalue != ""):
        reg_value_object = RegistryValue()
        reg_value_object.data = regvalue
        reg_value_object.data.condition = "Equals"
        reg_object.values = RegistryValues(reg_value_object)
    return reg_object

def generateSimpleObservable(indicator, attribute):
    cyboxName = misp_cybox_name[attribute["type"]]
    constructor = getattr(this_module, cyboxName, None)
    indicatorType = misp_indicator_type[attribute["type"]]
    if (indicatorType != ""):
        indicator.add_indicator_type(indicatorType)
    new_object = constructor()
    setattr(new_object, cybox_name_attribute[cyboxName], attribute["value"])
    return new_object

def generateTM(indicator, attribute):
    if (attribute["type"] == "snort"):
        tm = SnortTestMechanism()
        tm.rules = [attribute["value"]]
    else:
        # remove the following line and uncomment the code below once yara test mechanisms get added to python-stix
        return indicator
        #tm = SnortTestMechanism()
        #tm.rules = [attribute["value"]]
    indicator.test_mechanisms = [tm]

def resolveEmailObservable(indicator, attribute):
    indicator.add_indicator_type("Malicious E-mail")
    new_object = EmailMessage()
    email_header = EmailHeader()
    if (attribute["type"] == "email-src"):
        email_header.from_ = attribute["value"]
    elif(attribute["type"] == "email-dst"):
        email_header.to = attribute["value"]
    else:
        email_header.subject = attribute["value"]
    new_object.header = email_header
    return new_object

def resolveHTTPObservable(indicator, attribute):
    request_response = HTTPRequestResponse()
    client_request = HTTPClientRequest()
    if (attribute["type"] == "user-agent"):
        header = HTTPRequestHeader()
        header_fields = HTTPRequestHeaderFields()
        header_fields.user_agent = attribute["value"]
        header.parsed_header = header_fields
        client_request.http_request_header = header
    else:
        line = HTTPRequestLine()
        line.http_method = attribute["value"]
        client_request.http_request_line = line
    request_response.http_client_request = client_request
    new_object = HTTPSession()
    request_response.to_xml()
    new_object.http_request_response = [request_response]
    return new_object

# use this when implementing pattern in memory and pattern in traffic
def resolvePatternObservable(indicator, attribute):
    new_object = None
    if attribute["type"] == "pattern-in-file":
        byte_run = ByteRun()
        byte_run.byte_run_data = attribute["value"]
        new_object = File()
        new_object.byte_runs = ByteRuns(byte_run)
    # custom properties are not implemented in the API yet
    # elif attribute["type"] == "pattern-in-memory":
    # elif attribute["type"] == "pattern-in-traffic":
    return new_object

# create an artifact object for the malware-sample type.
def createArtifactObject(indicator, attribute):
    artifact = Artifact(data = attribute["data"])
    artifact.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":artifact-" + attribute["uuid"]
    observable = Observable(artifact)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-artifact-" + attribute["uuid"]
    indicator.add_observable(observable)

# return either a composition if data is set in attribute, or just an observable with a filename if it's not set
def returnAttachmentComposition(attribute):
    file_object = File()
    file_object.file_name = attribute["value"]
    file_object.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":file-" + attribute["uuid"]
    observable = Observable()
    if "data" in attribute:
        artifact = Artifact(data = attribute["data"])
        artifact.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":artifact-" + attribute["uuid"]
        observable_artifact = Observable(artifact)
        observable_artifact.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-artifact-" + attribute["uuid"]
        observable_file = Observable(file_object)
        observable_file.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-file-" + attribute["uuid"]
        composition = ObservableComposition(observables = [observable_artifact, observable_file])
        observable.observable_composition = composition
    else:
        observable = Observable(file_object)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
    if attribute["comment"] != "":
        observable.description = attribute["comment"]
    return observable

# email-attachment are mapped to an email message observable that contains the attachment as a file object
def generateEmailAttachmentObject(indicator, attribute):
    file_object = File()
    file_object.file_name = attribute["value"]
    email = EmailMessage()
    email.attachments = Attachments()
    email.add_related(file_object, "Contains", inline=True)
    file_object.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":file-" + attribute["uuid"]
    email.attachments.append(file_object.parent.id_)
    email.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":EmailMessage-" + attribute["uuid"]
    observable = Observable(email)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
    indicator.observable = observable

# split registry string into hive and key
def resolveRegHive(regStr):
    regStr = regStr.lstrip('\\')
    regStrU = regStr.upper()
    for hive in misp_reghive.iterkeys():
        if regStrU.startswith(hive):
            return misp_reghive[hive], regStr[len(hive):]
    return None, regStr
<commit_msg>Set conditions for simple observables<commit_after>from cybox.core import Object, Observable, ObservableComposition
from cybox.objects.file_object import File
from cybox.objects.address_object import Address
from cybox.objects.hostname_object import Hostname
from cybox.objects.uri_object import URI
from cybox.objects.pipe_object import Pipe
from cybox.objects.mutex_object import Mutex
from cybox.objects.artifact_object import Artifact
from cybox.objects.memory_object import Memory
from cybox.objects.email_message_object import EmailMessage, EmailHeader, Attachments
from cybox.objects.domain_name_object import DomainName
from cybox.objects.win_registry_key_object import *
from cybox.common import Hash, ByteRun, ByteRuns
from cybox.objects.http_session_object import *
from cybox.objects.as_object import AutonomousSystem
from stix.extensions.test_mechanism.snort_test_mechanism import *
import ntpath, socket, sys
from stix.indicator import Indicator

this_module = sys.modules[__name__]

simple_type_to_method = {}
simple_type_to_method.update(dict.fromkeys(["md5", "sha1", "sha256", "filename", "filename|md5", "filename|sha1", "filename|sha256", "malware-sample", "attachment"], "resolveFileObservable"))
simple_type_to_method.update(dict.fromkeys(["ip-src", "ip-dst"], "generateIPObservable"))
simple_type_to_method.update(dict.fromkeys(["regkey", "regkey|value"], "generateRegkeyObservable"))
simple_type_to_method.update(dict.fromkeys(["hostname", "domain", "url", "AS", "mutex", "named pipe", "link"], "generateSimpleObservable"))
simple_type_to_method.update(dict.fromkeys(["email-src", "email-dst", "email-subject"], "resolveEmailObservable"))
simple_type_to_method.update(dict.fromkeys(["http-method", "user-agent"], "resolveHTTPObservable"))
simple_type_to_method.update(dict.fromkeys(["pattern-in-file", "pattern-in-traffic", "pattern-in-memory"], "resolvePatternObservable"))

# mapping for the attributes that can go through the simpleobservable script
misp_cybox_name = {"domain" : "DomainName", "hostname" : "Hostname", "url" : "URI", "AS" : "AutonomousSystem", "mutex" : "Mutex", "named pipe" : "Pipe", "link" : "URI"}
cybox_name_attribute = {"DomainName" : "value", "Hostname" : "hostname_value", "URI" : "value", "AutonomousSystem" : "number", "Pipe" : "name", "Mutex" : "name"}
misp_indicator_type = {"domain" : "Domain Watchlist", "hostname" : "Domain Watchlist", "url" : "URL Watchlist", "AS" : "", "mutex" : "Host Characteristics", "named pipe" : "Host Characteristics", "link" : ""}

# mapping Windows Registry Hives and their abbreviations
# see https://cybox.mitre.org/language/version2.1/xsddocs/objects/Win_Registry_Key_Object_xsd.html#RegistryHiveEnum
# the dict keys must be UPPER CASE and end with \\
misp_reghive = {
    "HKEY_CLASSES_ROOT\\"                : "HKEY_CLASSES_ROOT",
    "HKCR\\"                             : "HKEY_CLASSES_ROOT",
    "HKEY_CURRENT_CONFIG\\"              : "HKEY_CURRENT_CONFIG",
    "HKCC\\"                             : "HKEY_CURRENT_CONFIG",
    "HKEY_CURRENT_USER\\"                : "HKEY_CURRENT_USER",
    "HKCU\\"                             : "HKEY_CURRENT_USER",
    "HKEY_LOCAL_MACHINE\\"               : "HKEY_LOCAL_MACHINE",
    "HKLM\\"                             : "HKEY_LOCAL_MACHINE",
    "HKEY_USERS\\"                       : "HKEY_USERS",
    "HKU\\"                              : "HKEY_USERS",
    "HKEY_CURRENT_USER_LOCAL_SETTINGS\\" : "HKEY_CURRENT_USER_LOCAL_SETTINGS",
    "HKCULS\\"                           : "HKEY_CURRENT_USER_LOCAL_SETTINGS",
    "HKEY_PERFORMANCE_DATA\\"            : "HKEY_PERFORMANCE_DATA",
    "HKPD\\"                             : "HKEY_PERFORMANCE_DATA",
    "HKEY_PERFORMANCE_NLSTEXT\\"         : "HKEY_PERFORMANCE_NLSTEXT",
    "HKPN\\"                             : "HKEY_PERFORMANCE_NLSTEXT",
    "HKEY_PERFORMANCE_TEXT\\"            : "HKEY_PERFORMANCE_TEXT",
    "HKPT\\"                             : "HKEY_PERFORMANCE_TEXT",
}

def generateObservable(indicator, attribute):
    if (attribute["type"] in ("snort", "yara")):
        generateTM(indicator, attribute)
    else:
        observable = None;
        if (attribute["type"] in simple_type_to_method.keys()):
            action = getattr(this_module, simple_type_to_method[attribute["type"]], None)
            if (action != None):
                property = action(indicator, attribute)
		property.condition = "Equals"
                object = Object(property)
                object.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":" + property.__class__.__name__ + "-" + attribute["uuid"]
                observable = Observable(object)
                observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
                indicator.add_observable(observable)

def resolveFileObservable(indicator, attribute):
    hashValue = ""
    filenameValue = ""
    if (attribute["type"] in ("filename|md5", "filename|sha1", "filename|sha256", "malware-sample")):
        values = attribute["value"].split('|')
        filenameValue = values[0]
        hashValue = values[1]
        indicator.add_indicator_type("File Hash Watchlist")
    else:
        if (attribute["type"] in ("filename", "attachment")):
            filenameValue = attribute["value"]
        else:
            hashValue = attribute["value"]
            indicator.add_indicator_type("File Hash Watchlist")
    observable = generateFileObservable(filenameValue, hashValue)
    return observable

def generateFileObservable(filenameValue, hashValue):
    file_object = File()
    if (filenameValue != ""):
        if (("/" in filenameValue) or ("\\" in filenameValue)):
            file_object.file_path = ntpath.dirname(filenameValue)
            file_object.file_name = ntpath.basename(filenameValue)
        else:
            file_object.file_name = filenameValue
    if (hashValue != ""):
        file_object.add_hash(Hash(hash_value=hashValue, exact=True))
    return file_object

def generateIPObservable(indicator, attribute):
    indicator.add_indicator_type("IP Watchlist")
    address_object = Address()
    cidr = False
    if ("/" in attribute["value"]):
        ip = attribute["value"].split('/')[0]
        cidr = True
    else:
        ip = attribute["value"]
    try:
        socket.inet_aton(ip)
        ipv4 = True
    except socket.error:
        ipv4 = False
    if (cidr == True):
        address_object.category = "cidr"
    elif (ipv4 == True):
        address_object.category = "ipv4-addr"
    else:
        address_object.category = "ipv6-addr"
    if (attribute["type"] == "ip-src"):
        address_object.is_source = True
    else:
        address_object.is_source = False
    address_object.address_value = attribute["value"]
    return address_object

def generateRegkeyObservable(indicator, attribute):
    indicator.add_indicator_type("Host Characteristics")
    regkey = ""
    regvalue = ""
    if (attribute["type"] == "regkey|value"):
        regkey = attribute["value"].split('|')[0]
        regvalue = attribute["value"].split('|')[1]
    else:
        regkey = attribute["value"]
    reghive, regkey = resolveRegHive(regkey)
    reg_object = WinRegistryKey()
    reg_object.key = regkey
    reg_object.key.condition = "Equals"
    if (reghive != None):
        reg_object.hive = reghive
        reg_object.hive.condition = "Equals"
    if (regvalue != ""):
        reg_value_object = RegistryValue()
        reg_value_object.data = regvalue
        reg_value_object.data.condition = "Equals"
        reg_object.values = RegistryValues(reg_value_object)
    return reg_object

def generateSimpleObservable(indicator, attribute):
    cyboxName = misp_cybox_name[attribute["type"]]
    constructor = getattr(this_module, cyboxName, None)
    indicatorType = misp_indicator_type[attribute["type"]]
    if (indicatorType != ""):
        indicator.add_indicator_type(indicatorType)
    new_object = constructor()
    setattr(new_object, cybox_name_attribute[cyboxName], attribute["value"])
    setattr(getattr(new_object, cybox_name_attribute[cyboxName]), "condition", "Equals")
    return new_object

def generateTM(indicator, attribute):
    if (attribute["type"] == "snort"):
        tm = SnortTestMechanism()
        tm.rules = [attribute["value"]]
    else:
        # remove the following line and uncomment the code below once yara test mechanisms get added to python-stix
        return indicator
        #tm = SnortTestMechanism()
        #tm.rules = [attribute["value"]]
    indicator.test_mechanisms = [tm]

def resolveEmailObservable(indicator, attribute):
    indicator.add_indicator_type("Malicious E-mail")
    new_object = EmailMessage()
    email_header = EmailHeader()
    if (attribute["type"] == "email-src"):
        email_header.from_ = attribute["value"]
    elif(attribute["type"] == "email-dst"):
        email_header.to = attribute["value"]
    else:
        email_header.subject = attribute["value"]
    new_object.header = email_header
    return new_object

def resolveHTTPObservable(indicator, attribute):
    request_response = HTTPRequestResponse()
    client_request = HTTPClientRequest()
    if (attribute["type"] == "user-agent"):
        header = HTTPRequestHeader()
        header_fields = HTTPRequestHeaderFields()
        header_fields.user_agent = attribute["value"]
        header.parsed_header = header_fields
        client_request.http_request_header = header
    else:
        line = HTTPRequestLine()
        line.http_method = attribute["value"]
        client_request.http_request_line = line
    request_response.http_client_request = client_request
    new_object = HTTPSession()
    request_response.to_xml()
    new_object.http_request_response = [request_response]
    return new_object

# use this when implementing pattern in memory and pattern in traffic
def resolvePatternObservable(indicator, attribute):
    new_object = None
    if attribute["type"] == "pattern-in-file":
        byte_run = ByteRun()
        byte_run.byte_run_data = attribute["value"]
        new_object = File()
        new_object.byte_runs = ByteRuns(byte_run)
    # custom properties are not implemented in the API yet
    # elif attribute["type"] == "pattern-in-memory":
    # elif attribute["type"] == "pattern-in-traffic":
    return new_object

# create an artifact object for the malware-sample type.
def createArtifactObject(indicator, attribute):
    artifact = Artifact(data = attribute["data"])
    artifact.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":artifact-" + attribute["uuid"]
    observable = Observable(artifact)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-artifact-" + attribute["uuid"]
    indicator.add_observable(observable)

# return either a composition if data is set in attribute, or just an observable with a filename if it's not set
def returnAttachmentComposition(attribute):
    file_object = File()
    file_object.file_name = attribute["value"]
    file_object.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":file-" + attribute["uuid"]
    observable = Observable()
    if "data" in attribute:
        artifact = Artifact(data = attribute["data"])
        artifact.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":artifact-" + attribute["uuid"]
        observable_artifact = Observable(artifact)
        observable_artifact.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-artifact-" + attribute["uuid"]
        observable_file = Observable(file_object)
        observable_file.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-file-" + attribute["uuid"]
        composition = ObservableComposition(observables = [observable_artifact, observable_file])
        observable.observable_composition = composition
    else:
        observable = Observable(file_object)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
    if attribute["comment"] != "":
        observable.description = attribute["comment"]
    return observable

# email-attachment are mapped to an email message observable that contains the attachment as a file object
def generateEmailAttachmentObject(indicator, attribute):
    file_object = File()
    file_object.file_name = attribute["value"]
    email = EmailMessage()
    email.attachments = Attachments()
    email.add_related(file_object, "Contains", inline=True)
    file_object.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":file-" + attribute["uuid"]
    email.attachments.append(file_object.parent.id_)
    email.parent.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":EmailMessage-" + attribute["uuid"]
    observable = Observable(email)
    observable.id_ = cybox.utils.idgen.__generator.namespace.prefix + ":observable-" + attribute["uuid"]
    indicator.observable = observable

# split registry string into hive and key
def resolveRegHive(regStr):
    regStr = regStr.lstrip('\\')
    regStrU = regStr.upper()
    for hive in misp_reghive.iterkeys():
        if regStrU.startswith(hive):
            return misp_reghive[hive], regStr[len(hive):]
    return None, regStr
<|endoftext|>