<commit_before>"""
The ``wind_turbine`` module contains the class WindTurbine that implements
a wind turbine in the windpowerlib and functions needed for the modelling of a
wind turbine.

"""

import pandas as pd
import logging
import sys
import os
import numpy as np

__copyright__ = "Copyright oemof developer group"
__license__ = "GPLv3"

class WindTurbine(object):
    r"""
    Defines a standard set of wind turbine attributes.

    Parameters
    ----------
    turbine_name : string
        Name of the wind turbine type.
        Use get_turbine_types() to see a list of all wind turbines for which
        power (coefficient) curve data is provided.
    hub_height : float
        Hub height of the wind turbine in m.
    d_rotor : float
        Diameter of the rotor in m.
    cp_values : pandas.DataFrame
        Power coefficient curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power coefficient curve, the power coefficient values are listed in
        the column 'cp'. Default: None.
    p_values : pandas.DataFrame
        Power curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power curve, the power values are listed in the column 'P'.
        Default: None.
    nominal_power : float
        The nominal output of the wind turbine in kW.
    fetch_curve : string
        Parameter to specify whether the power or power coefficient curve
        should be retrieved from the provided turbine data. Default: cp.

    Attributes
    ----------
    turbine_name : string
        Name of the wind turbine type.
        Use get_turbine_types() to see a list of all wind turbines for which
        power (coefficient) curve data is provided.
    hub_height : float
        Hub height of the wind turbine in m.
    d_rotor : float
        Diameter of the rotor in m.
    cp_values : pandas.DataFrame
        Power coefficient curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power coefficient curve, the power coefficient values are listed in
        the column 'cp'. Default: None.
    p_values : pandas.DataFrame
        Power curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power curve, the power values are listed in the column 'P'.
        Default: None.
    nominal_power : float
        The nominal output of the wind turbine in kW.
    power_output : pandas.Series
        The calculated power output of the wind turbine.

    Examples
    --------
    >>> from windpowerlib import wind_turbine
    >>> enerconE126 = {
    ...    'hub_height': 135,
    ...    'd_rotor': 127,
    ...    'turbine_name': 'ENERCON E 126 7500'}
    >>> e126 = wind_turbine.WindTurbine(**enerconE126)
    >>> print(e126.d_rotor)
    127

    """

    def __init__(self, turbine_name, hub_height, d_rotor, cp_values=None,
                 p_values=None, nominal_power=None, fetch_curve='cp'):

        self.turbine_name = turbine_name
        self.hub_height = hub_height
        self.d_rotor = d_rotor
        self.cp_values = cp_values
        self.p_values = p_values
        self.nominal_power = nominal_power
        self.fetch_curve = fetch_curve

        self.power_output = None

        if (self.cp_values is None and self.p_values is None):
            p_nom = self.fetch_turbine_data()
            if self.nominal_power is None:
                self.nominal_power = p_nom

    def fetch_turbine_data(self):
        r"""
        Fetches data of the requested wind turbine.

        Method fetches nominal power as well as power coefficient curve or
        power curve from a data file provided along with the windpowerlib.

        Returns
        -------
        float
            Nominal power of the requested wind turbine.

        Examples
        --------
        >>> from windpowerlib import wind_turbine
        >>> enerconE126 = {
        ...    'hub_height': 135,
        ...    'd_rotor': 127,
        ...    'turbine_name': 'ENERCON E 126 7500'}
        >>> e126 = wind_turbine.WindTurbine(**enerconE126)
        >>> print(e126.cp_values.cp[5.0])
        0.423
        >>> print(e126.nominal_power)
        7500000.0

        """

        def restructure_data():
            df = read_turbine_data(filename=filename)
            wpp_df = df[df.turbine_id == self.turbine_name]
            if wpp_df.shape[0] == 0:
                pd.set_option('display.max_rows', len(df))
                logging.info('Possible types: \n{0}'.format(df.turbine_id))
                pd.reset_option('display.max_rows')
                sys.exit('Cannot find the wind converter type: {0}'.format(
                    self.turbine_name))

            ncols = ['turbine_id', 'p_nom', 'source', 'modificationtimestamp']
            data = np.array([0, 0])
            for col in wpp_df.keys():
                if col not in ncols:
                    if wpp_df[col].iloc[0] is not None and not np.isnan(
                            float(wpp_df[col].iloc[0])):
                        data = np.vstack((data, np.array(
                            [float(col), float(wpp_df[col])])))
            data = np.delete(data, 0, 0)
            df = pd.DataFrame(data, columns=['v_wind', self.fetch_curve])
            df.set_index('v_wind', drop=True, inplace=True)
            nominal_power = wpp_df['p_nom'].iloc[0] * 1000
            return df, nominal_power
        if self.fetch_curve == 'P':
            filename = 'P_curves.csv'
            self.p_values, p_nom = restructure_data()
        else:
            filename = 'cp_curves.csv'
            self.cp_values, p_nom = restructure_data()
        return p_nom

def read_turbine_data(**kwargs):
    r"""
    Fetches power coefficient curve or power curve from a file.

    The data files are provided along with the windpowerlib and are located in
    the directory windpowerlib/data.

    Other Parameters
    ----------------
    datapath : string, optional
        Path where the data file is stored. Default: './data'
    filename : string, optional
        Name of data file. Default: 'cp_curves.csv'

    Returns
    -------
    pandas.DataFrame
        Power coefficient curve values or power curve values with the
        corresponding wind speeds as indices.

    """
    if 'datapath' not in kwargs:
        kwargs['datapath'] = os.path.join(os.path.dirname(__file__), 'data')

    if 'filename' not in kwargs:
        kwargs['filename'] = 'cp_curves.csv'

    df = pd.read_csv(os.path.join(kwargs['datapath'], kwargs['filename']),
                     index_col=0)
    return df

def get_turbine_types(print_out=True, **kwargs):
    r"""
    Get the names of all possible wind turbine types for which the power
    coefficient curve or power curve is provided in the data files in
    the directory windpowerlib/data.

    Parameters
    ----------
    print_out : boolean
        Directly prints the list of types if set to True. Default: True

    Examples
    --------
    >>> from windpowerlib import wind_turbine
    >>> valid_types_df = wind_turbine.get_turbine_types(print_out=False)
    >>> print(valid_types_df.iloc[5])
    turbine_id    DEWIND D8 2000
    p_nom                   2000
    Name: 5, dtype: object

    """
    df = read_turbine_data(**kwargs)

    if print_out:
        pd.set_option('display.max_rows', len(df))
        print(df[['turbine_id', 'p_nom']])
        pd.reset_option('display.max_rows')
    return df[['turbine_id', 'p_nom']]
<commit_msg>Minor changes to docstring in wind_turbine<commit_after>"""
The ``wind_turbine`` module contains the class WindTurbine that implements
a wind turbine in the windpowerlib and functions needed for the modelling of a
wind turbine.

"""

import pandas as pd
import logging
import sys
import os
import numpy as np

__copyright__ = "Copyright oemof developer group"
__license__ = "GPLv3"

class WindTurbine(object):
    r"""
    Defines a standard set of wind turbine attributes.

    Parameters
    ----------
    turbine_name : string
        Name of the wind turbine type.
        Use get_turbine_types() to see a list of all wind turbines for which
        power (coefficient) curve data is provided.
    hub_height : float
        Hub height of the wind turbine in m.
    d_rotor : float
        Diameter of the rotor in m.
    cp_values : pandas.DataFrame
        Power coefficient curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power coefficient curve, the power coefficient values are listed in
        the column 'cp'. Default: None.
    p_values : pandas.DataFrame
        Power curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power curve, the power values are listed in the column 'P'.
        Default: None.
    nominal_power : float
        The nominal output of the wind turbine in kW.
    fetch_curve : string
        Parameter to specify whether the power or power coefficient curve
        should be retrieved from the provided turbine data. Default: cp.

    Attributes
    ----------
    turbine_name : string
        Name of the wind turbine type.
        Use get_turbine_types() to see a list of all wind turbines for which
        power (coefficient) curve data is provided.
    hub_height : float
        Hub height of the wind turbine in m.
    d_rotor : float
        Diameter of the rotor in m.
    cp_values : pandas.DataFrame
        Power coefficient curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power coefficient curve, the power coefficient values are listed in
        the column 'cp'. Default: None.
    p_values : pandas.DataFrame
        Power curve of the wind turbine.
        The indices of the DataFrame are the corresponding wind speeds of the
        power curve, the power values are listed in the column 'P'.
        Default: None.
    nominal_power : float
        The nominal output of the wind turbine in kW.
    fetch_curve : string
        Parameter to specify whether the power or power coefficient curve
        should be retrieved from the provided turbine data. Default: cp.
    power_output : pandas.Series
        The calculated power output of the wind turbine.

    Examples
    --------
    >>> from windpowerlib import wind_turbine
    >>> enerconE126 = {
    ...    'hub_height': 135,
    ...    'd_rotor': 127,
    ...    'turbine_name': 'ENERCON E 126 7500'}
    >>> e126 = wind_turbine.WindTurbine(**enerconE126)
    >>> print(e126.d_rotor)
    127

    """

    def __init__(self, turbine_name, hub_height, d_rotor, cp_values=None,
                 p_values=None, nominal_power=None, fetch_curve='cp'):

        self.turbine_name = turbine_name
        self.hub_height = hub_height
        self.d_rotor = d_rotor
        self.cp_values = cp_values
        self.p_values = p_values
        self.nominal_power = nominal_power
        self.fetch_curve = fetch_curve

        self.power_output = None

        if (self.cp_values is None and self.p_values is None):
            p_nom = self.fetch_turbine_data()
            if self.nominal_power is None:
                self.nominal_power = p_nom

    def fetch_turbine_data(self):
        r"""
        Fetches data of the requested wind turbine.

        Method fetches nominal power as well as power coefficient curve or
        power curve from a data file provided along with the windpowerlib.

        Returns
        -------
        float
            Nominal power of the requested wind turbine.

        Examples
        --------
        >>> from windpowerlib import wind_turbine
        >>> enerconE126 = {
        ...    'hub_height': 135,
        ...    'd_rotor': 127,
        ...    'turbine_name': 'ENERCON E 126 7500'}
        >>> e126 = wind_turbine.WindTurbine(**enerconE126)
        >>> print(e126.cp_values.cp[5.0])
        0.423
        >>> print(e126.nominal_power)
        7500000.0

        """

        def restructure_data():
            df = read_turbine_data(filename=filename)
            wpp_df = df[df.turbine_id == self.turbine_name]
            if wpp_df.shape[0] == 0:
                pd.set_option('display.max_rows', len(df))
                logging.info('Possible types: \n{0}'.format(df.turbine_id))
                pd.reset_option('display.max_rows')
                sys.exit('Cannot find the wind converter type: {0}'.format(
                    self.turbine_name))

            ncols = ['turbine_id', 'p_nom', 'source', 'modificationtimestamp']
            data = np.array([0, 0])
            for col in wpp_df.keys():
                if col not in ncols:
                    if wpp_df[col].iloc[0] is not None and not np.isnan(
                            float(wpp_df[col].iloc[0])):
                        data = np.vstack((data, np.array(
                            [float(col), float(wpp_df[col])])))
            data = np.delete(data, 0, 0)
            df = pd.DataFrame(data, columns=['v_wind', self.fetch_curve])
            df.set_index('v_wind', drop=True, inplace=True)
            nominal_power = wpp_df['p_nom'].iloc[0] * 1000
            return df, nominal_power
        if self.fetch_curve == 'P':
            filename = 'P_curves.csv'
            self.p_values, p_nom = restructure_data()
        else:
            filename = 'cp_curves.csv'
            self.cp_values, p_nom = restructure_data()
        return p_nom

def read_turbine_data(**kwargs):
    r"""
    Fetches power coefficient curve or power curve from a file.

    The data files are provided along with the windpowerlib and are located in
    the directory windpowerlib/data.

    Other Parameters
    ----------------
    datapath : string, optional
        Path where the data file is stored. Default: './data'
    filename : string, optional
        Name of data file. Default: 'cp_curves.csv'

    Returns
    -------
    pandas.DataFrame
        Power coefficient curve values or power curve values with the
        corresponding wind speeds as indices.

    """
    if 'datapath' not in kwargs:
        kwargs['datapath'] = os.path.join(os.path.dirname(__file__), 'data')

    if 'filename' not in kwargs:
        kwargs['filename'] = 'cp_curves.csv'

    df = pd.read_csv(os.path.join(kwargs['datapath'], kwargs['filename']),
                     index_col=0)
    return df

def get_turbine_types(print_out=True, **kwargs):
    r"""
    Get the names of all possible wind turbine types for which the power
    coefficient curve or power curve is provided in the data files in
    the directory windpowerlib/data.

    Parameters
    ----------
    print_out : boolean
        Directly prints the list of types if set to True. Default: True

    Examples
    --------
    >>> from windpowerlib import wind_turbine
    >>> valid_types_df = wind_turbine.get_turbine_types(print_out=False)
    >>> print(valid_types_df.iloc[5])
    turbine_id    DEWIND D8 2000
    p_nom                   2000
    Name: 5, dtype: object

    """
    df = read_turbine_data(**kwargs)

    if print_out:
        pd.set_option('display.max_rows', len(df))
        print(df[['turbine_id', 'p_nom']])
        pd.reset_option('display.max_rows')
    return df[['turbine_id', 'p_nom']]
<|endoftext|>


<commit_before>from unittest import TestCase
import mock

import jsonpickle

from cloudshell.cp.azure.common.deploy_data_holder import DeployDataHolder

from tests.helpers.test_helper import TestHelper

from cloudshell.cp.azure.domain.services.key_pair import KeyPairService

from cloudshell.cp.azure.domain.services.tags import TagService

from cloudshell.cp.azure.domain.services.network_service import NetworkService

from cloudshell.cp.azure.domain.services.virtual_machine_service import VirtualMachineService

from cloudshell.cp.azure.domain.services.storage_service import StorageService

from cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation import \
    PrepareConnectivityOperation
from mock import MagicMock, Mock

class TestPrepareConnectivity(TestCase):
    def setUp(self):
        self.storage_service = StorageService()
        self.vm_service = VirtualMachineService()
        self.network_service = NetworkService()
        self.tag_service = TagService()
        self.key_pair_service = KeyPairService(storage_service=self.storage_service)
        self.security_group_service = MagicMock()
        self.logger = MagicMock()

        self.prepare_connectivity_operation = PrepareConnectivityOperation(
            logger=self.logger,
            vm_service=self.vm_service,
            network_service=self.network_service,
            storage_service=self.storage_service,
            tags_service=self.tag_service,
            key_pair_service=self.key_pair_service,
            security_group_service=self.security_group_service)

    @mock.patch("cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation.OperationsHelper")
    def test_prepare_connectivity(self, operation_helper_class):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.network_service.get_virtual_network_by_tag = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = [1]
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act
        result = Mock()
        result.wait=Mock()

        network_client = MagicMock()
        network_client.subnets.create_or_update = Mock(return_value=result)
        network_client.virtual_networks.list = Mock(return_value="test")

        self.prepare_connectivity_operation.prepare_connectivity(
            reservation=MagicMock(),
            cloud_provider_model=MagicMock(),
            storage_client=MagicMock(),
            resource_client=MagicMock(),
            network_client=network_client,
            logger=self.logger,
            request=prepare_connectivity_request)

        # Verify
        # Created resource Group
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.vm_service.create_resource_group))

        # created Storage account
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.storage_service.create_storage_account))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_network_by_tag, 2))

        # key pair created
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.key_pair_service.save_key_pair))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(network_client.virtual_networks.list ))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(network_client.subnets.create_or_update))

    def test_extract_cidr_throws_error(self):
        action = MagicMock()

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

        action = Mock()
        att = Mock()
        att.attributeName = 'Network'
        att.attributeValue = ''
        action.customActionAttributes = [att]

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

    def test_prepare_connectivity_throes_exception_on_unavailable_VNETs(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()
        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [None, Mock()]

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = '10.0.0.0/12'
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)

        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [Mock(), None]

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)
<commit_msg>updated tests<commit_after>from unittest import TestCase
import mock

import jsonpickle

from cloudshell.cp.azure.common.deploy_data_holder import DeployDataHolder
from cloudshell.cp.azure.domain.services.security_group import SecurityGroupService

from tests.helpers.test_helper import TestHelper

from cloudshell.cp.azure.domain.services.key_pair import KeyPairService

from cloudshell.cp.azure.domain.services.tags import TagService

from cloudshell.cp.azure.domain.services.network_service import NetworkService

from cloudshell.cp.azure.domain.services.virtual_machine_service import VirtualMachineService

from cloudshell.cp.azure.domain.services.storage_service import StorageService

from cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation import \
    PrepareConnectivityOperation
from mock import MagicMock, Mock

class TestPrepareConnectivity(TestCase):
    def setUp(self):
        self.storage_service = StorageService()
        self.vm_service = VirtualMachineService()
        self.network_service = NetworkService()
        self.tag_service = TagService()
        self.key_pair_service = KeyPairService(storage_service=self.storage_service)
        self.security_group_service = SecurityGroupService()
        self.logger = MagicMock()

        self.prepare_connectivity_operation = PrepareConnectivityOperation(
            logger=self.logger,
            vm_service=self.vm_service,
            network_service=self.network_service,
            storage_service=self.storage_service,
            tags_service=self.tag_service,
            key_pair_service=self.key_pair_service,
            security_group_service=self.security_group_service)

    @mock.patch("cloudshell.cp.azure.domain.vm_management.operations.prepare_connectivity_operation.OperationsHelper")
    def test_prepare_connectivity(self, operation_helper_class):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.network_service.get_virtual_network_by_tag = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.security_group_service.create_network_security_group = Mock(return_value=Mock())
        self.security_group_service.create_network_security_group_custom_rule = Mock()

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = [1]
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act
        result = Mock()
        result.wait = Mock()

        network_client = MagicMock()
        network_client.subnets.create_or_update = Mock(return_value=result)
        network_client.virtual_networks.list = Mock(return_value="test")

        self.prepare_connectivity_operation.prepare_connectivity(
            reservation=MagicMock(),
            cloud_provider_model=MagicMock(),
            storage_client=MagicMock(),
            resource_client=MagicMock(),
            network_client=network_client,
            logger=self.logger,
            request=prepare_connectivity_request)

        # Verify
        # Created resource Group
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.vm_service.create_resource_group))

        # created Storage account
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.storage_service.create_storage_account))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.network_service.get_virtual_network_by_tag, 2))

        # key pair created
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.key_pair_service.save_key_pair))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(network_client.virtual_networks.list))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(network_client.subnets.create_or_update))

        # check the creation of the NSG
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.security_group_service.create_network_security_group))
        self.assertTrue(TestHelper.CheckMethodCalledXTimes(self.security_group_service.create_network_security_group_custom_rule,2))

    def test_extract_cidr_throws_error(self):
        action = MagicMock()

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

        action = Mock()
        att = Mock()
        att.attributeName = 'Network'
        att.attributeValue = ''
        action.customActionAttributes = [att]

        self.assertRaises(ValueError,
                          self.prepare_connectivity_operation._extract_cidr,
                          action)

    def test_prepare_connectivity_throes_exception_on_unavailable_VNETs(self):
        # Arrange
        self.key_pair_service.save_key_pair = MagicMock()
        self.key_pair_service.generate_key_pair = MagicMock()
        self.storage_service.create_storage_account = MagicMock()
        self.vm_service.create_resource_group = MagicMock()
        self.network_service.get_virtual_networks = MagicMock()
        self.network_service.create_subnet = MagicMock()
        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [None, Mock()]

        req = MagicMock()
        action = MagicMock()
        att = MagicMock()
        att.attributeName = 'Network'
        att.attributeValue = '10.0.0.0/12'
        action.customActionAttributes = [att]
        req.actions = [action]
        prepare_connectivity_request = req
        # Act

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)

        self.network_service.get_virtual_network_by_tag = Mock(return_value=None)
        self.network_service.get_virtual_network_by_tag.side_effect = [Mock(), None]

        self.assertRaises(Exception,
                          self.prepare_connectivity_operation.prepare_connectivity,
                          reservation=MagicMock(),
                          cloud_provider_model=MagicMock(),
                          storage_client=MagicMock(),
                          resource_client=MagicMock(),
                          network_client=MagicMock(),
                          logger=self.logger,
                          request=prepare_connectivity_request)
<|endoftext|>


<commit_before>import sublime
import sublime_plugin
import json
import os
import re
import gscommon as gs
import mg9
from os.path import basename
from os.path import dirname

AC_OPTS = sublime.INHIBIT_WORD_COMPLETIONS | sublime.INHIBIT_EXPLICIT_COMPLETIONS
REASONABLE_PKGNAME_PAT = re.compile(r'^\w+$')

last_gopath = ''
END_SELECTOR_PAT = re.compile(r'.*?((?:[\w.]+\.)?(\w+))$')
START_SELECTOR_PAT = re.compile(r'^([\w.]+)')
DOMAIN = 'GsComplete'
SNIPPET_VAR_PAT = re.compile(r'\$\{([a-zA-Z]\w*)\}')

def snippet_match(ctx, m):
	try:
		for k,p in m.get('match', {}).iteritems():
			q = ctx.get(k, '')
			if p and gs.is_a_string(p):
				if not re.search(p, str(q)):
					return False
			elif p != q:
				return False
	except:
		gs.notice(DOMAIN, gs.traceback())
	return True

def expand_snippet_vars(vars, text, title, value):
	sub = lambda m: vars.get(m.group(1), '')
	return (
		SNIPPET_VAR_PAT.sub(sub, text),
		SNIPPET_VAR_PAT.sub(sub, title),
		SNIPPET_VAR_PAT.sub(sub, value)
	)

def resolve_snippets(ctx):
	cl = set()
	types = [''] if ctx.get('local') else ctx.get('types')
	vars = {}
	for k,v in ctx.iteritems():
		if gs.is_a_string(v):
			vars[k] = v

	try:
		snips = []
		snips.extend(gs.setting('default_snippets', []))
		snips.extend(gs.setting('snippets', []))
		for m in snips:
			try:
				if snippet_match(ctx, m):
					for ent in m.get('snippets', []):
						text = ent.get('text', '')
						title = ent.get('title', '')
						value = ent.get('value', '')
						if text and value:
							for typename in types:
								vars['typename'] = typename
								vars['typename_abbr'] = typename[0].lower() if typename else ''
								txt, ttl, val = expand_snippet_vars(vars, text, title, value)
								s = u'%s\t%s \u0282' % (txt, ttl)
								cl.add((s, val))
			except:
				gs.notice(DOMAIN, gs.traceback())
	except:
		gs.notice(DOMAIN, gs.traceback())
	return list(cl)

class GoSublime(sublime_plugin.EventListener):
	gocode_set = False
	def on_query_completions(self, view, prefix, locations):
		pos = locations[0]
		scopes = view.scope_name(pos).split()
		if ('source.go' not in scopes) or (gs.setting('gscomplete_enabled', False) is not True):
			return []

		if gs.IGNORED_SCOPES.intersection(scopes):
			return ([], AC_OPTS)

		types = []
		for r in view.find_by_selector('source.go keyword.control.go'):
			if view.substr(r) == 'type':
				end = r.end()
				r = view.find(r'\s+(\w+)', end)
				if r.begin() == end:
					types.append(view.substr(r).lstrip())

		try:
			default_pkgname = basename(dirname(view.file_name()))
		except Exception:
			default_pkgname = ''

		if not REASONABLE_PKGNAME_PAT.match(default_pkgname):
			default_pkgname = ''

		r = view.find('package\s+(\w+)', 0)
		ctx = {
			'global': True,
			'pkgname': view.substr(view.word(r.end())) if r else '',
			'types': types or [''],
			'has_types': len(types) > 0,
			'default_pkgname': default_pkgname,
			'fn': view.file_name() or '',
		}
		show_snippets = gs.setting('autocomplete_snippets', True) is True

		if not ctx.get('pkgname'):
			return (resolve_snippets(ctx), AC_OPTS) if show_snippets else ([], AC_OPTS)

		# gocode is case-sesitive so push the location back to the 'dot' so it gives
		# gives us everything then st2 can pick the matches for us
		offset = pos - len(prefix)
		src = view.substr(sublime.Region(0, view.size()))

		fn = view.file_name() or '<stdin>'
		if not src:
			return ([], AC_OPTS)

		nc = view.substr(sublime.Region(pos, pos+1))
		cl = self.complete(fn, offset, src, nc.isalpha() or nc == "(")

		pc = view.substr(sublime.Region(pos-1, pos))
		if show_snippets and (pc.isspace() or pc.isalpha()):
			if scopes[-1] == 'source.go':
				cl.extend(resolve_snippets(ctx))
			elif scopes[-1] == 'meta.block.go' and ('meta.function.plain.go' in scopes or 'meta.function.receiver.go' in scopes):
				ctx['global'] = False
				ctx['local'] = True
				cl.extend(resolve_snippets(ctx))
		return (cl, AC_OPTS)

	def find_end_pt(self, view, pat, start, end, flags=sublime.LITERAL):
		r = view.find(pat, start, flags)
		return r.end() if r and r.end() < end else -1

	def complete(self, fn, offset, src, func_name_only):
		comps = []
		offset = 'c%s' % offset
		js, err, _ = mg9.gocode(["-f=json", "autocomplete", fn, offset], input=src)
		if err:
			gs.notice(DOMAIN, err)
		else:
			try:
				js = json.loads(js)
				if js and js[1]:
					for ent in js[1]:
						tn = ent['type']
						cn = ent['class']
						nm = ent['name']
						sfx = self.typeclass_prefix(cn, tn)
						if cn == 'func':
							if nm in ('main', 'init'):
								continue
							act = gs.setting('autocomplete_tests', False)
							if not act and nm.startswith(('Test', 'Benchmark', 'Example')):
								continue

							params, ret = declex(tn)
							ret = ret.strip('() ')
							if func_name_only:
								a = nm
							else:
								a = []
								for i, p in enumerate(params):
									n, t = p
									if t.startswith('...'):
										n = '...'
									a.append('${%d:%s}' % (i+1, n))
								a = '%s(%s)' % (nm, ', '.join(a))
							comps.append(('%s\t%s %s' % (nm, ret, sfx), a))
						elif cn != 'PANIC':
							comps.append(('%s\t%s %s' % (nm, tn, sfx), nm))
			except KeyError as e:
				gs.notice(DOMAIN, 'Error while running gocode, possibly malformed data returned: %s' % e)
			except ValueError as e:
				gs.notice(DOMAIN, "Error while decoding gocode output: %s" % e)
		return comps

	def typeclass_prefix(self, typeclass, typename):
		return gs.NAME_PREFIXES.get(typename, gs.CLASS_PREFIXES.get(typeclass, ' '))

def declex(s):
	params = []
	ret = ''
	if s.startswith('func('):
		lp = len(s)
		sp = 5
		ep = sp
		dc = 1
		names = []
		while ep < lp and dc > 0:
			c = s[ep]
			if dc == 1 and c in (',', ')'):
				if sp < ep:
					n, _, t = s[sp:ep].strip().partition(' ')
					t = t.strip()
					if t:
						for name in names:
							params.append((name, t))
						names = []
						params.append((n, t))
					else:
						names.append(n)
					sp = ep + 1
			if c == '(':
				dc += 1
			elif c == ')':
				dc -= 1
			ep += 1
		ret = s[ep:].strip() if ep < lp else ''
	return (params, ret)

class GsShowCallTip(sublime_plugin.TextCommand):
	def is_enabled(self):
		return gs.is_go_source_view(self.view)

	def show_hint(self, s):
		dmn = '%s.completion-hint' % DOMAIN
		gs.show_output(dmn, s, print_output=False, syntax_file='GsDoc')

	def run(self, edit):
		view = self.view
		pt = gs.sel(view).begin()
		if view.substr(sublime.Region(pt-1, pt)) == '(':
			depth = 1
		else:
			depth = 0
		c = ''
		while True:
			line = view.line(pt)
			scope = view.scope_name(pt)
			if 'string' in scope or 'comment' in scope:
				pt = view.extract_scope(pt).begin() - 1
				continue

			c = view.substr(sublime.Region(pt-1, pt))
			if not c:
				pt = -1
				break

			if c.isalpha() and depth >= 0:
				while c.isalpha() or c == '.':
					pt += 1
					c = view.substr(sublime.Region(pt-1, pt))

				# curly braces ftw
				break # break outer while loop
			if c == ')':
				depth -= 1
			elif c == '(':
				depth += 1
				i = pt
				while True:
					pc = view.substr(sublime.Region(i-1, i))
					if pc == '.' or pc.isalpha():
						i -= 1
					else:
						break

				if i != pt:
					pt = i
					continue

			pt -= 1
			if pt <= line.begin():
				pt = -1
				break

		while not c.isalpha() and pt > 0:
			pt -= 1
			c = view.substr(sublime.Region(pt-1, pt))

		if pt <= 0 or view.scope_name(pt).strip() == 'source.go':
			self.show_hint("// can't find selector")
			return

		line = view.line(pt)
		line_start = line.begin()

		s = view.substr(line)
		if not s:
			self.show_hint('// no source')
			return

		scopes = [
			'support.function.any-method.go',
			'meta.function-call.go',
			'support.function.builtin.go',
		]
		found = False
		while True:
			scope = view.scope_name(pt).strip()
			for s in scopes:
				if scope.endswith(s):
					found = True
					break

			if found or pt <= line_start:
				break

			pt -= 1

		if not found:
			self.show_hint("// can't find function call")
			return

		s = view.substr(sublime.Region(line_start, pt))
		m = END_SELECTOR_PAT.match(s)
		if not m:
			self.show_hint("// can't match selector")
			return

		offset = (line_start + m.end())
		coffset = 'c%d' % offset
		sel = m.group(1)
		name = m.group(2)
		candidates = []
		src = view.substr(sublime.Region(0, view.size()))
		fn = view.file_name() or '<stdin>'
		cmd = gs.setting('gocode_cmd', 'gocode')
		args = [cmd, "-f=json", "autocomplete", fn, coffset]
		js, err, _ = gs.runcmd(args, src)
		if err:
			gs.notice(DOMAIN, err)
		else:
			try:
				js = json.loads(js)
				if js and js[1]:
					candidates = js[1]
			except:
				pass

		c = {}
		for i in candidates:
			if i['name'] == name:
				if c:
					c = None
					break
				c = i

		if not c:
			self.show_hint('// no candidates found')
			return

		s = '// %s %s\n%s' % (c['name'], c['class'], c['type'])
		self.show_hint(s)

<commit_msg>* impl autocomplete_closures<commit_after>import sublime
import sublime_plugin
import json
import os
import re
import gscommon as gs
import mg9
from os.path import basename
from os.path import dirname

AC_OPTS = sublime.INHIBIT_WORD_COMPLETIONS | sublime.INHIBIT_EXPLICIT_COMPLETIONS
REASONABLE_PKGNAME_PAT = re.compile(r'^\w+$')

last_gopath = ''
END_SELECTOR_PAT = re.compile(r'.*?((?:[\w.]+\.)?(\w+))$')
START_SELECTOR_PAT = re.compile(r'^([\w.]+)')
DOMAIN = 'GsComplete'
SNIPPET_VAR_PAT = re.compile(r'\$\{([a-zA-Z]\w*)\}')

def snippet_match(ctx, m):
	try:
		for k,p in m.get('match', {}).iteritems():
			q = ctx.get(k, '')
			if p and gs.is_a_string(p):
				if not re.search(p, str(q)):
					return False
			elif p != q:
				return False
	except:
		gs.notice(DOMAIN, gs.traceback())
	return True

def expand_snippet_vars(vars, text, title, value):
	sub = lambda m: vars.get(m.group(1), '')
	return (
		SNIPPET_VAR_PAT.sub(sub, text),
		SNIPPET_VAR_PAT.sub(sub, title),
		SNIPPET_VAR_PAT.sub(sub, value)
	)

def resolve_snippets(ctx):
	cl = set()
	types = [''] if ctx.get('local') else ctx.get('types')
	vars = {}
	for k,v in ctx.iteritems():
		if gs.is_a_string(v):
			vars[k] = v

	try:
		snips = []
		snips.extend(gs.setting('default_snippets', []))
		snips.extend(gs.setting('snippets', []))
		for m in snips:
			try:
				if snippet_match(ctx, m):
					for ent in m.get('snippets', []):
						text = ent.get('text', '')
						title = ent.get('title', '')
						value = ent.get('value', '')
						if text and value:
							for typename in types:
								vars['typename'] = typename
								vars['typename_abbr'] = typename[0].lower() if typename else ''
								txt, ttl, val = expand_snippet_vars(vars, text, title, value)
								s = u'%s\t%s \u0282' % (txt, ttl)
								cl.add((s, val))
			except:
				gs.notice(DOMAIN, gs.traceback())
	except:
		gs.notice(DOMAIN, gs.traceback())
	return list(cl)

class GoSublime(sublime_plugin.EventListener):
	gocode_set = False
	def on_query_completions(self, view, prefix, locations):
		pos = locations[0]
		scopes = view.scope_name(pos).split()
		if ('source.go' not in scopes) or (gs.setting('gscomplete_enabled', False) is not True):
			return []

		if gs.IGNORED_SCOPES.intersection(scopes):
			return ([], AC_OPTS)

		types = []
		for r in view.find_by_selector('source.go keyword.control.go'):
			if view.substr(r) == 'type':
				end = r.end()
				r = view.find(r'\s+(\w+)', end)
				if r.begin() == end:
					types.append(view.substr(r).lstrip())

		try:
			default_pkgname = basename(dirname(view.file_name()))
		except Exception:
			default_pkgname = ''

		if not REASONABLE_PKGNAME_PAT.match(default_pkgname):
			default_pkgname = ''

		r = view.find('package\s+(\w+)', 0)
		ctx = {
			'global': True,
			'pkgname': view.substr(view.word(r.end())) if r else '',
			'types': types or [''],
			'has_types': len(types) > 0,
			'default_pkgname': default_pkgname,
			'fn': view.file_name() or '',
		}
		show_snippets = gs.setting('autocomplete_snippets', True) is True

		if not ctx.get('pkgname'):
			return (resolve_snippets(ctx), AC_OPTS) if show_snippets else ([], AC_OPTS)

		# gocode is case-sesitive so push the location back to the 'dot' so it gives
		# gives us everything then st2 can pick the matches for us
		offset = pos - len(prefix)
		src = view.substr(sublime.Region(0, view.size()))

		fn = view.file_name() or '<stdin>'
		if not src:
			return ([], AC_OPTS)

		nc = view.substr(sublime.Region(pos, pos+1))
		cl = self.complete(fn, offset, src, nc.isalpha() or nc == "(")

		pc = view.substr(sublime.Region(pos-1, pos))
		if show_snippets and (pc.isspace() or pc.isalpha()):
			if scopes[-1] == 'source.go':
				cl.extend(resolve_snippets(ctx))
			elif scopes[-1] == 'meta.block.go' and ('meta.function.plain.go' in scopes or 'meta.function.receiver.go' in scopes):
				ctx['global'] = False
				ctx['local'] = True
				cl.extend(resolve_snippets(ctx))
		return (cl, AC_OPTS)

	def find_end_pt(self, view, pat, start, end, flags=sublime.LITERAL):
		r = view.find(pat, start, flags)
		return r.end() if r and r.end() < end else -1

	def complete(self, fn, offset, src, func_name_only):
		comps = []
		offset = 'c%s' % offset
		autocomplete_tests = gs.setting('autocomplete_tests', False)
		autocomplete_closures = gs.setting('autocomplete_closures', False)
		js, err, _ = mg9.gocode(["-f=json", "autocomplete", fn, offset], input=src)
		if err:
			gs.notice(DOMAIN, err)
		else:
			try:
				js = json.loads(js)
				if js and js[1]:
					for ent in js[1]:
						tn = ent['type']
						cn = ent['class']
						nm = ent['name']
						is_func = (cn == 'func')
						is_func_type = (cn == 'type' and tn.startswith('func('))

						if is_func:
							if nm in ('main', 'init'):
								continue

							if not autocomplete_tests and nm.startswith(('Test', 'Benchmark', 'Example')):
								continue

						if is_func or is_func_type:
							s_sfx = u'\u0282'
							t_sfx = gs.CLASS_PREFIXES.get('type', '')
							f_sfx = gs.CLASS_PREFIXES.get('func', '')
							params, ret = declex(tn)
							decl = []
							for i, p in enumerate(params):
								n, t = p
								if t.startswith('...'):
									n = '...'
								decl.append('${%d:%s}' % (i+1, n))
							decl = ', '.join(decl)
							ret = ret.strip('() ')

							if is_func:
								if func_name_only:
									comps.append((
										'%s\t%s %s' % (nm, ret, f_sfx),
										nm,
									))
								else:
									comps.append((
										'%s\t%s %s' % (nm, ret, f_sfx),
										'%s(%s)' % (nm, decl),
									))
							else:
								comps.append((
									'%s\t%s %s' % (nm, tn, t_sfx),
									nm,
								))

								if autocomplete_closures:
									comps.append((
										'%s {}\tfunc() {...} %s' % (nm, s_sfx),
										'%s {\n\t${0}\n}' % tn,
									))
						elif cn != 'PANIC':
							comps.append((
								'%s\t%s %s' % (nm, tn, self.typeclass_prefix(cn, tn)),
								nm,
							))
			except KeyError as e:
				gs.notice(DOMAIN, 'Error while running gocode, possibly malformed data returned: %s' % e)
			except ValueError as e:
				gs.notice(DOMAIN, "Error while decoding gocode output: %s" % e)
		return comps

	def typeclass_prefix(self, typeclass, typename):
		return gs.NAME_PREFIXES.get(typename, gs.CLASS_PREFIXES.get(typeclass, ' '))

def declex(s):
	params = []
	ret = ''
	if s.startswith('func('):
		lp = len(s)
		sp = 5
		ep = sp
		dc = 1
		names = []
		while ep < lp and dc > 0:
			c = s[ep]
			if dc == 1 and c in (',', ')'):
				if sp < ep:
					n, _, t = s[sp:ep].strip().partition(' ')
					t = t.strip()
					if t:
						for name in names:
							params.append((name, t))
						names = []
						params.append((n, t))
					else:
						names.append(n)
					sp = ep + 1
			if c == '(':
				dc += 1
			elif c == ')':
				dc -= 1
			ep += 1
		ret = s[ep:].strip() if ep < lp else ''
	return (params, ret)

class GsShowCallTip(sublime_plugin.TextCommand):
	def is_enabled(self):
		return gs.is_go_source_view(self.view)

	def show_hint(self, s):
		dmn = '%s.completion-hint' % DOMAIN
		gs.show_output(dmn, s, print_output=False, syntax_file='GsDoc')

	def run(self, edit):
		view = self.view
		pt = gs.sel(view).begin()
		if view.substr(sublime.Region(pt-1, pt)) == '(':
			depth = 1
		else:
			depth = 0
		c = ''
		while True:
			line = view.line(pt)
			scope = view.scope_name(pt)
			if 'string' in scope or 'comment' in scope:
				pt = view.extract_scope(pt).begin() - 1
				continue

			c = view.substr(sublime.Region(pt-1, pt))
			if not c:
				pt = -1
				break

			if c.isalpha() and depth >= 0:
				while c.isalpha() or c == '.':
					pt += 1
					c = view.substr(sublime.Region(pt-1, pt))

				# curly braces ftw
				break # break outer while loop
			if c == ')':
				depth -= 1
			elif c == '(':
				depth += 1
				i = pt
				while True:
					pc = view.substr(sublime.Region(i-1, i))
					if pc == '.' or pc.isalpha():
						i -= 1
					else:
						break

				if i != pt:
					pt = i
					continue

			pt -= 1
			if pt <= line.begin():
				pt = -1
				break

		while not c.isalpha() and pt > 0:
			pt -= 1
			c = view.substr(sublime.Region(pt-1, pt))

		if pt <= 0 or view.scope_name(pt).strip() == 'source.go':
			self.show_hint("// can't find selector")
			return

		line = view.line(pt)
		line_start = line.begin()

		s = view.substr(line)
		if not s:
			self.show_hint('// no source')
			return

		scopes = [
			'support.function.any-method.go',
			'meta.function-call.go',
			'support.function.builtin.go',
		]
		found = False
		while True:
			scope = view.scope_name(pt).strip()
			for s in scopes:
				if scope.endswith(s):
					found = True
					break

			if found or pt <= line_start:
				break

			pt -= 1

		if not found:
			self.show_hint("// can't find function call")
			return

		s = view.substr(sublime.Region(line_start, pt))
		m = END_SELECTOR_PAT.match(s)
		if not m:
			self.show_hint("// can't match selector")
			return

		offset = (line_start + m.end())
		coffset = 'c%d' % offset
		sel = m.group(1)
		name = m.group(2)
		candidates = []
		src = view.substr(sublime.Region(0, view.size()))
		fn = view.file_name() or '<stdin>'
		cmd = gs.setting('gocode_cmd', 'gocode')
		args = [cmd, "-f=json", "autocomplete", fn, coffset]
		js, err, _ = gs.runcmd(args, src)
		if err:
			gs.notice(DOMAIN, err)
		else:
			try:
				js = json.loads(js)
				if js and js[1]:
					candidates = js[1]
			except:
				pass

		c = {}
		for i in candidates:
			if i['name'] == name:
				if c:
					c = None
					break
				c = i

		if not c:
			self.show_hint('// no candidates found')
			return

		s = '// %s %s\n%s' % (c['name'], c['class'], c['type'])
		self.show_hint(s)

<|endoftext|>


<commit_before># Copyright (C) 2015 eSentire, Inc (jacob.gajek@esentire.com).
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import datetime
import logging
import random
import re
import requests
import ssl
import time

from lib.cuckoo.common.abstracts import Machinery
from lib.cuckoo.common.exceptions import CuckooMachineError
from lib.cuckoo.common.exceptions import CuckooDependencyError
from lib.cuckoo.common.exceptions import CuckooCriticalError

try:
    from pyVim.connect import SmartConnection
    HAVE_PYVMOMI = True
except ImportError:
    HAVE_PYVMOMI = False

log = logging.getLogger(__name__)
logging.getLogger("requests").setLevel(logging.WARNING)

class vSphere(Machinery):
    """vSphere/ESXi machinery class based on pyVmomi Python SDK."""

    # VM states
    RUNNING = "poweredOn"
    POWEROFF = "poweredOff"
    SUSPENDED = "suspended"
    ABORTED = "aborted"

    def __init__(self):
        if not HAVE_PYVMOMI:
            raise CuckooDependencyError(
                "Couldn't import pyVmomi, please install it (using "
                "`pip install -U pyvmomi`)"
            )

        super(vSphere, self).__init__()

    def _initialize(self, module_name):
        """Read configuration.
        @param module_name: module name.
        """
        super(vSphere, self)._initialize(module_name)

        # Initialize random number generator
        random.seed()

    def _initialize_check(self):
        """Runs checks against virtualization software when a machine manager
        is initialized.
        @raise CuckooCriticalError: if a misconfiguration or unsupported state
                                    is found.
        """
        self.connect_opts = {}

        if self.options.vsphere.host:
            self.connect_opts["host"] = self.options.vsphere.host
        else:
            raise CuckooCriticalError(
                "vSphere host address setting not found, please add it "
                "to the config file."
            )

        if self.options.vsphere.port:
            self.connect_opts["port"] = self.options.vsphere.port
        else:
            raise CuckooCriticalError(
                "vSphere port setting not found, please add it to the "
                "config file."
            )

        if self.options.vsphere.user:
            self.connect_opts["user"] = self.options.vsphere.user
        else:
            raise CuckooCriticalError(
                "vSphere username setting not found, please add it to "
                "the config file."
            )

        if self.options.vsphere.pwd:
            self.connect_opts["pwd"] = self.options.vsphere.pwd
        else:
            raise CuckooCriticalError(
                "vSphere password setting not found, please add it to "
                "the config file."
            )

        # Workaround for PEP-0476 issues in recent Python versions
        if self.options.vsphere.unverified_ssl:
            sslContext = ssl.SSLContext(ssl.PROTOCOL_TLSv1)
            sslContext.verify_mode = ssl.CERT_NONE
            self.connect_opts["sslContext"] = sslContext
            log.warn("Turning off SSL certificate verification!")

        # Check that a snapshot is configured for each machine
        # and that it was taken in a powered-on state
        try:
            with SmartConnection(**self.connect_opts) as conn:
                for machine in self.machines():
                    if not machine.snapshot:
                        raise CuckooCriticalError(
                            "Snapshot name not specified for machine %s, "
                            "please add it to the config file." %
                            machine.label
                        )

                    vm = self._get_virtual_machine_by_label(conn, machine.label)
                    if not vm:
                        raise CuckooCriticalError(
                            "Unable to find machine %s on vSphere host, "
                            "please update your configuration." %
                            machine.label
                        )

                    state = self._get_snapshot_power_state(vm, machine.snapshot)
                    if not state:
                        raise CuckooCriticalError(
                            "Unable to find snapshot %s for machine %s, "
                            "please update your configuration." %
                            (machine.snapshot, machine.label)
                        )

                    if state != self.RUNNING:
                        raise CuckooCriticalError(
                            "Snapshot for machine %s not in powered-on "
                            "state, please create one." % machine.label
                        )
        except CuckooCriticalError as e:
            raise e
        except Exception as e:
            raise CuckooCriticalError(
                "Couldn't connect to vSphere host: %s" % e
            )

        super(vSphere, self)._initialize_check()

    def start(self, label, task):
        """Start a machine.
        @param label: machine name.
        @param task: task object.
        @raise CuckooMachineError: if unable to start machine.
        """
        name = self.db.view_machine_by_label(label).snapshot
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._revert_snapshot(vm, name)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def stop(self, label):
        """Stop a machine.
        @param label: machine name.
        @raise CuckooMachineError: if unable to stop machine
        """
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._stop_virtual_machine(vm)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def dump_memory(self, label, path):
        """Take a memory dump of a machine.
        @param path: path to where to store the memory dump
        @raise CuckooMachineError: if error taking the memory dump
        """
        name = "cuckoo_memdump_{0}".format(random.randint(100000, 999999))
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._create_snapshot(vm, name)
                self._download_snapshot(conn, vm, name, path)
                self._delete_snapshot(vm, name)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def _list(self):
        """List virtual machines on vSphere host"""
        ret = []
        with SmartConnection(**self.connect_opts) as conn:
            for vm in self._get_virtual_machines(conn):
                ret.append(vm.summary.config.name)
        return ret

    def _status(self, label):
        """Get power state of vm from vSphere host.
        @param label: virtual machine name
        @raise CuckooMachineError: if error getting status or machine not found
        """
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if not vm:
                raise CuckooMachineError(
                    "Machine %s not found on server" % label
                )

            status = vm.runtime.powerState
            self.set_status(label, status)
            return status

    def _get_virtual_machines(self, conn):
        """Iterate over all VirtualMachine managed objects on vSphere host"""
        def traverseDCFolders(conn, nodes, path=""):
            for node in nodes:
                if hasattr(node, "childEntity"):
                    for child, childpath in traverseDCFolders(conn, node.childEntity, path + node.name + "/"):
                        yield child, childpath
                else:
                    yield node, path + node.name

        def traverseVMFolders(conn, nodes):
            for node in nodes:
                if hasattr(node, "childEntity"):
                    for child in traverseVMFolders(conn, node.childEntity):
                        yield child
                else:
                    yield node

        self.VMtoDC = {}

        for dc, dcpath in traverseDCFolders(conn, conn.content.rootFolder.childEntity):
            for vm in traverseVMFolders(conn, dc.vmFolder.childEntity):
                self.VMtoDC[vm.summary.config.name] = dcpath
                yield vm

    def _get_virtual_machine_by_label(self, conn, label):
        """Return the named VirtualMachine managed object"""
        for vm in self._get_virtual_machines(conn):
            if vm.summary.config.name == label:
                return vm

    def _get_snapshot_by_name(self, vm, name):
        """Return the named VirtualMachineSnapshot managed object for
        a virtual machine"""
        for ss in self._traverseSnapshots(vm.snapshot.rootSnapshotList):
            if ss.name == name:
                return ss.snapshot

    def _get_snapshot_power_state(self, vm, name):
        """Return the power state for a named VirtualMachineSnapshot object"""
        for ss in self._traverseSnapshots(vm.snapshot.rootSnapshotList):
            if ss.name == name:
                return ss.state

    def _create_snapshot(self, vm, name):
        """Create named snapshot of virtual machine"""
        log.info(
            "Creating snapshot %s for machine %s",
            name, vm.summary.config.name
        )

        task = vm.CreateSnapshot_Task(name=name,
                                      description="Created by Cuckoo sandbox",
                                      memory=True,
                                      quiesce=False)
        try:
            self._wait_task(task)
        except CuckooMachineError as e:
            raise CuckooMachineError("CreateSnapshot: %s" % e)

    def _delete_snapshot(self, vm, name):
        """Remove named snapshot of virtual machine"""
        snapshot = self._get_snapshot_by_name(vm, name)
        if snapshot:
            log.info(
                "Removing snapshot %s for machine %s",
                name, vm.summary.config.name
            )

            task = snapshot.RemoveSnapshot_Task(removeChildren=True)
            try:
                self._wait_task(task)
            except CuckooMachineError as e:
                log.error("RemoveSnapshot: {0}".format(e))
        else:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

    def _revert_snapshot(self, vm, name):
        """Revert virtual machine to named snapshot"""
        snapshot = self._get_snapshot_by_name(vm, name)
        if snapshot:
            log.info(
                "Reverting machine %s to snapshot %s",
                vm.summary.config.name, name
            )

            task = snapshot.RevertToSnapshot_Task()
            try:
                self._wait_task(task)
            except CuckooMachineError as e:
                raise CuckooMachineError("RevertToSnapshot: %s" % e)
        else:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

    def _download_snapshot(self, conn, vm, name, path):
        """Download snapshot file from host to local path"""

        # Get filespec to .vmsn or .vmem file of named snapshot
        snapshot = self._get_snapshot_by_name(vm, name)
        if not snapshot:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

        memorykey = datakey = filespec = None
        for s in vm.layoutEx.snapshot:
            if s.key == snapshot:
                memorykey = s.memoryKey
                datakey = s.dataKey
                break

        for f in vm.layoutEx.file:
            if f.key == memorykey and (f.type == "snapshotMemory" or
                                       f.type == "suspendMemory"):
                filespec = f.name
                break

        if not filespec:
            for f in vm.layoutEx.file:
                if f.key == datakey and f.type == "snapshotData":
                    filespec = f.name
                    break

        if not filespec:
            raise CuckooMachineError("Could not find snapshot memory file")

        log.info("Downloading memory dump %s to %s", filespec, path)

        # Parse filespec to get datastore and file path.
        datastore, filepath = re.match(r"\[([^\]]*)\] (.*)", filespec).groups()

        # Construct URL request
        params = {
            "dsName": datastore,
            "dcPath": self.VMtoDC.get(vm.summary.config.name, "ha-datacenter")
        }
        headers = {
            "Cookie": conn._stub.cookie,
        }
        url = "https://%s:%s/folder/%s" % (
            self.connect_opts["host"], self.connect_opts["port"], filepath
        )

        # Stream download to specified local path
        try:
            response = requests.get(url, params=params, headers=headers,
                                    verify=False, stream=True)

            response.raise_for_status()

            with open(path, "wb") as localfile:
                for chunk in response.iter_content(16*1024):
                    localfile.write(chunk)

        except Exception as e:
            raise CuckooMachineError(
                "Error downloading memory dump %s: %s" %
                (filespec, e)
            )

    def _stop_virtual_machine(self, vm):
        """Power off a virtual machine"""
        log.info("Powering off virtual machine %s", vm.summary.config.name)
        task = vm.PowerOffVM_Task()
        try:
            self._wait_task(task)
        except CuckooMachineError as e:
            log.error("PowerOffVM: %s", e)

    def _wait_task(self, task):
        """Wait for a task to complete with timeout"""
        limit = datetime.timedelta(
            seconds=int(self.options_globals.timeouts.vm_state)
        )
        start = datetime.datetime.utcnow()

        while True:
            if task.info.state == "error":
                raise CuckooMachineError("Task error")

            if task.info.state == "success":
                break

            if datetime.datetime.utcnow() - start > limit:
                raise CuckooMachineError("Task timed out")

            time.sleep(1)

    def _traverseSnapshots(self, root):
        """Recursive depth-first traversal of snapshot tree"""
        for node in root:
            if len(node.childSnapshotList) > 0:
                for child in self._traverseSnapshots(node.childSnapshotList):
                    yield child
            yield node
<commit_msg>Fixed AttributeError when traversing through VMs.<commit_after># Copyright (C) 2015 eSentire, Inc (jacob.gajek@esentire.com).
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import datetime
import logging
import random
import re
import requests
import ssl
import time

from lib.cuckoo.common.abstracts import Machinery
from lib.cuckoo.common.exceptions import CuckooMachineError
from lib.cuckoo.common.exceptions import CuckooDependencyError
from lib.cuckoo.common.exceptions import CuckooCriticalError

try:
    from pyVim.connect import SmartConnection
    HAVE_PYVMOMI = True
except ImportError:
    HAVE_PYVMOMI = False

log = logging.getLogger(__name__)
logging.getLogger("requests").setLevel(logging.WARNING)

class vSphere(Machinery):
    """vSphere/ESXi machinery class based on pyVmomi Python SDK."""

    # VM states
    RUNNING = "poweredOn"
    POWEROFF = "poweredOff"
    SUSPENDED = "suspended"
    ABORTED = "aborted"

    def __init__(self):
        if not HAVE_PYVMOMI:
            raise CuckooDependencyError(
                "Couldn't import pyVmomi, please install it (using "
                "`pip install -U pyvmomi`)"
            )

        super(vSphere, self).__init__()

    def _initialize(self, module_name):
        """Read configuration.
        @param module_name: module name.
        """
        super(vSphere, self)._initialize(module_name)

        # Initialize random number generator
        random.seed()

    def _initialize_check(self):
        """Runs checks against virtualization software when a machine manager
        is initialized.
        @raise CuckooCriticalError: if a misconfiguration or unsupported state
                                    is found.
        """
        self.connect_opts = {}

        if self.options.vsphere.host:
            self.connect_opts["host"] = self.options.vsphere.host
        else:
            raise CuckooCriticalError(
                "vSphere host address setting not found, please add it "
                "to the config file."
            )

        if self.options.vsphere.port:
            self.connect_opts["port"] = self.options.vsphere.port
        else:
            raise CuckooCriticalError(
                "vSphere port setting not found, please add it to the "
                "config file."
            )

        if self.options.vsphere.user:
            self.connect_opts["user"] = self.options.vsphere.user
        else:
            raise CuckooCriticalError(
                "vSphere username setting not found, please add it to "
                "the config file."
            )

        if self.options.vsphere.pwd:
            self.connect_opts["pwd"] = self.options.vsphere.pwd
        else:
            raise CuckooCriticalError(
                "vSphere password setting not found, please add it to "
                "the config file."
            )

        # Workaround for PEP-0476 issues in recent Python versions
        if self.options.vsphere.unverified_ssl:
            sslContext = ssl.SSLContext(ssl.PROTOCOL_TLSv1)
            sslContext.verify_mode = ssl.CERT_NONE
            self.connect_opts["sslContext"] = sslContext
            log.warn("Turning off SSL certificate verification!")

        # Check that a snapshot is configured for each machine
        # and that it was taken in a powered-on state
        try:
            with SmartConnection(**self.connect_opts) as conn:
                for machine in self.machines():
                    if not machine.snapshot:
                        raise CuckooCriticalError(
                            "Snapshot name not specified for machine %s, "
                            "please add it to the config file." %
                            machine.label
                        )

                    vm = self._get_virtual_machine_by_label(conn, machine.label)
                    if not vm:
                        raise CuckooCriticalError(
                            "Unable to find machine %s on vSphere host, "
                            "please update your configuration." %
                            machine.label
                        )

                    state = self._get_snapshot_power_state(vm, machine.snapshot)
                    if not state:
                        raise CuckooCriticalError(
                            "Unable to find snapshot %s for machine %s, "
                            "please update your configuration." %
                            (machine.snapshot, machine.label)
                        )

                    if state != self.RUNNING:
                        raise CuckooCriticalError(
                            "Snapshot for machine %s not in powered-on "
                            "state, please create one." % machine.label
                        )
        except CuckooCriticalError as e:
            raise e
        except Exception as e:
            raise CuckooCriticalError(
                "Couldn't connect to vSphere host: %s" % e
            )

        super(vSphere, self)._initialize_check()

    def start(self, label, task):
        """Start a machine.
        @param label: machine name.
        @param task: task object.
        @raise CuckooMachineError: if unable to start machine.
        """
        name = self.db.view_machine_by_label(label).snapshot
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._revert_snapshot(vm, name)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def stop(self, label):
        """Stop a machine.
        @param label: machine name.
        @raise CuckooMachineError: if unable to stop machine
        """
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._stop_virtual_machine(vm)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def dump_memory(self, label, path):
        """Take a memory dump of a machine.
        @param path: path to where to store the memory dump
        @raise CuckooMachineError: if error taking the memory dump
        """
        name = "cuckoo_memdump_{0}".format(random.randint(100000, 999999))
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if vm:
                self._create_snapshot(vm, name)
                self._download_snapshot(conn, vm, name, path)
                self._delete_snapshot(vm, name)
            else:
                raise CuckooMachineError(
                    "Machine %s not found on host" % label
                )

    def _list(self):
        """List virtual machines on vSphere host"""
        ret = []
        with SmartConnection(**self.connect_opts) as conn:
            for vm in self._get_virtual_machines(conn):
                ret.append(vm.summary.config.name)
        return ret

    def _status(self, label):
        """Get power state of vm from vSphere host.
        @param label: virtual machine name
        @raise CuckooMachineError: if error getting status or machine not found
        """
        with SmartConnection(**self.connect_opts) as conn:
            vm = self._get_virtual_machine_by_label(conn, label)
            if not vm:
                raise CuckooMachineError(
                    "Machine %s not found on server" % label
                )

            status = vm.runtime.powerState
            self.set_status(label, status)
            return status

    def _get_virtual_machines(self, conn):
        """Iterate over all VirtualMachine managed objects on vSphere host"""
        def traverseDCFolders(conn, nodes, path=""):
            for node in nodes:
                if hasattr(node, "childEntity"):
                    for child, childpath in traverseDCFolders(conn, node.childEntity, path + node.name + "/"):
                        yield child, childpath
                else:
                    yield node, path + node.name

        def traverseVMFolders(conn, nodes):
            for node in nodes:
                if hasattr(node, "childEntity"):
                    for child in traverseVMFolders(conn, node.childEntity):
                        yield child
                else:
                    yield node

        self.VMtoDC = {}

        for dc, dcpath in traverseDCFolders(conn, conn.content.rootFolder.childEntity):
            for vm in traverseVMFolders(conn, dc.vmFolder.childEntity):
                try:
                    self.VMtoDC[vm.summary.config.name] = dcpath
                    yield vm
                except AttributeError:
                    pass

    def _get_virtual_machine_by_label(self, conn, label):
        """Return the named VirtualMachine managed object"""
        for vm in self._get_virtual_machines(conn):
            try:
                if vm.summary.config.name == label:
                    return vm
            except AttributeError:
                pass

    def _get_snapshot_by_name(self, vm, name):
        """Return the named VirtualMachineSnapshot managed object for
        a virtual machine"""
        for ss in self._traverseSnapshots(vm.snapshot.rootSnapshotList):
            if ss.name == name:
                return ss.snapshot

    def _get_snapshot_power_state(self, vm, name):
        """Return the power state for a named VirtualMachineSnapshot object"""
        for ss in self._traverseSnapshots(vm.snapshot.rootSnapshotList):
            if ss.name == name:
                return ss.state

    def _create_snapshot(self, vm, name):
        """Create named snapshot of virtual machine"""
        log.info(
            "Creating snapshot %s for machine %s",
            name, vm.summary.config.name
        )

        task = vm.CreateSnapshot_Task(name=name,
                                      description="Created by Cuckoo sandbox",
                                      memory=True,
                                      quiesce=False)
        try:
            self._wait_task(task)
        except CuckooMachineError as e:
            raise CuckooMachineError("CreateSnapshot: %s" % e)

    def _delete_snapshot(self, vm, name):
        """Remove named snapshot of virtual machine"""
        snapshot = self._get_snapshot_by_name(vm, name)
        if snapshot:
            log.info(
                "Removing snapshot %s for machine %s",
                name, vm.summary.config.name
            )

            task = snapshot.RemoveSnapshot_Task(removeChildren=True)
            try:
                self._wait_task(task)
            except CuckooMachineError as e:
                log.error("RemoveSnapshot: {0}".format(e))
        else:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

    def _revert_snapshot(self, vm, name):
        """Revert virtual machine to named snapshot"""
        snapshot = self._get_snapshot_by_name(vm, name)
        if snapshot:
            log.info(
                "Reverting machine %s to snapshot %s",
                vm.summary.config.name, name
            )

            task = snapshot.RevertToSnapshot_Task()
            try:
                self._wait_task(task)
            except CuckooMachineError as e:
                raise CuckooMachineError("RevertToSnapshot: %s" % e)
        else:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

    def _download_snapshot(self, conn, vm, name, path):
        """Download snapshot file from host to local path"""

        # Get filespec to .vmsn or .vmem file of named snapshot
        snapshot = self._get_snapshot_by_name(vm, name)
        if not snapshot:
            raise CuckooMachineError(
                "Snapshot %s for machine %s not found" %
                (name, vm.summary.config.name)
            )

        memorykey = datakey = filespec = None
        for s in vm.layoutEx.snapshot:
            if s.key == snapshot:
                memorykey = s.memoryKey
                datakey = s.dataKey
                break

        for f in vm.layoutEx.file:
            if f.key == memorykey and (f.type == "snapshotMemory" or
                                       f.type == "suspendMemory"):
                filespec = f.name
                break

        if not filespec:
            for f in vm.layoutEx.file:
                if f.key == datakey and f.type == "snapshotData":
                    filespec = f.name
                    break

        if not filespec:
            raise CuckooMachineError("Could not find snapshot memory file")

        log.info("Downloading memory dump %s to %s", filespec, path)

        # Parse filespec to get datastore and file path.
        datastore, filepath = re.match(r"\[([^\]]*)\] (.*)", filespec).groups()

        # Construct URL request
        params = {
            "dsName": datastore,
            "dcPath": self.VMtoDC.get(vm.summary.config.name, "ha-datacenter")
        }
        headers = {
            "Cookie": conn._stub.cookie,
        }
        url = "https://%s:%s/folder/%s" % (
            self.connect_opts["host"], self.connect_opts["port"], filepath
        )

        # Stream download to specified local path
        try:
            response = requests.get(url, params=params, headers=headers,
                                    verify=False, stream=True)

            response.raise_for_status()

            with open(path, "wb") as localfile:
                for chunk in response.iter_content(16*1024):
                    localfile.write(chunk)

        except Exception as e:
            raise CuckooMachineError(
                "Error downloading memory dump %s: %s" %
                (filespec, e)
            )

    def _stop_virtual_machine(self, vm):
        """Power off a virtual machine"""
        log.info("Powering off virtual machine %s", vm.summary.config.name)
        task = vm.PowerOffVM_Task()
        try:
            self._wait_task(task)
        except CuckooMachineError as e:
            log.error("PowerOffVM: %s", e)

    def _wait_task(self, task):
        """Wait for a task to complete with timeout"""
        limit = datetime.timedelta(
            seconds=int(self.options_globals.timeouts.vm_state)
        )
        start = datetime.datetime.utcnow()

        while True:
            if task.info.state == "error":
                raise CuckooMachineError("Task error")

            if task.info.state == "success":
                break

            if datetime.datetime.utcnow() - start > limit:
                raise CuckooMachineError("Task timed out")

            time.sleep(1)

    def _traverseSnapshots(self, root):
        """Recursive depth-first traversal of snapshot tree"""
        for node in root:
            if len(node.childSnapshotList) > 0:
                for child in self._traverseSnapshots(node.childSnapshotList):
                    yield child
            yield node
<|endoftext|>


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

import os
import hashlib

from lib.cuckoo.common.abstracts import Report
from lib.cuckoo.common.exceptions import CuckooDependencyError, CuckooReportError
from lib.cuckoo.common.objects import File

try:
    from pymongo.connection import Connection
    from pymongo.errors import ConnectionFailure, InvalidDocument
    from gridfs import GridFS
    from gridfs.errors import FileExists
except ImportError:
    raise CuckooDependencyError("Unable to import pymongo")

class MongoDB(Report):
    """Stores report in MongoDB."""

    def run(self, results):
        """Writes report.
        @param results: Cuckoo results dict.
        @raise CuckooReportError: if fails to connect or write to MongoDB.
        """
        if not results["success"]:
            return

        self._connect()

        # Set an unique index on stored files, to avoid duplicates.
        # From pymongo docs:
        #  Returns the name of the created index if an index is actually created.
        #  Returns None if the index already exists.
        self._db.fs.files.ensure_index("md5", unique=True, name="md5_unique")

        # Add pcap file, check for dups and in case add only reference.
        pcap_file = os.path.join(self.analysis_path, "dump.pcap")
        pcap = File(pcap_file)
        if pcap.valid():
            pcap_id = self.store_file(pcap)

            # Preventive key check.
            if "network" in results and isinstance(results["network"], dict):
                results["network"]["pcap_id"] = pcap_id
            else:
                results["network"] = {"pcap_id": pcap_id}

        # Add dropped files, check for dups and in case add only reference.
        dropped_files = {}
        for dir_name, dir_names, file_names in os.walk(os.path.join(self.analysis_path, "files")):
            for file_name in file_names:
                file_path = os.path.join(dir_name, file_name)
                drop = File(file_path)
                dropped_files[drop.get_md5()] = drop

        result_files = dict((dropped.get("md5", None), dropped) for dropped in results["dropped"])

        # hopefully the md5s in dropped_files and result_files should be the same
        if set(dropped_files.keys()) - set(result_files.keys()):
            log.warning("Dropped files in result dict are different from those in storage.")

        # store files in gridfs
        for md5, fileobj in dropped_files.items():
            # only store in db if we have a filename for it in results (should be all)
            resultsdrop = result_files.get(md5, None)
            if resultsdrop and fileobj.valid():
                drop_id = self.store_file(fileobj, filename=resultsdrop["name"])
                resultsdrop["dropped_id"] = drop_id

        # Add screenshots.
        results["shots"] = []
        shots_path = os.path.join(self.analysis_path, "shots")
        if os.path.exists(shots_path):
            shots = [f for f in os.listdir(shots_path) if f.endswith(".jpg")]
            for shot_file in sorted(shots):
                shot_path = os.path.join(self.analysis_path, "shots", shot_file)
                shot = File(shot_path)
                if shot.valid():
                    shot_id = self.store_file(shot)
                    results["shots"].append(shot_id)

        # Save all remaining results.
        try:
            self._db.analysis.save(results, manipulate=False)
        except InvalidDocument:
            # The document is too big, we need to shrink it and re-save it.
            results["behavior"]["processes"] = None

            # Let's add an error message to the debug block.
            error = ("The analysis results were too big to be stored, " +
                     "the detailed behavioral analysis has been stripped out.")
            results["debug"]["errors"].append(error)

            # Try again to store, if it fails, just abort.
            try:
                self._db.analysis.save(results)
            except Exception as e:
                raise CuckooReportError("Failed to store the document into MongoDB: %s" % e)

    def store_file(self, fileobj, filename=None):
        if filename == None: filename = fileobj.get_name()

        existing = self._db.fs.files.find_one({"md5": fileobj.get_md5()})
        if not existing:
            gfsfile = self._fs.new_file(filename=filename)
            for chunk in fileobj.get_chunks():
                gfsfile.write(chunk)
            gfsfile.close()

            return gfsfile._id

        return existing["_id"]

    def _connect(self):
        """Connects to Mongo database, loads options and set connectors.
        @raise CuckooReportError: if unable to connect.
        """
        host = self.options.get("host", "127.0.0.1")
        port = self.options.get("port", 27017)

        try:
            self._conn = Connection(host, port)
            self._db = self._conn.cuckoo
            self._fs = GridFS(self._db)
        except TypeError:
            raise CuckooReportError("Mongo connection port must be integer")
        except ConnectionFailure:
            raise CuckooReportError("Cannot connect to MongoDB")
<commit_msg>Removed not used import<commit_after># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os

from lib.cuckoo.common.abstracts import Report
from lib.cuckoo.common.exceptions import CuckooDependencyError, CuckooReportError
from lib.cuckoo.common.objects import File

try:
    from pymongo.connection import Connection
    from pymongo.errors import ConnectionFailure, InvalidDocument
    from gridfs import GridFS
    from gridfs.errors import FileExists
except ImportError:
    raise CuckooDependencyError("Unable to import pymongo")

class MongoDB(Report):
    """Stores report in MongoDB."""

    def run(self, results):
        """Writes report.
        @param results: Cuckoo results dict.
        @raise CuckooReportError: if fails to connect or write to MongoDB.
        """
        if not results["success"]:
            return

        self._connect()

        # Set an unique index on stored files, to avoid duplicates.
        # From pymongo docs:
        #  Returns the name of the created index if an index is actually created.
        #  Returns None if the index already exists.
        self._db.fs.files.ensure_index("md5", unique=True, name="md5_unique")

        # Add pcap file, check for dups and in case add only reference.
        pcap_file = os.path.join(self.analysis_path, "dump.pcap")
        pcap = File(pcap_file)
        if pcap.valid():
            pcap_id = self.store_file(pcap)

            # Preventive key check.
            if "network" in results and isinstance(results["network"], dict):
                results["network"]["pcap_id"] = pcap_id
            else:
                results["network"] = {"pcap_id": pcap_id}

        # Add dropped files, check for dups and in case add only reference.
        dropped_files = {}
        for dir_name, dir_names, file_names in os.walk(os.path.join(self.analysis_path, "files")):
            for file_name in file_names:
                file_path = os.path.join(dir_name, file_name)
                drop = File(file_path)
                dropped_files[drop.get_md5()] = drop

        result_files = dict((dropped.get("md5", None), dropped) for dropped in results["dropped"])

        # hopefully the md5s in dropped_files and result_files should be the same
        if set(dropped_files.keys()) - set(result_files.keys()):
            log.warning("Dropped files in result dict are different from those in storage.")

        # store files in gridfs
        for md5, fileobj in dropped_files.items():
            # only store in db if we have a filename for it in results (should be all)
            resultsdrop = result_files.get(md5, None)
            if resultsdrop and fileobj.valid():
                drop_id = self.store_file(fileobj, filename=resultsdrop["name"])
                resultsdrop["dropped_id"] = drop_id

        # Add screenshots.
        results["shots"] = []
        shots_path = os.path.join(self.analysis_path, "shots")
        if os.path.exists(shots_path):
            shots = [f for f in os.listdir(shots_path) if f.endswith(".jpg")]
            for shot_file in sorted(shots):
                shot_path = os.path.join(self.analysis_path, "shots", shot_file)
                shot = File(shot_path)
                if shot.valid():
                    shot_id = self.store_file(shot)
                    results["shots"].append(shot_id)

        # Save all remaining results.
        try:
            self._db.analysis.save(results, manipulate=False)
        except InvalidDocument:
            # The document is too big, we need to shrink it and re-save it.
            results["behavior"]["processes"] = None

            # Let's add an error message to the debug block.
            error = ("The analysis results were too big to be stored, " +
                     "the detailed behavioral analysis has been stripped out.")
            results["debug"]["errors"].append(error)

            # Try again to store, if it fails, just abort.
            try:
                self._db.analysis.save(results)
            except Exception as e:
                raise CuckooReportError("Failed to store the document into MongoDB: %s" % e)

    def store_file(self, fileobj, filename=None):
        if filename == None: filename = fileobj.get_name()

        existing = self._db.fs.files.find_one({"md5": fileobj.get_md5()})
        if not existing:
            gfsfile = self._fs.new_file(filename=filename)
            for chunk in fileobj.get_chunks():
                gfsfile.write(chunk)
            gfsfile.close()

            return gfsfile._id

        return existing["_id"]

    def _connect(self):
        """Connects to Mongo database, loads options and set connectors.
        @raise CuckooReportError: if unable to connect.
        """
        host = self.options.get("host", "127.0.0.1")
        port = self.options.get("port", 27017)

        try:
            self._conn = Connection(host, port)
            self._db = self._conn.cuckoo
            self._fs = GridFS(self._db)
        except TypeError:
            raise CuckooReportError("Mongo connection port must be integer")
        except ConnectionFailure:
            raise CuckooReportError("Cannot connect to MongoDB")
<|endoftext|>


<commit_before>import rave.log
import rave.events
from ..common import events_for, handle_events as handle_system_events

_log = rave.log.get(__name__)

def handle_events():
    # Handle system events.
    handle_system_events()

    for ev in events_for('video'):
        # Dispatch events.
        if ev.type == sdl2.SDL_WINDOWEVENT:
            window = ev.window.windowID

            if ev.window.event == sdl2.SDL_WINDOWEVENT_SHOWN:
                rave.events.emit('video.window.shown', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_HIDDEN:
                rave.events.emit('video.window.hidden', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_EXPOSED:
                rave.events.emit('video.window.exposed', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MOVED:
                rave.events.emit('video.window.moved', window, ev.window.data1, ev.window.data2)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_SIZE_CHANGED:
                # Handled in SDL_WINDOWEVENT_RESIZED.
                pass
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_RESIZED:
                rave.events.emit('video.window.resized', window, ev.window.data1, ev.window.data2)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MINIMIZED:
                rave.events.emit('video.window.minimized', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MAXIMIZED:
                rave.events.emit('video.window.maximized', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_RESTORED:
                rave.events.emit('video.window.restored', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_ENTER:
                rave.events.emit('video.window.focused', window, 'mouse')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_LEAVE:
                rave.events.emit('video.window.unfocused', window, 'mouse')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_FOCUS_GAINED:
                rave.events.emit('video.window.focused', window, 'keyboard')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_FOCUS_LOST:
                rave.events.emit('video.window.unfocused', window, 'keyboard')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_CLOSE:
                rave.events.emit('video.window.close', window)
            else:
                _log.warn('Got unknown SDL_WINDOWEVENT of type {}.', ev.window.event)
        elif ev.type == sdl2.SDL_SYSWMEVENT:
            # Nothing for now.
            pass
        elif ev.type == sdl2.SDL_RENDER_TARGETS_RESET:
            # Nothing for now.
            pass
        elif ev.type == sdl2.SDL_RENDER_DEVICE_RESET:
            # Nothing for now.
            pass
        else:
            _log.warn('Got unknown event of type {}.', ev.type)
<commit_msg>modules/sdl2: Fix broken import.<commit_after>import sdl2
import rave.log
import rave.events
from ..common import events_for, handle_events as handle_system_events

_log = rave.log.get(__name__)

def handle_events():
    # Handle system events.
    handle_system_events()

    for ev in events_for('video'):
        # Dispatch events.
        if ev.type == sdl2.SDL_WINDOWEVENT:
            window = ev.window.windowID

            if ev.window.event == sdl2.SDL_WINDOWEVENT_SHOWN:
                rave.events.emit('video.window.shown', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_HIDDEN:
                rave.events.emit('video.window.hidden', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_EXPOSED:
                rave.events.emit('video.window.exposed', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MOVED:
                rave.events.emit('video.window.moved', window, ev.window.data1, ev.window.data2)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_SIZE_CHANGED:
                # Handled in SDL_WINDOWEVENT_RESIZED.
                pass
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_RESIZED:
                rave.events.emit('video.window.resized', window, ev.window.data1, ev.window.data2)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MINIMIZED:
                rave.events.emit('video.window.minimized', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_MAXIMIZED:
                rave.events.emit('video.window.maximized', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_RESTORED:
                rave.events.emit('video.window.restored', window)
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_ENTER:
                rave.events.emit('video.window.focused', window, 'mouse')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_LEAVE:
                rave.events.emit('video.window.unfocused', window, 'mouse')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_FOCUS_GAINED:
                rave.events.emit('video.window.focused', window, 'keyboard')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_FOCUS_LOST:
                rave.events.emit('video.window.unfocused', window, 'keyboard')
            elif ev.window.event == sdl2.SDL_WINDOWEVENT_CLOSE:
                rave.events.emit('video.window.close', window)
            else:
                _log.warn('Got unknown SDL_WINDOWEVENT of type {}.', ev.window.event)
        elif ev.type == sdl2.SDL_SYSWMEVENT:
            # Nothing for now.
            pass
        elif ev.type == sdl2.SDL_RENDER_TARGETS_RESET:
            # Nothing for now.
            pass
        elif ev.type == sdl2.SDL_RENDER_DEVICE_RESET:
            # Nothing for now.
            pass
        else:
            _log.warn('Got unknown event of type {}.', ev.type)
<|endoftext|>